✅ [Update] type hint and mAP calculation func

tests/conftest.py

@@ -4,15 +4,16 @@ from pathlib import Path
 import pytest
 import torch
 from hydra import compose, initialize
+from lightning import Trainer
 
 project_root = Path(__file__).resolve().parent.parent
 sys.path.append(str(project_root))
 
 from yolo import Anc2Box, Config, Vec2Box, create_converter, create_model
 from yolo.model.yolo import YOLO
-from yolo.tools.data_loader import StreamDataLoader, YoloDataLoader
+from yolo.tools.data_loader import StreamDataLoader, create_dataloader
 from yolo.tools.dataset_preparation import prepare_dataset
-from yolo.utils.logging_utils import ProgressLogger, set_seed
+from yolo.utils.logging_utils import set_seed, setup
 
 
 def pytest_configure(config):
@@ -52,18 +53,6 @@ def device():
     return torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
-@pytest.fixture(scope="session")
-def train_progress_logger(train_cfg: Config):
-    progress_logger = ProgressLogger(train_cfg, exp_name=train_cfg.name)
-    return progress_logger
-
-
-@pytest.fixture(scope="session")
-def validation_progress_logger(validation_cfg: Config):
-    progress_logger = ProgressLogger(validation_cfg, exp_name=validation_cfg.name)
-    return progress_logger
-
-
 @pytest.fixture(scope="session")
 def model(train_cfg: Config, device) -> YOLO:
     model = create_model(train_cfg.model)
@@ -76,6 +65,22 @@ def model_v7(inference_v7_cfg: Config, device) -> YOLO:
     return model.to(device)
 
 
+@pytest.fixture(scope="session")
+def solver(train_cfg: Config) -> Trainer:
+    callbacks, loggers = setup(train_cfg)
+    trainer = Trainer(
+        accelerator="cuda",
+        max_epochs=getattr(train_cfg.task, "epoch", None),
+        precision="16-mixed",
+        callbacks=callbacks,
+        logger=loggers,
+        log_every_n_steps=1,
+        gradient_clip_val=10,
+        deterministic=True,
+    )
+    return trainer
+
+
 @pytest.fixture(scope="session")
 def vec2box(train_cfg: Config, model: YOLO, device) -> Vec2Box:
     vec2box = create_converter(train_cfg.model.name, model, train_cfg.model.anchor, train_cfg.image_size, device)
@@ -93,13 +98,13 @@ def anc2box(inference_v7_cfg: Config, model: YOLO, device) -> Anc2Box:
 @pytest.fixture(scope="session")
 def train_dataloader(train_cfg: Config):
     prepare_dataset(train_cfg.dataset, task="train")
-    return YoloDataLoader(train_cfg.task.data, train_cfg.dataset, train_cfg.task.task)
+    return create_dataloader(train_cfg.task.data, train_cfg.dataset, train_cfg.task.task)
 
 
 @pytest.fixture(scope="session")
 def validation_dataloader(validation_cfg: Config):
     prepare_dataset(validation_cfg.dataset, task="val")
-    return YoloDataLoader(validation_cfg.task.data, validation_cfg.dataset, validation_cfg.task.task)
+    return create_dataloader(validation_cfg.task.data, validation_cfg.dataset, validation_cfg.task.task)
 
 
 @pytest.fixture(scope="session")

tests/test_tools/test_loss_functions.py

@@ -1,4 +1,5 @@
 import sys
+from math import isinf, isnan
 from pathlib import Path
 
 import pytest
@@ -51,6 +52,6 @@ def test_yolo_loss(loss_function, data):
     predicts, targets = data
     loss, loss_dict = loss_function(predicts, predicts, targets)
     assert torch.isnan(loss)
-    assert torch.isnan(loss_dict["BoxLoss"])
-    assert torch.isnan(loss_dict["DFLoss"])
-    assert torch.isinf(loss_dict["BCELoss"])
+    assert isnan(loss_dict["Loss/BoxLoss"])
+    assert isnan(loss_dict["Loss/DFLLoss"])
+    assert isinf(loss_dict["Loss/BCELoss"])

tests/test_utils/test_bounding_box_utils.py

@@ -216,9 +216,8 @@ def test_calculate_map():
     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.5050)
+    expected_ap50_95 = tensor(0.2020)
 
-    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 isclose(mAP["map_50"], expected_ap50, atol=1e-4), f"AP50 mismatch"
+    assert isclose(mAP["map"], expected_ap50_95, atol=1e-4), f"Mean AP mismatch"

yolo/tools/loss_functions.py

@@ -119,7 +119,7 @@ class DualLoss:
 
     def __call__(
         self, aux_predicts: List[Tensor], main_predicts: List[Tensor], targets: Tensor
-    ) -> Tuple[Tensor, Dict[str, Tensor]]:
+    ) -> Tuple[Tensor, Dict[str, float]]:
         # TODO: Need Refactor this region, make it flexible!
         aux_iou, aux_dfl, aux_cls = self.loss(aux_predicts, targets)
         main_iou, main_dfl, main_cls = self.loss(main_predicts, targets)

yolo/utils/bounding_box_utils.py

@@ -4,10 +4,11 @@ from typing import Dict, List, Optional, Tuple, Union
 import torch
 import torch.nn.functional as F
 from einops import rearrange
-from torch import Tensor, arange, tensor
+from torch import Tensor, tensor
+from torchmetrics.detection import MeanAveragePrecision
 from torchvision.ops import batched_nms
 
-from yolo.config.config import AnchorConfig, MatcherConfig, ModelConfig, NMSConfig
+from yolo.config.config import AnchorConfig, MatcherConfig, NMSConfig
 from yolo.model.yolo import YOLO
 from yolo.utils.logger import logger
 
@@ -406,50 +407,9 @@ def bbox_nms(cls_dist: Tensor, bbox: Tensor, nms_cfg: NMSConfig, confidence: Opt
     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)),
-    }
+def calculate_map(predictions, ground_truths) -> Dict[str, Tensor]:
+    metric = MeanAveragePrecision(iou_type="bbox", box_format="xyxy")
+    mAP = metric([to_metrics_format(predictions)], [to_metrics_format(ground_truths)])
     return mAP