# %%
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from collections import defaultdict as ddict
import json
import torch
from torchvision import datasets, transforms as T
import cv2
import numpy as np
import os
import sys


sys.path.append(r"K:\Designs\ML\inaturalist_models\data_aug")
sys.path.append(r"K:\Designs\ML\inaturalist_models\vision")


model = torchvision.models.detection.fasterrcnn_resnet50_fpn(num_classes = 3)

model.load_state_dict(torch.load('K:\Designs\ML\inaturalist_models\model_weights_start_1.pth'))
model.eval()
model.to('cuda')

#img = r'D:\ishan\ml\inaturalist\test2017\00a903fa1d23b2f8f28248e81bc1c4a4.jpg'
#img = r'J:\hummingbird_imagenet\hummingbird\Hummingbird_01_20210617093423.mp4_023.jpg'
img = r'J:\hummingbird_imagenet\hummingbird\Hummingbird_01_20210609095848.mp4_133.jpg'
import random
rtdir = r'J:\hummingbird_imagenet\hummingbird'

ff = os.listdir(rtdir)
img = os.path.join(rtdir,random.choice(ff))
image = cv2.imread(img)[:,:,::-1].copy()
o = T.ToTensor()(image).cuda()
img = o[None, :, :, :]

ou = model(img)
from torchvision.utils import draw_bounding_boxes
import torch as t
oimage = t.tensor(image, dtype=t.uint8).permute([2,0,1])
import matplotlib.pyplot as plt
ox = draw_bounding_boxes(oimage, ou[0]['boxes'], width=1)
plt.imshow(ox.permute([1,2,0]))
# %%
from data import iNaturalistDataset
sd = iNaturalistDataset(validation=True)
# # %%



# %%