✨ [New] validation code!

yolo/tools/drawer.py

@@ -37,7 +37,7 @@ def draw_bboxes(
     font = ImageFont.load_default(30)
 
     for bbox in bboxes:
-        class_id, x_min, y_min, x_max, y_max = bbox
+        class_id, x_min, y_min, x_max, y_max, *conf = bbox
         if scaled_bbox:
             x_min = x_min * width
             x_max = x_max * width

yolo/tools/solver.py

@@ -30,6 +30,12 @@ class ModelTrainer:
         self.progress = ProgressTracker(cfg.name, save_path, cfg.use_wandb)
         self.num_epochs = cfg.task.epoch
 
+        validation_dataloader = create_dataloader(cfg.task.validation.data, cfg.dataset, cfg.task.validation.task)
+        anchor2box = AnchorBoxConverter(cfg.model, cfg.image_size, device)
+        self.validator = ModelValidator(
+            cfg.task.validation, model, save_path, device, self.progress, anchor2box, validation_dataloader
+        )
+
         if getattr(train_cfg.ema, "enabled", False):
             self.ema = ExponentialMovingAverage(model, decay=train_cfg.ema.decay)
         else:
@@ -89,9 +95,7 @@ class ModelTrainer:
                 epoch_loss = self.train_one_epoch(dataloader)
                 self.progress.finish_one_epoch()
 
-                logger.info(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}")
-                if (epoch + 1) % 5 == 0:
-                    self.save_checkpoint(epoch, f"checkpoint_epoch_{epoch+1}.pth")
+                self.validator.solve()
 
 
 class ModelTester:
@@ -100,7 +104,7 @@ class ModelTester:
         self.device = device
         self.progress = ProgressTracker(cfg, save_path, cfg.use_wandb)
 
-        self.anchor2box = AnchorBoxConverter(cfg, device)
+        self.anchor2box = AnchorBoxConverter(cfg.model, cfg.image_size, device)
         self.nms = cfg.task.nms
         self.save_path = save_path
 
@@ -125,3 +129,45 @@ class ModelTester:
             else:
                 raise e
         dataloader.stop()
+
+
+class ModelValidator:
+    def __init__(
+        self,
+        validation_cfg: ValidationConfig,
+        model: YOLO,
+        save_path: str,
+        device,
+        progress: ProgressTracker,
+        anchor2box,
+        validation_dataloader,
+    ):
+        self.model = model
+        self.device = device
+        self.progress = progress
+        self.save_path = save_path
+
+        self.anchor2box = anchor2box
+        self.nms = validation_cfg.nms
+        self.validdataloader = validation_dataloader
+
+    def solve(self):
+        # logger.info("🧪 Start Validation!")
+        self.model.eval()
+
+        iou_thresholds = torch.arange(0.5, 1.0, 0.05)
+        map_all = []
+        self.progress.start_one_epoch(len(self.validdataloader))
+        for data, targets in self.validdataloader:
+            data, targets = data.to(self.device), targets.to(self.device)
+            with torch.no_grad():
+                raw_output = self.model(data)
+            predict, _ = self.anchor2box(raw_output[0][3:], with_logits=True)
+
+            nms_out = bbox_nms(predict, self.nms)
+            for idx, predict in enumerate(nms_out):
+                map_value = calculate_map(predict, targets[idx], iou_thresholds)
+                map_all.append(map_value[0])
+            self.progress.one_batch(mapp=torch.Tensor(map_all).mean())
+
+        self.progress.finish_one_epoch()

yolo/utils/bounding_box_utils.py

@@ -297,6 +297,7 @@ def bbox_nms(predicts: Tensor, nms_cfg: NMSConfig):
     cls_val, cls_idx = cls_dist.max(dim=-1, keepdim=True)
     valid_mask = cls_val > nms_cfg.min_confidence
     valid_cls = cls_idx[valid_mask].float()
+    valid_con = cls_val[valid_mask].float()
     valid_box = bbox[valid_mask.repeat(1, 1, 4)].view(-1, 4)
 
     batch_idx, *_ = torch.where(valid_mask)
@@ -305,7 +306,52 @@ def bbox_nms(predicts: Tensor, nms_cfg: NMSConfig):
     for idx in range(predicts.size(0)):
         instance_idx = nms_idx[idx == batch_idx[nms_idx]]
 
-        predict_nms = torch.cat([valid_cls[instance_idx][:, None], valid_box[instance_idx]], dim=-1)
+        predict_nms = torch.cat(
+            [valid_cls[instance_idx][:, None], valid_con[instance_idx][:, None], valid_box[instance_idx]], dim=-1
+        )
 
         predicts_nms.append(predict_nms)
     return predicts_nms
+
+
+def calculate_map(predictions, ground_truths, iou_thresholds):
+    # TODO: Refactor this block
+    device = predictions.device
+    n_preds = predictions.size(0)
+    n_gts = (ground_truths[:, 0] != -1).sum()
+    ground_truths = ground_truths[:n_gts]
+    aps = []
+
+    ious = calculate_iou(predictions[:, 2:], ground_truths[:, 1:])  # [n_preds, n_gts]
+
+    for threshold in iou_thresholds:
+        tp = torch.zeros(n_preds, device=device)
+        fp = torch.zeros(n_preds, device=device)
+
+        max_iou, max_indices = torch.max(ious, dim=1)
+        above_threshold = max_iou >= threshold
+        matched_classes = predictions[:, 0] == ground_truths[max_indices, 0]
+        tp[above_threshold & matched_classes] = 1
+        fp[above_threshold & ~matched_classes] = 1
+        fp[max_iou < threshold] = 1
+
+        _, indices = torch.sort(predictions[:, 1], descending=True)
+        tp = tp[indices]
+        fp = fp[indices]
+
+        tp_cumsum = torch.cumsum(tp, dim=0)
+        fp_cumsum = torch.cumsum(fp, dim=0)
+
+        precision = tp_cumsum / (tp_cumsum + fp_cumsum + 1e-6)
+        recall = tp_cumsum / (n_gts + 1e-6)
+
+        recall_thresholds = torch.arange(0, 1, 0.1)
+        precision_at_recall = torch.zeros_like(recall_thresholds)
+        for i, r in enumerate(recall_thresholds):
+            precision_at_recall[i] = precision[recall >= r].max().item() if torch.any(recall >= r) else 0
+
+        ap = precision_at_recall.mean()
+        aps.append(ap)
+
+    mean_ap = torch.mean(torch.stack(aps))
+    return mean_ap, aps