124 lines
3.2 KiB
Python
124 lines
3.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 numpy as np
|
|
import os
|
|
import sys
|
|
import json
|
|
import cv2
|
|
import random
|
|
from model import Model
|
|
from torchvision.utils import draw_bounding_boxes
|
|
import torch as t
|
|
import matplotlib.pyplot as plt
|
|
|
|
|
|
|
|
|
|
|
|
device = 'cpu'
|
|
|
|
model_rt_path = '/home/thebears/Seafile/Designs/ML/inaturalist_models/models/'#0210701_202822.json
|
|
newest_model = os.path.join(model_rt_path, max(os.listdir(model_rt_path)).replace('.pth',''))
|
|
with open(newest_model + '.json','r') as nmj:
|
|
model_json = json.load(nmj)
|
|
|
|
cats = model_json['categories']
|
|
cats.sort(key=lambda x: x['new_id'])
|
|
num_cat = len(cats) + 1
|
|
model_type = model_json['model_type']
|
|
model = Model(num_cat, model_type)
|
|
labels = [x['name'] for x in cats]
|
|
model.load_state_dict(
|
|
torch.load(newest_model + '.pth', map_location=torch.device(device))
|
|
)
|
|
model.eval()
|
|
model.to(device)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
#rtdir = '/home/thebears/data/hummingbird_imagenet/hummingbird'
|
|
|
|
#ff = [x for x in os.listdir(rtdir) if x.endswith('.jpg')]
|
|
|
|
# img = os.path.join(rtdir, random.choice(ff))
|
|
# image = cv2.imread(img)[:, :, ::-1].copy()
|
|
# o = T.ToTensor()(image).to(device)
|
|
# img = o[None, :, :, :]
|
|
|
|
# ou = model(img)
|
|
|
|
#oimage = t.tensor(image, dtype=t.uint8).permute([2, 0, 1])
|
|
|
|
#matplotlib.use('Qt5Agg')
|
|
|
|
|
|
|
|
|
|
#vid_path = '/home/thebears/data/hummingbird_videos/Hummingbird_01_20210601055009.mp4'
|
|
#vid_path = '/home/thebears/data/hummingbird_videos/Hummingbird_01_20210617113038.mp4'
|
|
|
|
print('model loaded')
|
|
# %%
|
|
vid_path = '/home/thebears/data/hummingbird_videos/Hummingbird_01_20210617113038.mp4'
|
|
cap = cv2.VideoCapture(vid_path)
|
|
imgs = list()
|
|
|
|
#movie = cv2.VideoWriter('/home/thebears/Seafile/Designs/ML/inaturalist_models/output.avi',cv2.VideoWriter_fourcc('M','J','P','G'), 10, (2560,1920))
|
|
|
|
frame_num = 0
|
|
# %%
|
|
while cap.isOpened():
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
# if frame_num % 10 == 1:
|
|
img = cap.read()[1]
|
|
image = img[:, :, ::-1].copy()
|
|
o = T.ToTensor()(image).to(device)
|
|
img = o[None, :, :, :]
|
|
|
|
ou = model(img)
|
|
|
|
idx = ou[0]['labels']
|
|
label_names = [labels[x-1] for x in idx]
|
|
|
|
|
|
scores = ou[0]['scores']
|
|
oimage = t.tensor(255*img.squeeze(), dtype=t.uint8)
|
|
boxes = ou[0]['boxes']
|
|
|
|
if boxes.shape[0] > 1:
|
|
boxes = boxes[[1],:]
|
|
label_names = [label_names[0]]
|
|
|
|
if boxes.shape[0] > 0:
|
|
label_names[0] += ' {0:0.2f}'.format(scores[0])
|
|
|
|
|
|
ox = draw_bounding_boxes(oimage, boxes, width=5, labels = label_names,
|
|
font='Victor Mono SemiBold Nerd Font Complete Mono Windows Compatible',font_size=50, fill = False, colors = (255, 255, 100, 100))
|
|
fname = '/home/thebears/Seafile/Designs/ML/inaturalist_models/frames/frame_{0:06g}.jpg'.format(frame_num)
|
|
from PIL import Image
|
|
im = Image.fromarray(np.uint8(ox.permute([1,2,0]).numpy()))
|
|
im.save(fname)
|
|
|
|
# plt.imshow(ox.permute([1, 2, 0]))
|
|
frame_num += 1
|
|
print(frame_num)
|
|
|
|
|
|
|
|
# %%
|
|
|
|
|
|
|