✅ [Pass] and add tests for yolo.py

tests/model/test_yolo.py

@@ -1,41 +0,0 @@
-import pytest
-import torch
-import sys
-
-sys.path.append("./")
-from utils.tools import load_model_cfg
-from model.yolo import YOLO
-
-
-def test_load_model_configuration():
-    config_name = "v7-base"
-    model_cfg = load_model_cfg(config_name)
-
-    assert "model" in model_cfg
-    assert isinstance(model_cfg, dict)
-
-
-def test_build_model():
-    config_name = "v7-base"
-    model_cfg = load_model_cfg(config_name)
-    model = YOLO(model_cfg)
-    model.build_model(model_cfg["model"])
-    assert len(model.model) == 106
-
-
-def test_yolo_forward_output_shape():
-    config_name = "v7-base"
-    model_cfg = load_model_cfg(config_name)
-    model = YOLO(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, 3, 20, 20, 85]),
-        torch.Size([2, 3, 80, 80, 85]),
-        torch.Size([2, 3, 40, 40, 85]),
-    ]

tests/test_model/test_yolo.py

@@ -0,0 +1,48 @@
+import pytest
+import torch
+from hydra import initialize, compose
+from hydra.core.global_hydra import GlobalHydra
+from omegaconf import DictConfig, OmegaConf
+
+import sys
+
+sys.path.append("./")
+from model.yolo import YOLO, get_model
+
+config_path = "../../config/model"
+config_name = "v7-base"
+
+
+def test_build_model():
+
+    with initialize(config_path=config_path, version_base=None):
+        model_cfg = compose(config_name=config_name)
+        OmegaConf.set_struct(model_cfg, False)
+        model = YOLO(model_cfg)
+        model.build_model(model_cfg.model)
+        assert len(model.model) == 106
+
+
+def test_get_model():
+    with initialize(config_path=config_path, version_base=None):
+        model_cfg = compose(config_name=config_name)
+        model = get_model(model_cfg)
+        assert isinstance(model, YOLO)
+
+
+def test_yolo_forward_output_shape():
+    with initialize(config_path=config_path, version_base=None):
+        model_cfg = compose(config_name=config_name)
+
+        model = get_model(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, 3, 20, 20, 85]),
+            torch.Size([2, 3, 80, 80, 85]),
+            torch.Size([2, 3, 40, 40, 85]),
+        ]

tests/test_utils/test_dataaugment.py

@@ -0,0 +1,63 @@
+import pytest
+import torch
+from PIL import Image
+from torchvision.transforms import functional as TF
+import sys
+
+sys.path.append("./")
+from utils.dataargument import RandomHorizontalFlip, Compose, Mosaic
+
+
+def test_random_horizontal_flip():
+    # Create a mock image and bounding boxes
+    img = Image.new("RGB", (100, 100), color="red")
+    boxes = torch.tensor([[1, 0.1, 0.1, 0.9, 0.9]])  # class, xmin, ymin, xmax, ymax
+
+    flip_transform = RandomHorizontalFlip(prob=1)  # Set probability to 1 to ensure flip
+    flipped_img, flipped_boxes = flip_transform(img, boxes)
+
+    # Assert image is flipped by comparing it to a manually flipped image
+    assert TF.hflip(img) == flipped_img
+
+    # Assert bounding boxes are flipped correctly
+    expected_boxes = torch.tensor([[1, 0.1, 0.1, 0.9, 0.9]])
+    expected_boxes[:, [1, 3]] = 1 - expected_boxes[:, [3, 1]]
+    assert torch.allclose(flipped_boxes, expected_boxes), "Bounding boxes were not flipped correctly"
+
+
+def test_compose():
+    # Define two mock transforms that simply return the inputs
+    def mock_transform(image, boxes):
+        return image, boxes
+
+    compose = Compose([mock_transform, mock_transform])
+    img = Image.new("RGB", (10, 10), color="blue")
+    boxes = torch.tensor([[0, 0.2, 0.2, 0.8, 0.8]])
+
+    transformed_img, transformed_boxes = compose(img, boxes)
+
+    assert transformed_img == img, "Image should not be altered"
+    assert torch.equal(transformed_boxes, boxes), "Boxes should not be altered"
+
+
+def test_mosaic():
+    img = Image.new("RGB", (100, 100), color="green")
+    boxes = torch.tensor([[0, 0.25, 0.25, 0.75, 0.75]])
+
+    # Mock parent with image_size and get_more_data method
+    class MockParent:
+        image_size = 100
+
+        def get_more_data(self, num_images):
+            return [(img, boxes) for _ in range(num_images)]
+
+    mosaic = Mosaic(prob=1)  # Ensure mosaic is applied
+    mosaic.set_parent(MockParent())
+
+    mosaic_img, mosaic_boxes = mosaic(img, boxes)
+
+    # Checks here would depend on the exact expected behavior of the mosaic function,
+    # such as dimensions and content of the output image and boxes.
+
+    assert mosaic_img.size == (200, 200), "Mosaic image size should be doubled"
+    assert len(mosaic_boxes) > 0, "Should have some bounding boxes"