changed yolo

api/src/main.py

@@ -0,0 +1,4 @@
+from yolo.train_yolo_faces import train_yolo
+
+
+train_yolo()

api/src/r_cnn/cocoDetectionDataset.py

@@ -58,68 +58,3 @@ class CocoDetectionDataset(Dataset):
             image = self.transforms(image)
  
         return image, target
-    
-
-class CocoDetectionDatasetResized(Dataset):
-    def __init__(self, image_dir, annotation_path, img_size=64, transforms=None):
-        self.image_dir = image_dir
-        self.coco = COCO(annotation_path)
-        self.image_ids = list(self.coco.imgs.keys())
-        self.transforms = transforms
-        self.img_size = img_size
-
-    def __len__(self):
-        return len(self.image_ids)
-
-    def __getitem__(self, idx):
-        image_id = self.image_ids[idx]
-        image_info = self.coco.loadImgs(image_id)[0]
-        image_path = os.path.join(self.image_dir, image_info['file_name'])
-
-        image = Image.open(image_path).convert("RGB")
-
-        orig_w, orig_h = image.size
-
-        scale_w = self.img_size / orig_w
-        scale_h = self.img_size / orig_h
-
-        # Load annotations
-        annotation_ids = self.coco.getAnnIds(imgIds=image_id)
-        annotations = self.coco.loadAnns(annotation_ids)
-
-        boxes = []
-        labels = []
-
-        for obj in annotations:
-            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)
-
-            boxes.append([xmin, ymin, xmax, ymax])
-            labels.append(obj['category_id'])
-
-        boxes = torch.tensor(boxes, dtype=torch.float32)
-        labels = torch.tensor(labels, dtype=torch.long)
-
-        area = torch.tensor([obj['area'] * scale_w * scale_h for obj in annotations], dtype=torch.float32)
-        iscrowd = torch.tensor([obj.get('iscrowd', 0) for obj in annotations], dtype=torch.long)
-
-        target = {
-            "boxes": boxes,
-            "labels": labels,
-            "image_id": torch.tensor([image_id]),
-            "area": area,
-            "iscrowd": iscrowd
-        }
-
-        # resize image
-        image = image.resize((self.img_size, self.img_size))
-
-        if self.transforms:
-            image = self.transforms(image)
-
-        return image, target

api/src/yolo/train_yolo_faces.py

@@ -2,9 +2,9 @@ import os
 
 from torchvision import transforms
 from tqdm import tqdm
-from yolo_model import train, Yolo_model
+from yolo.yolo_dataset import YoloDataset
-from yolo_loss import YoloLoss
+from yolo.yolo_model import train, Yolo_model
-from cocoDetectionDataset import CocoDetectionDatasetResized
+from yolo.yolo_loss import YoloLoss
 from util import TransformedSubset, visualizeImage
 from torch.utils.data import DataLoader, random_split
 from torchvision import datasets
@@ -14,42 +14,41 @@ from torchvision.io import read_image
 from torch.utils.data import WeightedRandomSampler
 
 
+def train_yolo():
+    SAVE_PATH = "./saved_models"
+    IMAGE_SIZE = 64
+    GRID = 9
+    BATCH_SIZE = 256
 
-SAVE_PATH = "./saved_models"
+    transform = transforms.Compose([
-IMAGE_SIZE = 128
-GRID = 9
-BATCH_SIZE = 1
-
-transform = transforms.Compose([
         transforms.ToTensor()
-])
+    ])
 
-dataset = CocoDetectionDatasetResized(
+    dataset = YoloDataset(
         image_dir="data/faces/train", 
         annotation_path="data/faces/train/_annotations.coco.json",
         img_size=IMAGE_SIZE,
         transforms=transform
-)
+    )
 
-train_size = int(0.8 * len(dataset))
+    train_size = int(0.8 * len(dataset))
-val_size = len(dataset) - train_size
+    val_size = len(dataset) - train_size
 
-train_subset, val_subset = random_split(dataset, [train_size, val_size])
+    train_subset, val_subset = random_split(dataset, [train_size, val_size])
 
 
-train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))
+    train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=False)
-val_loader = DataLoader(val_subset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=lambda x: tuple(zip(*x)))
+    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")
+    device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
-print("Using device", device)
+    print("Using device", device)
 
-model = Yolo_model(c_in=3, c_hidden=32, boxes=2, img_size=IMAGE_SIZE, grid=GRID, labels=1)
+    model = Yolo_model(c_in=3, c_hidden=32, boxes=1, img_size=IMAGE_SIZE, grid=GRID, labels=1)
-model.to(device)
+    model.to(device)
-optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
-loss_module = YoloLoss()
+    loss_module = YoloLoss()
 
-train(model=model, loss_module=loss_module, train_loader=train_loader, val_loader=val_loader, 
+
+    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, 
-            model_name="face_detection_yolo", img_size=IMAGE_SIZE, num_classes=1, grid=GRID)
+                model_name="face_detection_yolo")
-
-exit()

api/src/yolo/yolo_dataset.py

@@ -0,0 +1,105 @@
+import os
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+from pycocotools.coco import COCO 
+import numpy as np
+
+
+def is_center_in_grid_cell(x, y, img_w, img_h, S, cell_i, cell_j):
+        cell_w = img_w / S
+        cell_h = img_h / S
+
+        # find which cell the center belongs to
+        gt_cell_i = int(y / cell_h)
+        gt_cell_j = int(x / cell_w)
+
+        return (gt_cell_i == cell_i) and (gt_cell_j == cell_j)
+
+def one_hot(index, num_classes):
+    if index < 0 or index >= num_classes:
+        raise ValueError(f"Index out of range, index: {index}, num_classes: {num_classes}")
+
+    encoding = [0] * num_classes
+    encoding[index] = 1
+    return encoding
+
+class YoloDataset(Dataset):
+    def __init__(self, image_dir, annotation_path, img_size=64, grid = 9, transforms=None):
+        self.image_dir = image_dir
+        self.coco = COCO(annotation_path)
+        self.image_ids = list(self.coco.imgs.keys())
+        self.transforms = transforms
+        self.img_size = img_size
+        self.grid = grid
+
+    def __len__(self):
+        return len(self.image_ids)
+
+    def __getitem__(self, idx):
+        image_id = self.image_ids[idx]
+        image_info = self.coco.loadImgs(image_id)[0]
+        image_path = os.path.join(self.image_dir, image_info['file_name'])
+
+        image = Image.open(image_path).convert("RGBA").convert("RGB")
+
+        orig_w, orig_h = image.size
+
+        scale_w = self.img_size / orig_w
+        scale_h = self.img_size / orig_h
+
+        # Load annotations
+        annotation_ids = self.coco.getAnnIds(imgIds=image_id)
+        annotations = self.coco.loadAnns(annotation_ids)
+
+        boxes = []
+        labels = []
+
+        for obj in annotations:
+            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)
+
+            boxes.append([xmin, ymin, xmax, ymax])
+            labels.append(obj['category_id'] - 1)
+
+        ground_truth = list(zip(boxes, labels))
+        num_labels = max(labels) + 1 if len(labels) > 0 else 1
+        #[S, S, (x+y+w+h+c+C)]
+        targets = np.zeros((self.grid, self.grid, 5 + num_labels))
+        for x in range(self.grid):
+            for y in range(self.grid):
+                for i in range(len(ground_truth)):
+                    box = ground_truth[i][0]
+                    label = ground_truth[i][1]
+
+                    if is_center_in_grid_cell(x=box[0]+box[2]/2, y=box[1]+box[3]/2, img_w=self.img_size, 
+                                              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, i + 5] = class_one_hot[i]
+
+                        del ground_truth[i]
+
+                        break
+
+
+        targets = torch.tensor(targets, dtype=torch.float32)
+
+        # resize image
+        image = image.resize((self.img_size, self.img_size))
+
+        if self.transforms:
+            image = self.transforms(image)
+
+        return image, targets

api/src/yolo/yolo_loss.py

@@ -5,103 +5,20 @@ class YoloLoss(nn.Module):
     def __init__(self):
         super(YoloLoss, self).__init__()
 
-    @staticmethod
+    def forward(self, predictions, targets, lambda_coord=5, lambda_noobj=0.5):
-    def is_center_in_grid_cell(x, y, img_w, img_h, S, cell_i, cell_j):
+        pred_boxes = predictions[..., :4]
-        cell_w = img_w / S
+        pred_conf = predictions[..., 4]
-        cell_h = img_h / S
+        pred_classes = predictions[..., 5:]
+        target_boxes = targets[..., :4]
+        target_conf = targets[..., 4]
+        target_classes = targets[..., 5:]
         
-        # find which cell the center belongs to
+        box_loss = lambda_coord * torch.mean((pred_boxes - target_boxes) ** 2)
-        gt_cell_i = int(y / cell_h)
-        gt_cell_j = int(x / cell_w)
 
-        return (gt_cell_i == cell_i) and (gt_cell_j == cell_j)
+        obj_loss = torch.mean((pred_conf[target_conf == 1] - target_conf[target_conf == 1]) ** 2)
+        noobj_loss = lambda_noobj * torch.mean((pred_conf[target_conf == 0]) ** 2)
 
-    @staticmethod
+        class_loss = torch.mean((pred_classes[target_conf == 1] - target_classes[target_conf == 1]) ** 2)
-    def save_sqrt(i):
-        return torch.sqrt(torch.clamp(i, min=1e-6))
 
-    @staticmethod
+        total_loss = box_loss + obj_loss + noobj_loss + class_loss
-    def iou(boxA, boxB):
+        return total_loss
-        xA = torch.max(boxA[0], boxB[0])
-        yA = torch.max(boxA[1], boxB[1])
-        xB = torch.min(boxA[2], boxB[2])
-        yB = torch.min(boxA[3], boxB[3])
-
-        inter_area = torch.clamp(xB - xA, min=0) * torch.clamp(yB - yA, min=0)
-
-        boxA_area = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
-        boxB_area = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])
-
-        union = boxA_area + boxB_area - inter_area
-
-        return inter_area / (union + 1e-6)
-    
-    @staticmethod
-    def xywh_to_xyxy(box):
-        x, y, w, h = box
-
-        return torch.stack([
-            x - w / 2,
-            y - h / 2,
-            x + w / 2,
-            y + h / 2
-        ])
-
-    def forward(self, pred, target, labels, cell_size, img_w, img_h, factor_no_object = 0.5, factor_object = 5):
-        #pred: [SxSx(B(x,y,w,h,con)C)]
-        #target: [B(x, y, w, h, C)]
-        
-        for s0 in range(pred.shape[0]):
-            for s1 in range(pred.shape[1]):
-                boxes = pred[s0, s1, :(-1*labels)]
-                pred_label = pred[s0, s1, -1:]
-
-                ground_truth_box = None
-                for b in range(target.shape[0] // (4 + labels)):
-                    box = target[(b)*(4 + labels):(b+1)*(4 + labels)]
-                    if self.is_center_in_grid_cell(box[0]+box[2]/2, box[1]+box[3]/2, img_w, img_h, cell_size, s0, s1):
-                        ground_truth_box = box[:4]
-                        ground_truth_label = box[4:]
-                        break
-
-
-                coord_loss = torch.tensor(0., device=pred.device)
-                size_loss = torch.tensor(0., device=pred.device)
-                prob_loss = torch.tensor(0., device=pred.device)
-                label_loss = torch.tensor(0., device=pred.device)
-
-                if ground_truth_box is not None:
-
-                    iou_stats = []
-
-                    for b in range(boxes.shape[0] // 5):
-                        box = boxes[(b)*5:(b+1)*5]
-
-                        iou_stats.append(self.iou(self.xywh_to_xyxy(box[0:4]), self.xywh_to_xyxy(ground_truth_box[0:4])))
-                    iou_tensor = torch.stack(iou_stats)
-                    responsible = torch.argmax(iou_tensor)
-
-                    for b in range(boxes.shape[0] // 5):
-                        box = boxes[(b)*5:(b+1)*5]
-
-                        is_responsible = (b == responsible).float()
-
-                        coord_loss += factor_object * is_responsible * ((box[0] - ground_truth_box[0])**2 + (box[1] - ground_truth_box[1])**2)
-                        size_loss += factor_object * is_responsible * ((self.save_sqrt(box[2]) - self.save_sqrt(ground_truth_box[2]))**2 + 
-                                                       (self.save_sqrt(box[3]) - self.save_sqrt(ground_truth_box[3]))**2)
-                        
-                        prob_loss += is_responsible * (box[4] - 1)**2
-                        prob_loss += factor_no_object * (1 - is_responsible) * box[4]**2
-
-
-                    for i in range(pred_label.shape[0]):
-                        label_loss += ((pred_label[i] - ground_truth_label[i])**2) / pred_label.shape[0]
-                        
-
-                else:
-                    for b in range(boxes.shape[0] // 5):
-                            box = boxes[(b)*5:(b+1)*5]
-
-                            prob_loss += factor_no_object * box[4] ** 2
-
-        return (label_loss + prob_loss + coord_loss + size_loss) / pred.shape[0]**2

api/src/yolo/yolo_model.py

@@ -1,11 +1,7 @@
 import os
-import cv2
 import torch
 from tqdm import tqdm
 import torch.nn as nn
-import torch.nn.functional as F
-import torchvision.transforms.functional as TF
-import time
 
 
 class Yolo_model(nn.Module):
@@ -16,16 +12,22 @@ class Yolo_model(nn.Module):
 
         self.model = nn.Sequential(
             nn.Conv2d(in_channels=c_in, out_channels=c_hidden, kernel_size=7, padding=3),
+            nn.BatchNorm2d(c_hidden),
             nn.LeakyReLU(),
             nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1),
+            nn.BatchNorm2d(c_hidden),
             nn.LeakyReLU(),
             nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1),
+            nn.BatchNorm2d(c_hidden),
             nn.LeakyReLU(),
             nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1),
+            nn.BatchNorm2d(c_hidden),
             nn.LeakyReLU(),
             nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1),
+            nn.BatchNorm2d(c_hidden),
             nn.LeakyReLU(),
             nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1),
+            nn.BatchNorm2d(c_hidden),
             nn.LeakyReLU(),
             nn.Flatten(),
             nn.Linear(in_features=img_size*img_size*c_hidden, out_features=grid*grid*(boxes*5+labels))
@@ -37,29 +39,8 @@ class Yolo_model(nn.Module):
         x = x.reshape(batch, self.grid, self.grid, -1)
         return x
     
-def create_stack(images, annotations, num_classes, device):
-    data = []
-    labels = []
-    for i in range(len(images)):
-        image = images[i]
 
-        annotation = annotations[i]
+def train(model, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, model_name, saving=True, device="cpu"):
-        
-        label_train = []
-        for box, label in tuple(zip(annotation["boxes"], annotation["labels"])):
-            box = box
-            label = 0
-            one_hot = F.one_hot(torch.tensor(label, dtype=torch.long), num_classes=num_classes)
-            label_train.append(torch.cat((box, one_hot), dim=0))
-        
-        label_torch = torch.stack(label_train).reshape(-1).to(device)
-        labels.append(label_torch)
-        data.append(image)
-
-    return data, labels
-    
-
-def train(model, num_classes, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, model_name, grid, saving=True, device="cpu", img_size=64):
     best_val = torch.finfo(torch.float32).max
 
     for epoch in range(200):
@@ -69,18 +50,13 @@ def train(model, num_classes, loss_module, train_loader, val_loader, optimizer,
         model.train()
 
         count, lossCount = 0, 0.
-        for images, annotations in tqdm(train_loader, desc=f"Train", leave=False):
+        for images, labels in tqdm(train_loader, desc=f"Train", leave=False):
-            data, labels = create_stack(images, annotations, num_classes, device)
+            images = images.to(device)
+            labels = labels.to(device)
 
-            data = torch.stack(data).to(device)
+            prediction = model(images)
 
-            prediction = model(data)
+            loss = loss_module(prediction, labels)
-
-            loss = []
-            for pred, label in tuple(zip(prediction, labels)):
-                loss.append(loss_module(pred = pred, target= label, labels=num_classes, cell_size=grid, img_w = img_size, img_h = img_size))
-            loss = torch.stack(loss).to(device)
-            loss = loss.sum()
             lossCount += loss.item()
 
             optimizer.zero_grad()
@@ -89,7 +65,7 @@ def train(model, num_classes, loss_module, train_loader, val_loader, optimizer,
 
             optimizer.step()
 
-            count += data.size(0)
+            count += images.size(0)
 
         train_loss = lossCount / count
 
@@ -101,23 +77,20 @@ def train(model, num_classes, loss_module, train_loader, val_loader, optimizer,
         model.eval()
 
         count, lossCount = 0, 0.
-        for images, annotations in tqdm(val_loader, desc=f"Test", leave=False):
+        for images, labels in tqdm(val_loader, desc=f"Test", leave=False):
             with torch.no_grad():
-                data, labels = create_stack(images, annotations, num_classes, device)
+                images = images.to(device)
+                labels = labels.to(device)
 
-            data = torch.stack(data).to(device)
+                prediction = model(images)
 
-            prediction = model(data)
+                loss = loss_module(prediction, labels)
-
+                lossCount += loss.item()
-            loss = []
-            for pred, label in tuple(zip(prediction, labels)):
-                loss.append(loss_module(pred = pred, target= label, labels=num_classes, cell_size=grid, img_w = img_size, img_h = img_size))
-            loss = torch.stack(loss).to(device)
 
 
                 lossCount += loss.sum().item()
 
-            count += data.size(0)
+                count += images.size(0)
 
         val_loss = lossCount / count