added yolo data viewer

api/src/yolo/train_yolo_faces.py

@@ -12,13 +12,38 @@ import torch
 import torch.nn as nn
 from torchvision.io import read_image
 from torch.utils.data import WeightedRandomSampler
+from torchvision.utils import draw_bounding_boxes
+
+def view_data(dataloader):
+    for images, labels in iter(dataloader):
+        for batch in range(images.shape[0]):
+            image = images[batch]
+            label = labels[batch]
+            image_size = image.shape[1]
+            grid_size = image_size // label.shape[0]
+
+            boxes_to_draw = []
+            grids_to_draw = []
+            for x in range(label.shape[0]):
+                for y in range(label.shape[1]):
+                    grids_to_draw.append([x*grid_size, y*grid_size, (x+1)*grid_size, (y+1)*grid_size])
+                    if label[x, y, 4].item() == 0:
+                        continue
+                    boxes_to_draw.append(label[x, y, :4])
+            if len(boxes_to_draw) == 0:
+                continue
+            boxes_to_draw = torch.stack(boxes_to_draw)
+            grids_to_draw = torch.tensor(grids_to_draw)
+            image = draw_bounding_boxes(image, grids_to_draw, colors=(0, 255, 0))
+            image = draw_bounding_boxes(image, boxes_to_draw, colors=(255, 0, 0))
+            visualizeImage(image)
 
 
 def train_yolo():
     SAVE_PATH = "./saved_models"
     IMAGE_SIZE = 64
     GRID = 9
-    BATCH_SIZE = 256
+    BATCH_SIZE = 1
 
     transform = transforms.Compose([
         transforms.ToTensor()
@@ -37,7 +62,7 @@ def train_yolo():
     train_subset, val_subset = random_split(dataset, [train_size, val_size])
 
 
-    train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=False)
+    train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=True)
     val_loader = DataLoader(val_subset, batch_size=BATCH_SIZE, shuffle=False)
 
     device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
@@ -48,6 +73,9 @@ def train_yolo():
     optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
     loss_module = YoloLoss()
 
+    #view_data(train_loader)
+    #exit()
+
 
     train(model=model, loss_module=loss_module, train_loader=train_loader, val_loader=val_loader, 
                 optimizer=optimizer, SAVE_PATH=SAVE_PATH, saving=True, device=device, 

api/src/yolo/yolo_dataset.py

@@ -1,3 +1,4 @@
+import math
 import os
 import torch
 from PIL import Image
@@ -59,10 +60,10 @@ class YoloDataset(Dataset):
             xmin, ymin, width, height = obj['bbox']
             xmin, ymin, width, height = float(xmin), float(ymin), float(width), float(height)
 
-            xmin = xmin * scale_w
-            ymin = ymin * scale_h
-            xmax = (xmin + width * scale_w)
-            ymax = (ymin + height * scale_h)
+            xmin = math.ceil(xmin * scale_w)
+            ymin = math.ceil(ymin * scale_h)
+            xmax = math.floor(xmin + max(width * scale_w, 1))
+            ymax = math.floor(ymin + max(height * scale_h, 1))
 
             boxes.append([xmin, ymin, xmax, ymax])
             labels.append(obj['category_id'] - 1)
@@ -81,12 +82,12 @@ class YoloDataset(Dataset):
                                               img_h=self.img_size, S=self.grid, cell_i=x, cell_j=y):
                         
                         class_one_hot = one_hot(label, num_labels)
-                        targets[x, y, 0] = box[0]
-                        targets[x, y, 1] = box[1]
-                        targets[x, y, 2] = box[2]
-                        targets[x, y, 3] = box[3]
-                        targets[x, y, 4] = 1
-                        for i in range(len(class_one_hot)):
+                        targets[x, y, 0] = box[0]  #x
+                        targets[x, y, 1] = box[1]  #y
+                        targets[x, y, 2] = box[2]  #w
+                        targets[x, y, 3] = box[3]  #h
+                        targets[x, y, 4] = 1       #confidence
+                        for i in range(len(class_one_hot)):  #label
                             targets[x, y, i + 5] = class_one_hot[i]
 
                         del ground_truth[i]