50 lines
1.2 KiB
Python
50 lines
1.2 KiB
Python
# %%
|
|
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)
|
|
|
|
|
|
|
|
# %%
|