[113] - Handle batches with no gt, Add images with no gt to the dataset (#120)

yolo/tools/data_loader.py

@@ -87,9 +87,6 @@ class YoloDataset(Dataset):
                     continue
                 annotations = annotations_index.get(image_info["id"], [])
                 image_seg_annotations = scale_segmentation(annotations, image_info)
-                if not image_seg_annotations:
-                    continue
-
             elif data_type == "txt":
                 label_path = labels_path / f"{image_id}.txt"
                 if not label_path.is_file():

yolo/tools/loss_functions.py

@@ -93,7 +93,7 @@ class YOLOLoss:
         targets_cls, targets_bbox = self.separate_anchor(align_targets)
         predicts_box = predicts_box / self.vec2box.scaler[None, :, None]
 
-        cls_norm = targets_cls.sum()
+        cls_norm = max(targets_cls.sum(), 1)
         box_norm = targets_cls.sum(-1)[valid_masks]
 
         ## -- CLS -- ##