✅ [Pass] the test and rename the function of v2 arch

tests/test_model/test_yolo.py

@@ -1,6 +1,7 @@
 import sys
 from pathlib import Path
 
+import pytest
 import torch
 from hydra import compose, initialize
 from omegaconf import OmegaConf
@@ -8,6 +9,7 @@ from omegaconf import OmegaConf
 project_root = Path(__file__).resolve().parent.parent.parent
 sys.path.append(str(project_root))
 
+from yolo.config.config import Config
 from yolo.model.yolo import YOLO, create_model
 
 config_path = "../../yolo/config"
@@ -16,37 +18,44 @@ config_name = "config"
 
 def test_build_model():
     with initialize(config_path=config_path, version_base=None):
-        cfg = compose(config_name=config_name)
+        cfg: Config = compose(config_name=config_name)
 
         OmegaConf.set_struct(cfg.model, False)
         cfg.weight = None
         model = YOLO(cfg.model)
-        assert len(model.model) == 38
+        assert len(model.model) == 39
 
 
-def test_create_model():
-    with initialize(config_path=config_path, version_base=None):
-        cfg = compose(config_name=config_name)
+@pytest.fixture
+def cfg() -> Config:
+    with initialize(config_path="../../yolo/config", version_base=None):
+        cfg: Config = compose(config_name="config")
         cfg.weight = None
-        model = create_model(cfg)
-        assert isinstance(model, YOLO)
+    return cfg
 
 
-def test_yolo_forward_output_shape():
-    with initialize(config_path=config_path, version_base=None):
-        cfg = compose(config_name=config_name)
-        model = create_model(cfg)
-        # 2 - batch size, 3 - number of channels, 640x640 - image dimensions
-        dummy_input = torch.rand(2, 3, 640, 640)
-
-        # Forward pass through the model
-        output = model(dummy_input)
-        output_shape = [x.shape for x in output[-1]]
-        assert output_shape == [
-            torch.Size([2, 144, 80, 80]),
-            torch.Size([2, 144, 40, 40]),
-            torch.Size([2, 144, 20, 20]),
-            torch.Size([2, 144, 80, 80]),
-            torch.Size([2, 144, 40, 40]),
-            torch.Size([2, 144, 20, 20]),
-        ]
+@pytest.fixture
+def model(cfg: Config):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = create_model(cfg.model)
+    return model.to(device)
+
+
+def test_model_basic_status(model):
+    assert isinstance(model, YOLO)
+    assert len(model.model) == 39
+
+
+def test_yolo_forward_output_shape(model):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # 2 - batch size, 3 - number of channels, 640x640 - image dimensions
+    dummy_input = torch.rand(2, 3, 640, 640, device=device)
+
+    # Forward pass through the model
+    output = model(dummy_input)
+    output_shape = [(cls.shape, anc.shape, box.shape) for cls, anc, box in output["Main"]]
+    assert output_shape == [
+        (torch.Size([2, 80, 80, 80]), torch.Size([2, 16, 4, 80, 80]), torch.Size([2, 4, 80, 80])),
+        (torch.Size([2, 80, 40, 40]), torch.Size([2, 16, 4, 40, 40]), torch.Size([2, 4, 40, 40])),
+        (torch.Size([2, 80, 20, 20]), torch.Size([2, 16, 4, 20, 20]), torch.Size([2, 4, 20, 20])),
+    ]

tests/test_utils/test_dataaugment.py

@@ -43,8 +43,9 @@ def test_compose():
     boxes = torch.tensor([[0, 0.2, 0.2, 0.8, 0.8]])
 
     transformed_img, transformed_boxes = compose(img, boxes)
+    tensor_img = TF.pil_to_tensor(img).to(torch.float32) / 255
 
-    assert transformed_img == img, "Image should not be altered"
+    assert (transformed_img == tensor_img).all(), "Image should not be altered"
     assert torch.equal(transformed_boxes, boxes), "Boxes should not be altered"
 
 

tests/test_utils/test_loss.py

@@ -8,32 +8,50 @@ from hydra import compose, initialize
 project_root = Path(__file__).resolve().parent.parent.parent
 sys.path.append(str(project_root))
 
-from yolo.tools.loss_functions import YOLOLoss
+from yolo.config.config import Config
+from yolo.model.yolo import create_model
+from yolo.tools.loss_functions import DualLoss, create_loss_function
+from yolo.utils.bounding_box_utils import Vec2Box
 
 
 @pytest.fixture
-def cfg():
+def cfg() -> Config:
     with initialize(config_path="../../yolo/config", version_base=None):
         cfg = compose(config_name="config")
     return cfg
 
 
 @pytest.fixture
-def loss_function(cfg) -> YOLOLoss:
-    return YOLOLoss(cfg)
+def model(cfg: Config):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = create_model(cfg.model)
+    return model.to(device)
+
+
+@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)
+
+
+@pytest.fixture
+def loss_function(cfg, vec2box) -> DualLoss:
+    return create_loss_function(cfg, vec2box)
 
 
 @pytest.fixture
 def data():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     targets = torch.zeros(1, 20, 5, device=device)
-    predicts = [torch.zeros(1, 144, 80 // i, 80 // i, device=device) for i in [1, 2, 4]]
+    predicts = [torch.zeros(1, 8400, *cn, device=device) for cn in [(80,), (4, 16), (4,)]]
     return predicts, targets
 
 
 def test_yolo_loss(loss_function, data):
     predicts, targets = data
-    loss_iou, loss_dfl, loss_cls = loss_function(predicts, targets)
-    assert torch.isnan(loss_iou)
-    assert torch.isnan(loss_dfl)
-    assert torch.isinf(loss_cls)
+    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"])

yolo/model/README.md

@@ -1 +0,0 @@
-main code

yolo/tools/loss_functions.py

@@ -132,7 +132,7 @@ class DualLoss:
         return loss_sum, loss_dict
 
 
-def get_loss_function(cfg: Config, vec2box) -> DualLoss:
+def create_loss_function(cfg: Config, vec2box) -> DualLoss:
     loss_function = DualLoss(cfg, vec2box)
     logger.info("✅ Success load loss function")
     return loss_function

yolo/tools/solver.py

@@ -9,7 +9,7 @@ from yolo.config.config import Config, TrainConfig, ValidationConfig
 from yolo.model.yolo import YOLO
 from yolo.tools.data_loader import StreamDataLoader, create_dataloader
 from yolo.tools.drawer import draw_bboxes
-from yolo.tools.loss_functions import get_loss_function
+from yolo.tools.loss_functions import create_loss_function
 from yolo.utils.bounding_box_utils import Vec2Box, bbox_nms, calculate_map
 from yolo.utils.logging_utils import ProgressTracker
 from yolo.utils.model_utils import (
@@ -27,7 +27,7 @@ class ModelTrainer:
         self.device = device
         self.optimizer = create_optimizer(model, train_cfg.optimizer)
         self.scheduler = create_scheduler(self.optimizer, train_cfg.scheduler)
-        self.loss_fn = get_loss_function(cfg, vec2box)
+        self.loss_fn = create_loss_function(cfg, vec2box)
         self.progress = ProgressTracker(cfg.name, save_path, cfg.use_wandb)
         self.num_epochs = cfg.task.epoch