✅ [Add] tests, increase test coverage

requirements.txt

@@ -5,6 +5,7 @@ loguru
 numpy
 opencv-python
 Pillow
+pycocotools
 pytest
 pyyaml
 requests

tests/test_model/test_module.py

@@ -43,13 +43,6 @@ def test_adown():
     assert out.shape == (1, OUT_CHANNELS, 32, 32)
 
 
-def test_adown():
-    adown = ADown(IN_CHANNELS, OUT_CHANNELS)
-    x = torch.randn(1, IN_CHANNELS, 64, 64)
-    out = adown(x)
-    assert out.shape == (1, OUT_CHANNELS, 32, 32)
-
-
 def test_cblinear():
     cblinear = CBLinear(IN_CHANNELS, [5, 5])
     x = torch.randn(1, IN_CHANNELS, 64, 64)

tests/test_model/test_yolo.py

@@ -16,7 +16,7 @@ config_path = "../../yolo/config"
 config_name = "config"
 
 
-def test_build_model():
+def test_build_model_v9c():
     with initialize(config_path=config_path, version_base=None):
         cfg: Config = compose(config_name=config_name)
 
@@ -26,6 +26,16 @@ def test_build_model():
         assert len(model.model) == 39
 
 
+def test_build_model_v9m():
+    with initialize(config_path=config_path, version_base=None):
+        cfg: Config = compose(config_name=config_name, overrides=[f"model=v9-m"])
+
+        OmegaConf.set_struct(cfg.model, False)
+        cfg.weight = None
+        model = YOLO(cfg.model)
+        assert len(model.model) == 39
+
+
 @pytest.fixture
 def cfg() -> Config:
     with initialize(config_path="../../yolo/config", version_base=None):

tests/{test_utils → test_tools}/test_loss_functions.py

@@ -31,7 +31,6 @@ def model(cfg: Config):
 @pytest.fixture
 def vec2box(cfg: Config, model):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    print(device)
     return Vec2Box(model, cfg.image_size, device)
 
 

tests/test_tools/test_solver.py

@@ -0,0 +1,114 @@
+import sys
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from hydra import compose, initialize
+
+project_root = Path(__file__).resolve().parent.parent.parent
+sys.path.append(str(project_root))
+
+from yolo.config.config import (
+    Config,
+    DataConfig,
+    LossConfig,
+    TrainConfig,
+    ValidationConfig,
+)
+from yolo.model.yolo import YOLO, create_model
+from yolo.tools.data_loader import create_dataloader
+from yolo.tools.loss_functions import create_loss_function
+from yolo.tools.solver import (  # Adjust the import to your module
+    ModelTester,
+    ModelTrainer,
+    ModelValidator,
+)
+from yolo.utils.bounding_box_utils import Vec2Box
+from yolo.utils.logging_utils import ProgressLogger
+from yolo.utils.model_utils import (
+    ExponentialMovingAverage,
+    create_optimizer,
+    create_scheduler,
+)
+
+
+@pytest.fixture
+def cfg() -> Config:
+    with initialize(config_path="../../yolo/config", version_base=None):
+        cfg: Config = compose(config_name="config")
+        cfg.weight = None
+    return cfg
+
+
+@pytest.fixture
+def cfg_validaion() -> Config:
+    with initialize(config_path="../../yolo/config", version_base=None):
+        cfg: Config = compose(config_name="config", overrides=["task=validation"])
+        cfg.weight = None
+    return cfg
+
+
+@pytest.fixture
+def cfg_inference() -> Config:
+    with initialize(config_path="../../yolo/config", version_base=None):
+        cfg: Config = compose(config_name="config", overrides=["task=inference"])
+        cfg.weight = None
+    return cfg
+
+
+@pytest.fixture
+def device() -> torch.device:
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return device
+
+
+@pytest.fixture
+def model(cfg: Config, device) -> YOLO:
+    model = create_model(cfg.model, weight_path=None)
+    return model.to(device)
+
+
+@pytest.fixture
+def vec2box(cfg: Config, model: YOLO, device) -> Vec2Box:
+    model = create_model(cfg.model, weight_path=None).to(device)
+    vec2box = Vec2Box(model, cfg.image_size, device)
+    return vec2box
+
+
+@pytest.fixture
+def progress_logger(cfg: Config):
+    progress_logger = ProgressLogger(cfg, exp_name=cfg.name)
+    return progress_logger
+
+
+def test_model_trainer_initialization(cfg: Config, model: YOLO, vec2box: Vec2Box, progress_logger, device):
+    trainer = ModelTrainer(cfg, model, vec2box, progress_logger, device, use_ddp=False)
+    assert trainer.model == model
+    assert trainer.device == device
+    assert trainer.optimizer is not None
+    assert trainer.scheduler is not None
+    assert trainer.loss_fn is not None
+    assert trainer.progress == progress_logger
+
+
+# def test_model_trainer_train_one_batch(config, model, vec2box, progress_logger, device):
+#     trainer = ModelTrainer(config, model, vec2box, progress_logger, device, use_ddp=False)
+#     images = torch.rand(1, 3, 224, 224)
+#     targets = torch.rand(1, 5)
+#     loss_item = trainer.train_one_batch(images, targets)
+#     assert isinstance(loss_item, dict)
+
+
+def test_model_validator_initialization(cfg_validaion: Config, model: YOLO, vec2box: Vec2Box, progress_logger, device):
+    validator = ModelValidator(cfg_validaion.task, model, vec2box, progress_logger, device)
+    assert validator.model == model
+    assert validator.device == device
+    assert validator.progress == progress_logger
+
+
+def test_model_tester_initialization(cfg_inference: Config, model: YOLO, vec2box: Vec2Box, progress_logger, device):
+    tester = ModelTester(cfg_inference, model, vec2box, progress_logger, device)
+    assert tester.model == model
+    assert tester.device == device
+    assert tester.progress == progress_logger

tests/test_utils/test_bounding_box_utils.py

@@ -0,0 +1,163 @@
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+from hydra import compose, initialize
+from torch import Tensor, allclose, float32, isclose, nn, tensor
+
+project_root = Path(__file__).resolve().parent.parent.parent
+sys.path.append(str(project_root))
+from yolo import Config, NMSConfig, create_model
+from yolo.utils.bounding_box_utils import (
+    Vec2Box,
+    bbox_nms,
+    calculate_iou,
+    calculate_map,
+    generate_anchors,
+    transform_bbox,
+)
+
+EPS = 1e-4
+
+
+@pytest.fixture
+def dummy_bboxes():
+    bbox1 = tensor([[50, 80, 150, 140], [30, 20, 100, 80]], dtype=float32)
+    bbox2 = tensor([[90, 70, 160, 160], [40, 40, 90, 120]], dtype=float32)
+    return bbox1, bbox2
+
+
+def test_calculate_iou_2d(dummy_bboxes):
+    bbox1, bbox2 = dummy_bboxes
+    iou = calculate_iou(bbox1, bbox2)
+    expected_iou = tensor([[0.4138, 0.1905], [0.0096, 0.3226]])
+    assert iou.shape == (2, 2)
+    assert allclose(iou, expected_iou, atol=EPS)
+
+
+def test_calculate_iou_3d(dummy_bboxes):
+    bbox1, bbox2 = dummy_bboxes
+    iou = calculate_iou(bbox1[None], bbox2[None])
+    expected_iou = tensor([[0.4138, 0.1905], [0.0096, 0.3226]])
+    assert iou.shape == (1, 2, 2)
+    assert allclose(iou, expected_iou, atol=EPS)
+
+
+def test_calculate_diou(dummy_bboxes):
+    bbox1, bbox2 = dummy_bboxes
+    iou = calculate_iou(bbox1, bbox2, "diou")
+    expected_diou = tensor([[0.3816, 0.0943], [-0.2048, 0.2622]])
+
+    assert iou.shape == (2, 2)
+    assert allclose(iou, expected_diou, atol=EPS)
+
+
+def test_calculate_ciou(dummy_bboxes):
+    bbox1, bbox2 = dummy_bboxes
+    iou = calculate_iou(bbox1, bbox2, metrics="ciou")
+    # TODO: check result!
+    expected_ciou = tensor([[0.3769, 0.0853], [-0.2050, 0.2602]])
+    assert iou.shape == (2, 2)
+    assert allclose(iou, expected_ciou, atol=EPS)
+
+    bbox1 = tensor([[50, 80, 150, 140], [30, 20, 100, 80]], dtype=float32)
+    bbox2 = tensor([[90, 70, 160, 160], [40, 40, 90, 120]], dtype=float32)
+
+
+def test_transform_bbox_xywh_to_Any(dummy_bboxes):
+    bbox1, _ = dummy_bboxes
+    transformed_bbox = transform_bbox(bbox1, "xywh -> xyxy")
+    expected_bbox = tensor([[50.0, 80.0, 200.0, 220.0], [30.0, 20.0, 130.0, 100.0]])
+    assert allclose(transformed_bbox, expected_bbox)
+
+
+def test_transform_bbox_xycwh_to_Any(dummy_bboxes):
+    bbox1, bbox2 = dummy_bboxes
+    transformed_bbox = transform_bbox(bbox1, "xycwh -> xycwh")
+    assert allclose(transformed_bbox, bbox1)
+
+    transformed_bbox = transform_bbox(bbox2, "xyxy -> xywh")
+    expected_bbox = tensor([[90.0, 70.0, 70.0, 90.0], [40.0, 40.0, 50.0, 80.0]])
+    assert allclose(transformed_bbox, expected_bbox)
+
+
+def test_transform_bbox_xyxy_to_Any(dummy_bboxes):
+    bbox1, bbox2 = dummy_bboxes
+    transformed_bbox = transform_bbox(bbox1, "xyxy -> xyxy")
+    assert allclose(transformed_bbox, bbox1)
+
+    transformed_bbox = transform_bbox(bbox2, "xyxy -> xycwh")
+    expected_bbox = tensor([[125.0, 115.0, 70.0, 90.0], [65.0, 80.0, 50.0, 80.0]])
+    assert allclose(transformed_bbox, expected_bbox)
+
+
+def test_transform_bbox_invalid_format(dummy_bboxes):
+    bbox, _ = dummy_bboxes
+
+    # Test invalid input format
+    with pytest.raises(ValueError, match="Invalid input or output format"):
+        transform_bbox(bbox, "invalid->xyxy")
+
+    # Test invalid output format
+    with pytest.raises(ValueError, match="Invalid input or output format"):
+        transform_bbox(bbox, "xywh->invalid")
+
+
+def test_generate_anchors():
+    image_size = [256, 256]
+    strides = [8, 16, 32]
+    anchors, scalers = generate_anchors(image_size, strides)
+    assert anchors.shape[0] == scalers.shape[0]
+    assert anchors.shape[1] == 2
+
+
+def test_vec2box_autoanchor():
+    with initialize(config_path="../../yolo/config", version_base=None):
+        cfg: Config = compose(config_name="config", overrides=["model=v9-m"])
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = create_model(cfg.model, weight_path=None).to(device)
+    vec2box = Vec2Box(model, cfg.image_size, device)
+    assert vec2box.strides == [8, 16, 32]
+
+    vec2box.update((320, 640))
+    assert vec2box.anchor_grid.shape == (4200, 2)
+    assert vec2box.scaler.shape == tuple([4200])
+
+
+def test_bbox_nms():
+    cls_dist = tensor(
+        [[[0.1, 0.7, 0.2], [0.6, 0.3, 0.1]], [[0.4, 0.4, 0.2], [0.5, 0.4, 0.1]]]  # Example class distribution
+    )
+    bbox = tensor(
+        [[[50, 50, 100, 100], [60, 60, 110, 110]], [[40, 40, 90, 90], [70, 70, 120, 120]]],  # Example bounding boxes
+        dtype=float32,
+    )
+    nms_cfg = NMSConfig(min_confidence=0.5, min_iou=0.5)
+
+    expected_output = [
+        tensor(
+            [
+                [1.0000, 50.0000, 50.0000, 100.0000, 100.0000, 0.6682],
+                [0.0000, 60.0000, 60.0000, 110.0000, 110.0000, 0.6457],
+            ]
+        )
+    ]
+
+    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]
+
+    mean_ap, first_ap = calculate_map(predictions, ground_truths)
+
+    expected_mean_ap = tensor(0.2)
+    expected_first_ap = tensor(0.5)
+
+    assert isclose(mean_ap, expected_mean_ap, atol=1e-5), f"Mean AP mismatch: {mean_ap} != {expected_mean_ap}"
+    assert isclose(first_ap, expected_first_ap, atol=1e-5), f"First AP mismatch: {first_ap} != {expected_first_ap}"

tests/{test_tools → test_utils}/test_module_utils.py

@@ -1,3 +1,4 @@
+import re
 import sys
 from pathlib import Path
 
@@ -6,7 +7,11 @@ from torch import nn
 
 project_root = Path(__file__).resolve().parent.parent.parent
 sys.path.append(str(project_root))
-from yolo.utils.module_utils import auto_pad, create_activation_function
+from yolo.utils.module_utils import (
+    auto_pad,
+    create_activation_function,
+    divide_into_chunks,
+)
 
 
 @pytest.mark.parametrize(
@@ -35,3 +40,34 @@ def test_get_activation(activation_name, expected_type):
 def test_get_activation_invalid():
     with pytest.raises(ValueError):
         create_activation_function("unsupported_activation")
+
+
+def test_divide_into_chunks():
+    input_list = [0, 1, 2, 3, 4, 5]
+    chunk_num = 2
+    expected_output = [[0, 1, 2], [3, 4, 5]]
+    assert divide_into_chunks(input_list, chunk_num) == expected_output
+
+
+def test_divide_into_chunks_non_divisible_length():
+    input_list = [0, 1, 2, 3, 4, 5]
+    chunk_num = 4
+    with pytest.raises(
+        ValueError,
+        match=re.escape("The length of the input list (6) must be exactly divisible by the number of chunks (4)."),
+    ):
+        divide_into_chunks(input_list, chunk_num)
+
+
+def test_divide_into_chunks_single_chunk():
+    input_list = [0, 1, 2, 3, 4, 5]
+    chunk_num = 1
+    expected_output = [[0, 1, 2, 3, 4, 5]]
+    assert divide_into_chunks(input_list, chunk_num) == expected_output
+
+
+def test_divide_into_chunks_equal_chunks():
+    input_list = [0, 1, 2, 3, 4, 5, 6, 7]
+    chunk_num = 4
+    expected_output = [[0, 1], [2, 3], [4, 5], [6, 7]]
+    assert divide_into_chunks(input_list, chunk_num) == expected_output

yolo/config/config.py

@@ -163,9 +163,6 @@ class YOLOLayer(nn.Module):
     layer_type: str
     usable: bool
 
-    def __post_init__(self):
-        super().__init__()
-
 
 IDX_TO_ID = [
     1,

yolo/tools/solver.py

@@ -122,6 +122,7 @@ class ModelTrainer:
             self.progress.finish_one_epoch(epoch_loss, epoch)
 
             self.validator.solve(self.validation_dataloader, epoch_idx=epoch)
+            # TODO: save model if result are better than before
         self.progress.finish_train()
 
 

yolo/utils/bounding_box_utils.py

@@ -46,7 +46,7 @@ def calculate_iou(bbox1, bbox2, metrics="iou") -> Tensor:
     # Calculate IoU
     iou = intersection_area / (union_area + EPS)
     if metrics == "iou":
-        return iou
+        return iou.to(dtype)
 
     # Calculate centroid distance
     cx1 = (bbox1[..., 2] + bbox1[..., 0]) / 2
@@ -62,7 +62,7 @@ def calculate_iou(bbox1, bbox2, metrics="iou") -> Tensor:
 
     diou = iou - (cent_dis / diag_dis)
     if metrics == "diou":
-        return diou
+        return diou.to(dtype)
 
     # Compute aspect ratio penalty term
     arctan = torch.atan((bbox1[..., 2] - bbox1[..., 0]) / (bbox1[..., 3] - bbox1[..., 1] + EPS)) - torch.atan(
@@ -268,7 +268,7 @@ class Vec2Box:
     def __init__(self, model: YOLO, image_size, device):
         self.device = device
 
-        if hasattr(model, "strides"):
+        if hasattr(model, "strides") and getattr(model, "strides"):
             logger.info(f"🈶 Found stride of model {model.strides}")
             self.strides = model.strides
         else:

yolo/utils/module_utils.py

@@ -68,8 +68,7 @@ def divide_into_chunks(input_list, chunk_num):
 
     if list_size % chunk_num != 0:
         raise ValueError(
-            f"The length of the input list ({list_size}) must be exactly\
-                            divisible by the number of chunks ({chunk_num})."
+            f"The length of the input list ({list_size}) must be exactly divisible by the number of chunks ({chunk_num})."
         )
 
     chunk_size = list_size // chunk_num