update

classifier_config.json

@@ -2,5 +2,5 @@
 	"n_neighbors": 10,
 	"classes" : ["black_red", "green_orange", "yellow_grey"],
 	"path_to_dataset": "./triplet_dataset",
-	"embedding_model": "./model_embedding.pt"
+	"embedding_model": "./embedding-output/model_embedding.pt"
 }

download-test-videos

@@ -0,0 +1,4 @@
+#!/bin/bash
+gdown 1Mm24z7fe1fkbcTt05IQpLlNdSALJQFRc
+unzip -q test_videos_2022.zip
+rm test_videos_2022.zip

embedding_config.json

@@ -1,5 +1,5 @@
 {
-	"epochs": 4,
+	"epochs": 1,
 	"embedding_dims": 128,
 	"batch_size": 32,
 	"classes" : ["black_red", "green_orange", "yellow_grey"],

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()
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

train-embedding.py

@@ -7,6 +7,7 @@ import torch.nn as nn
 import torch.optim as optim
 import matplotlib.pyplot as plt
 import matplotlib as mpl
+import os
 from pathlib import Path
 import json
 from torch.utils.data import Dataset, DataLoader
@@ -176,6 +177,8 @@ train_loss = []
 test_loss = []
 epoch_all_loss = []
 
+print('Training.')
+
 for epoch in range(epochs):
     
     train_one_epoch_loss = []
@@ -219,14 +222,27 @@ for epoch in range(epochs):
 
     print(f"Epoch: {epoch+1}/{epochs} - Training loss: {np.mean(train_one_epoch_loss):.4f} - Test loss: {np.mean(test_one_epoch_loss):.4f}")
 
+
+output_dir = Path('./embedding-output')
+os.makedirs(output_dir, exist_ok=True)
+
+torch.save(model.state_dict(), output_dir/'model_embedding.pt')
+
 train_loss = np.array(train_loss)
 test_loss = np.array(test_loss)
-q = 10
+
-plt.plot(np.convolve(train_loss, np.ones(len(train_loss) - len(test_loss) + q)/(len(train_loss) - len(test_loss) + q), mode = 'valid'), legend = 'Train loss')
+loss = {}
-plt.plot(np.convolve(test_loss, np.ones(q)/q, mode = 'valid'), legend = 'Test loss')
+loss["test_loss"] = list(test_loss.astype(float))
+loss["train_loss"] = list(train_loss.astype(float))
+with open(str(output_dir / 'loss.json'), 'w') as f:
+    json.dump(loss, f, indent = 4)
+
+q = 10 # plotting parameter
+plt.plot(np.convolve(train_loss, np.ones(len(train_loss) - len(test_loss) + q)/(len(train_loss) - len(test_loss) + q), mode = 'valid'), label = 'Train loss')
+plt.plot(np.convolve(test_loss, np.ones(q)/q, mode = 'valid'), label = 'Test loss')
 plt.legend()
 plt.title("Epoch loss")
-plt.savefig("embedding_loss.pdf")
+plt.savefig(output_dir/"embedding_loss.pdf")
 
 if config["visualize"]:
 
@@ -261,5 +277,6 @@ if config["visualize"]:
     for i in range(len(classes)):
       legend.legendHandles[i]._sizes = [30]
     
-    plt.savefig('visualized_simple.pdf')
+    plt.axis('off')
+    plt.savefig(output_dir/'visualized_simple.pdf')