Use torchmetrics to calculate map

requirements.txt

@@ -5,7 +5,7 @@ loguru
 numpy
 opencv-python
 Pillow
-pycocotools
+torchmetrics
 requests
 rich
 torch

tests/test_tools/test_data_loader.py

@@ -17,40 +17,26 @@ def test_create_dataloader_cache(train_cfg: Config):
 
     make_cache_loader = create_dataloader(train_cfg.task.data, train_cfg.dataset)
     load_cache_loader = create_dataloader(train_cfg.task.data, train_cfg.dataset)
-    m_batch_size, m_images, _, m_reverse_tensors, m_image_paths = next(iter(make_cache_loader))
-    l_batch_size, l_images, _, l_reverse_tensors, l_image_paths = next(iter(load_cache_loader))
+    m_batch_size, m_images, _, m_reverse_tensors = next(iter(make_cache_loader))
+    l_batch_size, l_images, _, l_reverse_tensors = next(iter(load_cache_loader))
     assert m_batch_size == l_batch_size
     assert m_images.shape == l_images.shape
     assert m_reverse_tensors.shape == l_reverse_tensors.shape
-    assert m_image_paths == l_image_paths
 
 
 def test_training_data_loader_correctness(train_dataloader: YoloDataLoader):
     """Test that the training data loader produces correctly shaped data and metadata."""
-    batch_size, images, _, reverse_tensors, image_paths = next(iter(train_dataloader))
+    batch_size, images, _, reverse_tensors = next(iter(train_dataloader))
     assert batch_size == 2
     assert images.shape == (2, 3, 640, 640)
     assert reverse_tensors.shape == (2, 5)
-    expected_paths = [
-        Path("tests/data/images/train/000000050725.jpg"),
-        Path("tests/data/images/train/000000167848.jpg"),
-    ]
-    assert list(image_paths) == list(expected_paths)
-
 
 def test_validation_data_loader_correctness(validation_dataloader: YoloDataLoader):
-    batch_size, images, targets, reverse_tensors, image_paths = next(iter(validation_dataloader))
+    batch_size, images, targets, reverse_tensors = next(iter(validation_dataloader))
     assert batch_size == 4
     assert images.shape == (4, 3, 640, 640)
     assert targets.shape == (4, 18, 5)
     assert reverse_tensors.shape == (4, 5)
-    expected_paths = [
-        Path("tests/data/images/val/000000151480.jpg"),
-        Path("tests/data/images/val/000000284106.jpg"),
-        Path("tests/data/images/val/000000323571.jpg"),
-        Path("tests/data/images/val/000000570456.jpg"),
-    ]
-    assert list(image_paths) == list(expected_paths)
 
 
 def test_file_stream_data_loader_frame(file_stream_data_loader: StreamDataLoader):

tests/test_tools/test_solver.py

@@ -17,7 +17,7 @@
 @pytest.fixture
 def model_validator(validation_cfg: Config, model: YOLO, vec2box: Vec2Box, validation_progress_logger, device):
     validator = ModelValidator(
-        validation_cfg.task, validation_cfg.dataset, model, vec2box, validation_progress_logger, device
+        validation_cfg.task, model, vec2box, validation_progress_logger, device
     )
     return validator
 
@@ -28,11 +28,10 @@ def test_model_validator_initialization(model_validator: ModelValidator):
 
 
 def test_model_validator_solve_mock_dataset(model_validator: ModelValidator, validation_dataloader: YoloDataLoader):
-    mAPs = model_validator.solve(validation_dataloader)
-    except_mAPs = {"mAP.5": tensor(0.6969), "mAP.5:.95": tensor(0.4195)}
-    assert allclose(mAPs["mAP.5"], except_mAPs["mAP.5"], rtol=0.1)
-    print(mAPs)
-    assert allclose(mAPs["mAP.5:.95"], except_mAPs["mAP.5:.95"], rtol=0.1)
+    metrics = model_validator.solve(validation_dataloader)
+    except_metrics = {"map_50": tensor(0.7515), "map": tensor(0.5986)}
+    assert allclose(metrics["map_50"], except_metrics["map_50"], rtol=0.1)
+    assert allclose(metrics["map"], except_metrics["map"], rtol=0.1)
 
 
 @pytest.fixture

tests/test_utils/test_bounding_box_utils.py

@@ -15,7 +15,6 @@
     Vec2Box,
     bbox_nms,
     calculate_iou,
-    calculate_map,
     generate_anchors,
     transform_bbox,
 )
@@ -167,17 +166,4 @@ def test_bbox_nms():
     output = bbox_nms(cls_dist, bbox, nms_cfg)
 
     for out, exp in zip(output, expected_output):
-        assert allclose(out, exp, atol=1e-4), f"Output: {out} Expected: {exp}"
-
-
-def test_calculate_map():
-    predictions = tensor([[0, 60, 60, 160, 160, 0.5], [0, 40, 40, 120, 120, 0.5]])  # [class, x1, y1, x2, y2]
-    ground_truths = tensor([[0, 50, 50, 150, 150], [0, 30, 30, 100, 100]])  # [class, x1, y1, x2, y2]
-
-    mAP = calculate_map(predictions, ground_truths)
-
-    expected_ap50 = tensor(0.5)
-    expected_ap50_95 = tensor(0.2)
-
-    assert isclose(mAP["mAP.5"], expected_ap50, atol=1e-5), f"AP50 mismatch"
-    assert isclose(mAP["mAP.5:.95"], expected_ap50_95, atol=1e-5), f"Mean AP mismatch"
+        assert allclose(out, exp, atol=1e-4), f"Output: {out} Expected: {exp}"

yolo/lazy.py

@@ -26,13 +26,12 @@ def main(cfg: Config):
     else:
         model = create_model(cfg.model, class_num=cfg.dataset.class_num, weight_path=cfg.weight)
         model = model.to(device)
-
     converter = create_converter(cfg.model.name, model, cfg.model.anchor, cfg.image_size, device)
 
     if cfg.task.task == "train":
         solver = ModelTrainer(cfg, model, converter, progress, device, use_ddp)
     if cfg.task.task == "validation":
-        solver = ModelValidator(cfg.task, cfg.dataset, model, converter, progress, device)
+        solver = ModelValidator(cfg.task, model, converter, progress, device)
     if cfg.task.task == "inference":
         solver = ModelTester(cfg, model, converter, progress, device)
     progress.start()

yolo/model/yolo.py

@@ -123,7 +123,7 @@ def save_load_weights(self, weights: Union[Path, OrderedDict]):
             weights: A OrderedDict containing the new weights.
         """
         if isinstance(weights, Path):
-            weights = torch.load(weights, map_location=torch.device("cpu"))
+            weights = torch.load(weights, map_location=torch.device("cpu"), weights_only=False)
         if "model_state_dict" in weights:
             weights = weights["model_state_dict"]
 

yolo/tools/data_loader.py

@@ -134,16 +134,16 @@ def load_valid_labels(self, label_path: str, seg_data_one_img: list) -> Union[Te
     def get_data(self, idx):
         img_path, bboxes = self.data[idx]
         img = Image.open(img_path).convert("RGB")
-        return img, bboxes, img_path
+        return img, bboxes
 
     def get_more_data(self, num: int = 1):
         indices = torch.randint(0, len(self), (num,))
-        return [self.get_data(idx)[:2] for idx in indices]
+        return [self.get_data(idx) for idx in indices]
 
     def __getitem__(self, idx) -> Tuple[Image.Image, Tensor, Tensor, List[str]]:
-        img, bboxes, img_path = self.get_data(idx)
+        img, bboxes = self.get_data(idx)
         img, bboxes, rev_tensor = self.transform(img, bboxes)
-        return img, bboxes, rev_tensor, img_path
+        return img, bboxes, rev_tensor
 
     def __len__(self) -> int:
         return len(self.data)
@@ -189,11 +189,11 @@ def collate_fn(self, batch: List[Tuple[Tensor, Tensor]]) -> Tuple[Tensor, List[T
             batch_targets[idx, : min(target_size, 100)] = batch[idx][1][:100]
         batch_targets[:, :, 1:] *= self.image_size
 
-        batch_images, _, batch_reverse, batch_path = zip(*batch)
+        batch_images, _, batch_reverse = zip(*batch)
         batch_images = torch.stack(batch_images)
         batch_reverse = torch.stack(batch_reverse)
 
-        return batch_size, batch_images, batch_targets, batch_reverse, batch_path
+        return batch_size, batch_images, batch_targets, batch_reverse
 
 
 def create_dataloader(data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: str = "train", use_ddp: bool = False):

yolo/tools/solver.py

@@ -1,37 +1,30 @@
-import contextlib
-import io
-import json
 import os
 import time
 from collections import defaultdict
-from pathlib import Path
 from typing import Dict, Optional
 
 import torch
 from loguru import logger
-from pycocotools.coco import COCO
 from torch import Tensor, distributed
 from torch.cuda.amp import GradScaler, autocast
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data import DataLoader
+from torchmetrics.detection import MeanAveragePrecision
 
 from yolo.config.config import Config, DatasetConfig, TrainConfig, ValidationConfig
 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, calculate_map
-from yolo.utils.dataset_utils import locate_label_paths
+from yolo.utils.bounding_box_utils import Vec2Box
 from yolo.utils.logging_utils import ProgressLogger, log_model_structure
+from yolo.utils.solver_utils import format_prediction, format_target
 from yolo.utils.model_utils import (
     ExponentialMovingAverage,
     PostProccess,
-    collect_prediction,
     create_optimizer,
     create_scheduler,
-    predicts_to_json,
 )
-from yolo.utils.solver_utils import calculate_ap
 
 
 class ModelTrainer:
@@ -58,7 +51,7 @@ def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, progress: Progres
         self.validation_dataloader = create_dataloader(
             cfg.task.validation.data, cfg.dataset, cfg.task.validation.task, use_ddp
         )
-        self.validator = ModelValidator(cfg.task.validation, cfg.dataset, model, vec2box, progress, device)
+        self.validator = ModelValidator(cfg.task.validation, model, vec2box, progress, device)
 
         if getattr(train_cfg.ema, "enabled", False):
             self.ema = ExponentialMovingAverage(model, decay=train_cfg.ema.decay)
@@ -89,7 +82,7 @@ def train_one_epoch(self, dataloader):
         total_loss = defaultdict(float)
         total_samples = 0
         self.optimizer.next_epoch(len(dataloader))
-        for batch_size, images, targets, *_ in dataloader:
+        for batch_size, images, targets, _ in dataloader:
             self.optimizer.next_batch()
             loss_each = self.train_one_batch(images, targets)
 
@@ -213,7 +206,6 @@ class ModelValidator:
     def __init__(
         self,
         validation_cfg: ValidationConfig,
-        dataset_cfg: DatasetConfig,
         model: YOLO,
         vec2box: Vec2Box,
         progress: ProgressLogger,
@@ -222,46 +214,28 @@ def __init__(
         self.model = model
         self.device = device
         self.progress = progress
-
         self.post_proccess = PostProccess(vec2box, validation_cfg.nms)
-        self.json_path = self.progress.save_path / "predict.json"
-
-        with contextlib.redirect_stdout(io.StringIO()):
-            # TODO: load with config file
-            json_path, _ = locate_label_paths(Path(dataset_cfg.path), dataset_cfg.get("validation", "val"))
-            if json_path:
-                self.coco_gt = COCO(json_path)
 
     def solve(self, dataloader, epoch_idx=1):
-        # logger.info("🧪 Start Validation!")
+        logger.info("🧪 Start Validation!")
+        metric = MeanAveragePrecision(iou_type="bbox", box_format="xyxy")
         self.model.eval()
-        predict_json, mAPs = [], defaultdict(list)
         self.progress.start_one_epoch(len(dataloader), task="Validate")
-        for batch_size, images, targets, rev_tensor, img_paths in dataloader:
+        for _, images, targets, rev_tensor 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)
-                for idx, predict in enumerate(predicts):
-                    mAP = calculate_map(predict, targets[idx])
-                    for mAP_key, mAP_val in mAP.items():
-                        mAPs[mAP_key].append(mAP_val)
-
-            avg_mAPs = {key: 100 * torch.mean(torch.stack(val)) for key, val in mAPs.items()}
-            self.progress.one_batch(avg_mAPs)
-
-            predict_json.extend(predicts_to_json(img_paths, predicts, rev_tensor))
-        self.progress.finish_one_epoch(avg_mAPs, epoch_idx=epoch_idx)
+                batch_metrics = metric([format_prediction(predict) for predict in predicts], 
+                       [format_target(target) for target in targets])
+
+            self.progress.one_batch({
+                "map": batch_metrics["map"],
+                "map_50": batch_metrics["map_50"],
+            })
+
+        epoch_metrics = metric.compute()
+        del epoch_metrics['classes']
+        self.progress.finish_one_epoch(epoch_metrics, epoch_idx=epoch_idx)
         self.progress.visualize_image(images, targets, predicts, epoch_idx=epoch_idx)
-
-        with open(self.json_path, "w") as f:
-            predict_json = collect_prediction(predict_json, self.progress.local_rank)
-            if self.progress.local_rank != 0:
-                return
-            json.dump(predict_json, f)
-        if hasattr(self, "coco_gt"):
-            self.progress.start_pycocotools()
-            result = calculate_ap(self.coco_gt, predict_json)
-            self.progress.finish_pycocotools(result, epoch_idx)
-
-        return avg_mAPs
+        return epoch_metrics

yolo/utils/bounding_box_utils.py

@@ -227,7 +227,19 @@ def __call__(self, target: Tensor, predict: Tuple[Tensor]) -> Tuple[Tensor, Tens
         2. Select the targets
         2. Noramlize the class probilities of targets
         """
+
         predict_cls, predict_bbox = predict
+
+        # return if target has no gt information.
+        n_targets = target.shape[1]
+        if n_targets == 0:
+            device = predict_bbox.device
+            align_cls = torch.zeros_like(predict_cls, device=device)
+            align_bbox = torch.zeros_like(predict_bbox, device=device)
+            valid_mask = torch.zeros(predict_cls.shape[:2], dtype=bool, device=device)
+            anchor_matched_targets = torch.cat([align_cls, align_bbox], dim=-1)
+            return anchor_matched_targets, valid_mask
+
         target_cls, target_bbox = target.split([1, 4], dim=-1)  # B x N x (C B) -> B x N x C, B x N x B
         target_cls = target_cls.long().clamp(0)
 
@@ -262,7 +274,7 @@ def __call__(self, target: Tensor, predict: Tuple[Tensor]) -> Tuple[Tensor, Tens
         normalize_term = normalize_term.permute(0, 2, 1).gather(2, unique_indices)
         align_cls = align_cls * normalize_term * valid_mask[:, :, None]
 
-        return torch.cat([align_cls, align_bbox], dim=-1), valid_mask.bool()
+        return torch.cat([align_cls, align_bbox], dim=-1), valid_mask
 
 
 class Vec2Box:
@@ -398,51 +410,4 @@ def bbox_nms(cls_dist: Tensor, bbox: Tensor, nms_cfg: NMSConfig, confidence: Opt
         )
 
         predicts_nms.append(predict_nms)
-    return predicts_nms
-
-
-def calculate_map(predictions, ground_truths, iou_thresholds=arange(0.5, 1, 0.05)) -> Dict[str, Tensor]:
-    # 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()
-    ground_truths = ground_truths[:n_gts]
-    aps = []
-
-    ious = calculate_iou(predictions[:, 1:-1], ground_truths[:, 1:])  # [n_preds, n_gts]
-
-    for threshold in iou_thresholds:
-        tp = torch.zeros(n_preds, device=device, dtype=bool)
-
-        max_iou, max_indices = ious.max(dim=1)
-        above_threshold = max_iou >= threshold
-        matched_classes = predictions[:, 0] == ground_truths[max_indices, 0]
-        max_match = torch.zeros_like(ious)
-        max_match[arange(n_preds), max_indices] = max_iou
-        if max_match.size(0):
-            tp[max_match.argmax(dim=0)] = True
-        tp[~above_threshold | ~matched_classes] = False
-
-        _, indices = torch.sort(predictions[:, 1], descending=True)
-        tp = tp[indices]
-
-        tp_cumsum = torch.cumsum(tp, dim=0)
-        fp_cumsum = torch.cumsum(~tp, dim=0)
-
-        precision = tp_cumsum / (tp_cumsum + fp_cumsum + 1e-6)
-        recall = tp_cumsum / (n_gts + 1e-6)
-
-        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)
-
-        ap = torch.trapezoid(precision, recall)
-        aps.append(ap)
-
-    mAP = {
-        "mAP.5": aps[0],
-        "mAP.5:.95": torch.mean(torch.stack(aps)),
-    }
-    return mAP
+    return predicts_nms