f

test.py

@@ -0,0 +1,123 @@
+# %%
+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)
+
+
+
+# %%
+
+
+