✨ [New] local mAP calculation, AP.5:.95 - mAP.5

yolo/config/task/validation.yaml

@@ -8,5 +8,5 @@ data:
   pin_memory: True
   data_augment: {}
 nms:
-  min_confidence: 0.001
-  min_iou: 0.7
+  min_confidence: 0.05
+  min_iou: 0.9

yolo/tools/solver.py

@@ -18,7 +18,7 @@ from yolo.model.yolo import YOLO
 from yolo.tools.data_loader import StreamDataLoader, create_dataloader
 from yolo.tools.drawer import draw_bboxes, draw_model
 from yolo.tools.loss_functions import create_loss_function
-from yolo.utils.bounding_box_utils import Vec2Box
+from yolo.utils.bounding_box_utils import Vec2Box, calculate_map
 from yolo.utils.logging_utils import ProgressLogger, log_model_structure
 from yolo.utils.model_utils import (
     ExponentialMovingAverage,
@@ -198,16 +198,18 @@ class ModelValidator:
     def solve(self, dataloader, epoch_idx=-1):
         # logger.info("🧪 Start Validation!")
         self.model.eval()
-        predict_json = []
+        mAPs, predict_json = [], []
         self.progress.start_one_epoch(len(dataloader))
         for images, targets, rev_tensor, img_paths in dataloader:
             images, targets, rev_tensor = images.to(self.device), targets.to(self.device), rev_tensor.to(self.device)
             with torch.no_grad():
                 predicts = self.model(images)
-                predicts = self.post_proccess(predicts, rev_tensor)
-            self.progress.one_batch()
+                predicts = self.post_proccess(predicts)
+                for idx, predict in enumerate(predicts):
+                    mAPs.append(calculate_map(predict, targets[idx]))
+            self.progress.one_batch(mAP=Tensor(mAPs))
 
-            predict_json.extend(predicts_to_json(img_paths, predicts))
+            predict_json.extend(predicts_to_json(img_paths, predicts, rev_tensor))
         self.progress.finish_one_epoch()
         with open(self.json_path, "w") as f:
             json.dump(predict_json, f)

yolo/utils/bounding_box_utils.py

@@ -5,7 +5,7 @@ import torch
 import torch.nn.functional as F
 from einops import rearrange
 from loguru import logger
-from torch import Tensor
+from torch import Tensor, arange
 from torchvision.ops import batched_nms
 
 from yolo.config.config import MatcherConfig, ModelConfig, NMSConfig
@@ -338,8 +338,8 @@ def bbox_nms(cls_dist: Tensor, bbox: Tensor, nms_cfg: NMSConfig):
     return predicts_nms
 
 
-def calculate_map(predictions, ground_truths, iou_thresholds):
-    # TODO: Refactor this block
+def calculate_map(predictions, ground_truths, iou_thresholds=arange(0.5, 1, 0.05)):
+    # TODO: Refactor this block, Flexible for calculate different mAP condition?
     device = predictions.device
     n_preds = predictions.size(0)
     n_gts = (ground_truths[:, 0] != -1).sum()
@@ -369,13 +369,16 @@ def calculate_map(predictions, ground_truths, iou_thresholds):
         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
+        precision = torch.cat([torch.ones(1, device=device), precision, torch.zeros(1, device=device)])
+        recall = torch.cat([torch.zeros(1, device=device), recall, torch.ones(1, device=device)])
+
+        precision, _ = torch.cummax(precision.flip(0), dim=0)
+        precision = precision.flip(0)
+
+        indices = (recall[1:] != recall[:-1]).nonzero(as_tuple=True)[0]
+        ap = torch.sum((recall[indices + 1] - recall[indices]) * precision[indices + 1])
 
-        ap = precision_at_recall.mean()
         aps.append(ap)
 
     mean_ap = torch.mean(torch.stack(aps))
-    return mean_ap, aps
+    return mean_ap, aps[0]

yolo/utils/logging_utils.py

@@ -95,9 +95,11 @@ class ProgressLogger(Progress):
                 self.wandb.log({f"Learning Rate/{lr_name}": lr_value}, step=epoch_idx)
         self.batch_task = self.add_task("[green]Batches", total=num_batches)
 
-    def one_batch(self, loss_dict: Dict[str, Tensor] = None):
+    def one_batch(self, loss_dict: Dict[str, Tensor] = None, mAP: Tensor = None):
         if loss_dict is None:
-            self.update(self.batch_task, advance=1, description=f"[green]Validating")
+            # refactor this block & class
+            mAP_50, mAP_50_95 = mAP.mean(0)
+            self.update(self.batch_task, advance=1, description=f"[green]Validating {mAP_50: .2f} {mAP_50_95: .2f}")
             return
         if self.use_wandb:
             for loss_name, loss_value in loss_dict.items():

yolo/utils/model_utils.py

@@ -106,7 +106,7 @@ class PostProccess:
         self.vec2box = vec2box
         self.nms = nms_cfg
 
-    def __call__(self, predict, rev_tensor: Optional[Tensor]):
+    def __call__(self, predict, rev_tensor: Optional[Tensor] = None):
         pred_class, _, pred_bbox = self.vec2box(predict["Main"])
         if rev_tensor is not None:
             pred_bbox = (pred_bbox - rev_tensor[:, None, 1:]) / rev_tensor[:, 0:1, None]
@@ -114,13 +114,15 @@ class PostProccess:
         return pred_bbox
 
 
-def predicts_to_json(img_paths, predicts):
+def predicts_to_json(img_paths, predicts, rev_tensor):
     """
     TODO: function document
     turn a batch of imagepath and predicts(n x 6 for each image) to a List of diction(Detection output)
     """
     batch_json = []
-    for img_path, bboxes in zip(img_paths, predicts):
+    for img_path, bboxes, box_reverse in zip(img_paths, predicts, rev_tensor):
+        scale, shift = box_reverse.split([1, 4])
+        bboxes[:, 1:5] = (bboxes[:, 1:5] - shift[None]) / scale[None]
         bboxes[:, 1:5] = transform_bbox(bboxes[:, 1:5], "xyxy -> xywh")
         for cls, *pos, conf in bboxes:
             bbox = {