Update

model_inference.py

@@ -0,0 +1,176 @@
+import pickle
+import os
+import sys
+from pathlib import Path
+
+import torch
+import torch.backends.cudnn as cudnn
+import random
+import numpy as np
+import pandas as pd
+import torch.nn as nn
+import torch.optim as optim
+from tqdm.notebook import tqdm
+import matplotlib.pyplot as plt
+from torch.utils.data import Dataset, DataLoader
+from torchvision.io import read_image
+from torchvision import transforms
+import cv2
+
+if './yolov5/' not in sys.path:
+  sys.path.append('./yolov5/')
+
+from models.common import DetectMultiBackend
+from utils.augmentations import letterbox
+from utils.general import scale_coords, non_max_suppression, check_img_size
+from utils.dataloaders import LoadImages
+from utils.plots import Annotator, colors, save_one_box
+
+class EmbeddingModel(nn.Module):
+    def __init__(self, emb_dim=128):
+        super(EmbeddingModel, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(3, 16, 3),
+            nn.BatchNorm2d(16),
+            nn.PReLU(),
+            nn.MaxPool2d(2),
+
+            nn.Conv2d(16, 32, 3),
+            nn.BatchNorm2d(32),
+            nn.PReLU(32),
+            nn.MaxPool2d(2),
+
+            nn.Conv2d(32, 64, 3),
+            nn.PReLU(),
+            nn.BatchNorm2d(64),
+            nn.MaxPool2d(2)
+        )
+        
+        self.fc = nn.Sequential(
+            nn.Linear(64*6*6, 256),
+            nn.PReLU(),
+            nn.Linear(256, emb_dim)
+        )
+        
+    def forward(self, x):
+        x = self.conv(x)
+        x = x.view(-1, 64*6*6)
+        x = self.fc(x)
+        return x
+
+class SquarePad:
+	def __call__(self, image):
+		_, w, h = image.size()
+		max_wh = max(w, h)
+		hp = int((max_wh - w) / 2)
+		vp = int((max_wh - h) / 2)
+		padding = (vp, hp, vp, hp)
+		return transforms.functional.pad(image, padding, 0, 'constant')
+
+class Normalize01:
+  def __call__(self, image):
+    image -= image.min()
+    image /= image.max()
+    return image
+    
+def prepare_for_embedding(image):
+  image = torch.tensor(image).permute((2, 0, 1)).float().flip(0)
+  transform = transforms.Compose([
+                                  SquarePad(),
+                                  transforms.Resize((64, 64)),
+                                  Normalize01()
+                                  ])
+  image = transform(image)
+  return image.unsqueeze(0)
+
+
+if torch.cuda.is_available(): 
+    print('Using GPU.')
+    device = 'cuda'
+else:
+    print("CUDA not detected, using CPU.")
+    device = 'cpu'
+
+model_embedding = EmbeddingModel()
+model_embedding.load_state_dict(torch.load('./embedding-output/model_embedding.pt'))
+model_embedding.to(device)
+model_embedding.eval()
+
+with open('/model_classifier.obj','rb') as file:
+  model_classifier = pickle.load(file)
+
+classes = model_classifier.__getstate__()['classes_']
+
+video = Path('/content/test_videos_2022/2022-NLS-5-NLS_05_2022_Heli_UHD_01-000140-000155-Karussell.mp4')
+
+reader = cv2.VideoCapture(str(video))
+fps = reader.get(cv2.CAP_PROP_FPS)
+w = int(reader.get(cv2.CAP_PROP_FRAME_WIDTH))
+h = int(reader.get(cv2.CAP_PROP_FRAME_HEIGHT))
+reader.release()
+
+imgsz = check_img_size((w, h), s=model.stride)
+dataset = LoadImages(video, img_size=imgsz, stride=model.stride, auto=model.pt)
+
+weights_path = Path('./yolov5/best.pt')
+model = DetectMultiBackend(weights_path, device=torch.device(device))
+
+save_dir = Path('./detection-output/')
+os.makedirs(save_dir)
+
+writer = cv2.VideoWriter(str(save_dir / 'res.mp4'), cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
+
+for frame_n, (path, im, im0s, vid_cap, s) in enumerate(dataset):
+  im = torch.from_numpy(im).to(device)
+  im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
+  im /= 255  # 0 - 255 to 0.0 - 1.0
+  im = im[None]
+  pred = model(im)
+  pred = non_max_suppression(pred, conf_thres = 0.5, max_det = 100)
+  for i, det in enumerate(pred):
+    p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
+    imc = im0.copy()
+    annotator = Annotator(imc, line_width=3, example=str(model.names), pil = True, font_size=20 )
+
+    if len(det) == 0:
+      continue
+    det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
+
+    
+    for *xyxy, conf, cls in reversed(det):
+      crop = save_one_box(xyxy, im0, file=save_dir / 'crops' / f'frame{frame_n}_{i}.jpg', BGR=True)
+      image = prepare_for_embedding(crop).to(device)
+      embedding = model_embedding(image).cpu().detach().numpy()
+      
+      probabilities = model_classifier.predict_proba(embedding)[0]
+      best = np.argmax(probabilities)
+
+      annotator.text([xyxy[0] -20, xyxy[1] - 20], classes[best])
+      # print(classes[best])
+    # print()
+
+  im0 = annotator.result().copy()
+  writer.write(im0)
+
+writer.release()
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+