loss: 3.7

api/src/yolo/train_yolo_faces.py

@@ -15,7 +15,7 @@ import torch.nn.functional as F
 from production.yolo_model_production import convert_prediction
 
 def view_data(dataset):
-    dataloader = DataLoader(dataset=dataset, batch_size=1, shuffle=False)
+    dataloader = DataLoader(dataset=dataset, batch_size=1, shuffle=True)
     for images, labels in iter(dataloader):
         for batch in range(images.shape[0]):
             image = images[batch]
@@ -112,16 +112,16 @@ def train_yolo():
     BATCH_SIZE = 64
 
     dataset = YoloDataset(
-        image_dir="data/faces_2/train", 
-        annotation_path="data/faces_2/train/_annotations.coco.json",
+        image_dir="data/faces_scenery/train", 
+        annotation_path="data/faces_scenery/train/_annotations.coco.json",
         img_size=IMAGE_SIZE,
         transform=True,
         grid=GRID
     )
 
     dataset_valid = YoloDataset(
-        image_dir="data/faces_2/test", 
-        annotation_path="data/faces_2/test/_annotations.coco.json",
+        image_dir="data/faces_scenery/test", 
+        annotation_path="data/faces_scenery/test/_annotations.coco.json",
         img_size=IMAGE_SIZE,
         transform=True,
         grid=GRID
@@ -140,7 +140,7 @@ def train_yolo():
     loss_module = YoloLoss()
 
 
-    #view_data(dataset)
+    #view_data(dataset_valid)
     #exit()
 
 
@@ -152,8 +152,8 @@ def train_yolo():
     #exit()
 
 
-    use_webcam(GRID, IMAGE_SIZE)
-    exit()
+    #use_webcam(GRID, IMAGE_SIZE)
+    #exit()
 
 
     train(model=model, loss_module=loss_module, train_loader=train_loader, val_loader=val_loader, 

api/src/yolo/yolo_dataset.py

@@ -48,7 +48,7 @@ class YoloDataset(Dataset):
         self.image_ids = list(self.coco.imgs.keys())
         self.img_size = img_size
         self.grid = grid
-        self.num_classes = len(self.coco.cats)-1
+        self.num_classes = len(self.coco.cats)
         self.toTensor = transforms.ToTensor()
 
         if transform:

api/src/yolo/yolo_loss.py

@@ -6,7 +6,7 @@ class YoloLoss(nn.Module):
     def __init__(self):
         super(YoloLoss, self).__init__()
 
-    def forward(self, predictions, targets, lambda_coord=1, lambda_noobj=1):
+    def forward(self, predictions, targets):
         pred_boxes = predictions[..., :4]
         pred_conf = predictions[..., 4]
         pred_classes = predictions[..., 5:]
@@ -27,8 +27,6 @@ class YoloLoss(nn.Module):
         else:
             box_loss = torch.tensor(0.0, device=predictions.device)
 
-        box_loss = lambda_coord * box_loss
-
 
         obj_loss = F.mse_loss(
             pred_conf[obj_mask],
@@ -42,8 +40,6 @@ class YoloLoss(nn.Module):
             reduction="mean"
         ) if noobj_mask.any() else torch.tensor(0.0, device=predictions.device)
 
-        noobj_loss = lambda_noobj * noobj_loss
-
 
         class_loss = F.binary_cross_entropy_with_logits(
             pred_classes[obj_mask],

api/src/yolo/yolo_model.py

@@ -80,10 +80,7 @@ class Yolo_model(nn.Module):
 
     def forward(self, x):
         x = self.model(x).permute(0, 2, 3, 1)
-        center = F.sigmoid(x[..., :2])
-        size = torch.exp(x[..., 2:4])
-        conf_class = F.sigmoid(x[..., 4:])
-        return torch.cat([center, size, conf_class], dim=3)
+        return F.sigmoid(x)
     
 
 def train(model, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, model_name, saving=True):