changed yolo loss

api/src/main.py

@@ -1,3 +1,4 @@
+from r_cnn.r_cnn_test import train_cnn_test
 from yolo.train_yolo_faces import train_yolo
 
 

api/src/r_cnn/r_cnn.py

@@ -4,7 +4,7 @@ import torch
 import torchvision.transforms.functional as TF
 from tqdm import tqdm
 import torch.nn as nn
-from src.util import iou, visualizeImage
+from util import iou, visualizeImage
 import random
 import torch.nn.functional as F
 

api/src/r_cnn/r_cnn_test.py

@@ -8,12 +8,15 @@ import torch.nn as nn
 from .r_cnn import ObjectDetectionCNN, train, eval
 from .cocoDetectionDataset import CocoDetectionDataset
 
-SAVE_PATH = "./saved_models"
-PADDING = 20
 
 def get_transform():
     return ToTensor()
 
+def train_cnn_test():
+
+    SAVE_PATH = "./saved_models"
+    PADDING = 20
+
     train_dataset = CocoDetectionDataset(
         image_dir="data/football/train", 
         annotation_path="data/football/train/_annotations.coco.json",

api/src/yolo/yolo_loss.py

@@ -13,12 +13,15 @@ class YoloLoss(nn.Module):
         target_conf = targets[..., 4]
         target_classes = targets[..., 5:]
 
-        box_loss = lambda_coord * torch.mean((pred_boxes - target_boxes) ** 2)
+        obj_mask = targets[..., 4] == 1
+        noobj_mask = targets[..., 4] == 0
         
-        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)
+        box_loss = lambda_coord * torch.mean((pred_boxes[obj_mask] - target_boxes[obj_mask]) ** 2)
 
-        class_loss = torch.mean((pred_classes[target_conf == 1] - target_classes[target_conf == 1]) ** 2)
+        obj_loss = torch.mean((pred_conf[obj_mask] - target_conf[obj_mask]) ** 2)
+        noobj_loss = lambda_noobj * torch.mean((pred_conf[noobj_mask]) ** 2)
+
+        class_loss = torch.mean((pred_classes[obj_mask] - target_classes[obj_mask]) ** 2)
 
         total_loss = box_loss + obj_loss + noobj_loss + class_loss
         return total_loss

api/src/yolo/yolo_model.py

@@ -30,7 +30,8 @@ class Yolo_model(nn.Module):
             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))
+            nn.Linear(in_features=img_size*img_size*c_hidden, out_features=grid*grid*(boxes*5+labels)),
+            nn.Dropout(0.3)
         )
 
     def forward(self, x):