✅ [Pass] Test for using anchor config in Vec2Box

tests/conftest.py

@@ -67,7 +67,7 @@ def model(train_cfg: Config, device) -> YOLO:
 
 @pytest.fixture(scope="session")
 def vec2box(train_cfg: Config, model: YOLO, device) -> Vec2Box:
-    vec2box = Vec2Box(model, train_cfg.image_size, device)
+    vec2box = Vec2Box(model, train_cfg.model.anchor, train_cfg.image_size, device)
     return vec2box
 
 

tests/test_tools/test_data_loader.py

@@ -31,10 +31,10 @@ def test_training_data_loader_correctness(train_dataloader: YoloDataLoader):
     assert batch_size == 2
     assert images.shape == (2, 3, 640, 640)
     assert reverse_tensors.shape == (2, 5)
-    expected_paths = (
+    expected_paths = [
         Path("tests/data/images/train/000000050725.jpg"),
         Path("tests/data/images/train/000000167848.jpg"),
-    )
+    ]
     assert image_paths == expected_paths
 
 
@@ -44,12 +44,12 @@ def test_validation_data_loader_correctness(validation_dataloader: YoloDataLoade
     assert images.shape == (4, 3, 640, 640)
     assert targets.shape == (4, 18, 5)
     assert reverse_tensors.shape == (4, 5)
-    expected_paths = (
+    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 image_paths == expected_paths
 
 

tests/test_tools/test_loss_functions.py

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

tests/test_utils/test_bounding_box_utils.py

@@ -4,7 +4,7 @@ from pathlib import Path
 import pytest
 import torch
 from hydra import compose, initialize
-from torch import Tensor, allclose, float32, isclose, nn, tensor
+from torch import allclose, float32, isclose, tensor
 
 project_root = Path(__file__).resolve().parent.parent.parent
 sys.path.append(str(project_root))
@@ -117,7 +117,7 @@ def test_vec2box_autoanchor():
         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)
+    vec2box = Vec2Box(model, cfg.model.anchor, cfg.image_size, device)
     assert vec2box.strides == [8, 16, 32]
 
     vec2box.update((320, 640))

yolo/lazy.py

@@ -27,7 +27,7 @@ def main(cfg: Config):
         model = create_model(cfg.model, class_num=cfg.class_num, weight_path=cfg.weight)
         model = model.to(device)
 
-    vec2box = Vec2Box(model, cfg.image_size, device)
+    vec2box = Vec2Box(model, cfg.model.anchor, cfg.image_size, device)
     if cfg.task.task == "train":
         solver = ModelTrainer(cfg, model, vec2box, progress, device, use_ddp)
     if cfg.task.task == "validation":

yolo/utils/bounding_box_utils.py

@@ -364,7 +364,7 @@ class Anc2Box:
         return preds_cls, None, preds_box, preds_cnf.sigmoid()
 
 
-def bbox_nms(cls_dist: Tensor, bbox: Tensor, nms_cfg: NMSConfig, confidence: Optional[Tensor]):
+def bbox_nms(cls_dist: Tensor, bbox: Tensor, nms_cfg: NMSConfig, confidence: Optional[Tensor] = None):
     cls_dist = cls_dist.sigmoid() * (1 if confidence is None else confidence)
 
     # filter class by confidence