51 lines
1.5 KiB
Python
51 lines
1.5 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 = 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)
|
|
# # %%
|
|
|
|
|
|
|
|
# %%
|