# %%
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 = 2)

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


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)



# %%