✅ [Add] test, increase test coverage for dev mode

tests/conftest.py

@@ -0,0 +1,91 @@
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+from hydra import compose, initialize
+
+project_root = Path(__file__).resolve().parent.parent
+sys.path.append(str(project_root))
+
+from yolo import Config, Vec2Box, create_model
+from yolo.model.yolo import YOLO
+from yolo.tools.data_loader import StreamDataLoader, YoloDataLoader
+from yolo.utils.logging_utils import ProgressLogger, set_seed
+
+
+def pytest_configure(config):
+    config.addinivalue_line("markers", "requires_cuda: mark test to run only if CUDA is available")
+
+
+def get_cfg(overrides=[]) -> Config:
+    config_path = "../yolo/config"
+    with initialize(config_path=config_path, version_base=None):
+        cfg: Config = compose(config_name="config", overrides=overrides)
+        set_seed(cfg.lucky_number)
+        return cfg
+
+
+@pytest.fixture(scope="session")
+def train_cfg() -> Config:
+    return get_cfg(overrides=["task=train", "dataset=mock"])
+
+
+@pytest.fixture(scope="session")
+def validation_cfg():
+    return get_cfg(overrides=["task=validation", "dataset=mock"])
+
+
+@pytest.fixture(scope="session")
+def inference_cfg():
+    return get_cfg(overrides=["task=inference"])
+
+
+@pytest.fixture(scope="session")
+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)
+    return model.to(device)
+
+
+@pytest.fixture(scope="session")
+def vec2box(train_cfg: Config, model: YOLO, device) -> Vec2Box:
+    vec2box = Vec2Box(model, train_cfg.image_size, device)
+    return vec2box
+
+
+@pytest.fixture(scope="session")
+def train_dataloader(train_cfg: Config):
+    return YoloDataLoader(train_cfg.task.data, train_cfg.dataset, train_cfg.task.task)
+
+
+@pytest.fixture(scope="session")
+def validation_dataloader(validation_cfg: Config):
+    return YoloDataLoader(validation_cfg.task.data, validation_cfg.dataset, validation_cfg.task.task)
+
+
+@pytest.fixture(scope="session")
+def file_stream_data_loader(inference_cfg: Config):
+    return StreamDataLoader(inference_cfg.task.data)
+
+
+@pytest.fixture(scope="session")
+def directory_stream_data_loader(inference_cfg: Config):
+    inference_cfg.task.data.source = "tests/data/images/train"
+    return StreamDataLoader(inference_cfg.task.data)

tests/test_tools/test_data_loader.py

@@ -0,0 +1,71 @@
+import sys
+from pathlib import Path
+
+import pytest
+
+project_root = Path(__file__).resolve().parent.parent.parent
+sys.path.append(str(project_root))
+
+from yolo.config.config import Config, TrainConfig
+from yolo.tools.data_loader import StreamDataLoader, YoloDataLoader, create_dataloader
+
+
+def test_create_dataloader_cache(train_cfg: Config):
+    train_cfg.task.data.shuffle = False
+    train_cfg.task.data.batch_size = 2
+
+    cache_file = Path("tests/data/train.cache")
+    cache_file.unlink(missing_ok=True)
+
+    make_cache_loader = create_dataloader(train_cfg.task.data, train_cfg.dataset)
+    load_cache_loader = create_dataloader(train_cfg.task.data, train_cfg.dataset)
+    m_batch_size, m_images, _, m_reverse_tensors, m_image_paths = next(iter(make_cache_loader))
+    l_batch_size, l_images, _, l_reverse_tensors, l_image_paths = next(iter(load_cache_loader))
+    assert m_batch_size == l_batch_size
+    assert m_images.shape == l_images.shape
+    assert m_reverse_tensors.shape == l_reverse_tensors.shape
+    assert m_image_paths == l_image_paths
+
+
+def test_training_data_loader_correctness(train_dataloader: YoloDataLoader):
+    """Test that the training data loader produces correctly shaped data and metadata."""
+    batch_size, images, _, reverse_tensors, image_paths = next(iter(train_dataloader))
+    assert batch_size == 2
+    assert images.shape == (2, 3, 640, 640)
+    assert reverse_tensors.shape == (2, 5)
+    expected_paths = [
+        Path("tests/data/images/train/000000050725.jpg"),
+        Path("tests/data/images/train/000000167848.jpg"),
+    ]
+    assert image_paths == expected_paths
+
+
+def test_validation_data_loader_correctness(validation_dataloader: YoloDataLoader):
+    batch_size, images, targets, reverse_tensors, image_paths = next(iter(validation_dataloader))
+    assert batch_size == 4
+    assert images.shape == (4, 3, 640, 640)
+    assert targets.shape == (4, 18, 5)
+    assert reverse_tensors.shape == (4, 5)
+    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
+
+
+def test_file_stream_data_loader_frame(file_stream_data_loader: StreamDataLoader):
+    """Test the frame output from the file stream data loader."""
+    frame, rev_tensor, origin_frame = next(iter(file_stream_data_loader))
+    assert frame.shape == (1, 3, 640, 640)
+    assert rev_tensor.shape == (1, 5)
+    assert origin_frame.size == (1024, 768)
+
+
+def test_directory_stream_data_loader_frame(directory_stream_data_loader: StreamDataLoader):
+    """Test the frame output from the directory stream data loader."""
+    frame, rev_tensor, origin_frame = next(iter(directory_stream_data_loader))
+    assert frame.shape == (1, 3, 640, 640)
+    assert rev_tensor.shape == (1, 5)
+    assert origin_frame.size == (480, 640)

tests/test_tools/test_dataset_preparation.py

@@ -0,0 +1,29 @@
+import os
+import shutil
+import sys
+from pathlib import Path
+
+project_root = Path(__file__).resolve().parent.parent.parent
+sys.path.append(str(project_root))
+
+from yolo.config.config import Config
+from yolo.tools.dataset_preparation import prepare_dataset, prepare_weight
+
+
+def test_prepare_dataset(train_cfg: Config):
+    dataset_path = Path("tests/data")
+    if dataset_path.exists():
+        shutil.rmtree(dataset_path)
+    prepare_dataset(train_cfg.dataset, task="train")
+    prepare_dataset(train_cfg.dataset, task="val")
+
+    images_path = Path("tests/data/images")
+    for data_type in images_path.iterdir():
+        assert len(os.listdir(data_type)) == 5
+
+    annotations_path = Path("tests/data/annotations")
+    assert os.listdir(annotations_path) == ["instances_val.json", "instances_train.json"]
+
+
+def test_prepare_weight():
+    prepare_weight()

tests/test_tools/test_drawer.py

@@ -0,0 +1,29 @@
+import sys
+from pathlib import Path
+
+from PIL import Image
+from torch import tensor
+
+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
+from yolo.tools.drawer import draw_bboxes, draw_model
+
+
+def test_draw_model_by_config(train_cfg: Config):
+    """Test the drawing of a model based on a configuration."""
+    draw_model(model_cfg=train_cfg.model)
+
+
+def test_draw_model_by_model(model: YOLO):
+    """Test the drawing of a YOLO model."""
+    draw_model(model=model)
+
+
+def test_draw_bboxes():
+    """Test drawing bounding boxes on an image."""
+    predictions = tensor([[0, 60, 60, 160, 160, 0.5], [0, 40, 40, 120, 120, 0.5]])
+    pil_image = Image.open("tests/data/images/train/000000050725.jpg")
+    draw_bboxes(pil_image, [predictions])

tests/test_tools/test_solver.py

@@ -1,114 +1,70 @@
 import sys
 from pathlib import Path
-from unittest.mock import MagicMock, patch
 
 import pytest
-import torch
-from hydra import compose, initialize
+from torch import allclose, tensor
 
 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.config.config import Config
+from yolo.model.yolo import YOLO
+from yolo.tools.data_loader import StreamDataLoader, YoloDataLoader
+from yolo.tools.solver import 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
+def model_validator(validation_cfg: Config, model: YOLO, vec2box: Vec2Box, validation_progress_logger, device):
+    validator = ModelValidator(
+        validation_cfg.task, validation_cfg.dataset, model, vec2box, validation_progress_logger, device
+    )
+    return validator
 
 
-@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
+def test_model_validator_initialization(model_validator: ModelValidator):
+    assert isinstance(model_validator.model, YOLO)
+    assert hasattr(model_validator, "solve")
 
 
-@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
+def test_model_validator_solve_mock_dataset(model_validator: ModelValidator, validation_dataloader: YoloDataLoader):
+    mAPs = model_validator.solve(validation_dataloader)
+    except_mAPs = {"mAP.5": tensor(0.6969), "mAP.5:.95": tensor(0.4195)}
+    assert allclose(mAPs["mAP.5"], except_mAPs["mAP.5"], rtol=1e-4)
+    print(mAPs)
+    assert allclose(mAPs["mAP.5:.95"], except_mAPs["mAP.5:.95"], rtol=1e-4)
 
 
 @pytest.fixture
-def device() -> torch.device:
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    return device
+def model_tester(inference_cfg: Config, model: YOLO, vec2box: Vec2Box, validation_progress_logger, device):
+    tester = ModelTester(inference_cfg, model, vec2box, validation_progress_logger, device)
+    return tester
 
 
-@pytest.fixture
-def model(cfg: Config, device) -> YOLO:
-    model = create_model(cfg.model, weight_path=None)
-    return model.to(device)
+def test_model_tester_initialization(model_tester: ModelTester):
+    assert isinstance(model_tester.model, YOLO)
+    assert hasattr(model_tester, "solve")
 
 
-@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
+def test_model_tester_solve_single_image(model_tester: ModelTester, file_stream_data_loader: StreamDataLoader):
+    model_tester.solve(file_stream_data_loader)
 
 
 @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 model_trainer(train_cfg: Config, model: YOLO, vec2box: Vec2Box, train_progress_logger, device):
+    train_cfg.task.epoch = 2
+    trainer = ModelTrainer(train_cfg, model, vec2box, train_progress_logger, device, use_ddp=False)
+    return trainer
 
-# 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_trainer_initialization(model_trainer: ModelTrainer):
 
-def test_model_validator_initialization(cfg_validaion: Config, model: YOLO, vec2box: Vec2Box, progress_logger, device):
-    validator = ModelValidator(cfg_validaion.task, cfg_validaion.dataset, model, vec2box, progress_logger, device)
-    assert validator.model == model
-    assert validator.device == device
-    assert validator.progress == progress_logger
+    assert isinstance(model_trainer.model, YOLO)
+    assert hasattr(model_trainer, "solve")
+    assert model_trainer.optimizer is not None
+    assert model_trainer.scheduler is not None
+    assert model_trainer.loss_fn is not None
 
 
-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
+def test_model_trainer_solve_mock_dataset(model_trainer: ModelTrainer, train_dataloader: YoloDataLoader):
+    model_trainer.solve(train_dataloader)

yolo/tools/data_loader.py

@@ -111,7 +111,7 @@ class YoloDataset(Dataset):
         logger.info("Recorded {}/{} valid inputs", valid_inputs, len(images_list))
         return data
 
-    def load_valid_labels(self, label_path: str, seg_data_one_img: list) -> Union[torch.Tensor, None]:
+    def load_valid_labels(self, label_path: str, seg_data_one_img: list) -> Union[Tensor, None]:
         """
         Loads and validates bounding box data is [0, 1] from a label file.
 
@@ -119,7 +119,7 @@ class YoloDataset(Dataset):
             label_path (str): The filepath to the label file containing bounding box data.
 
         Returns:
-            torch.Tensor or None: A tensor of all valid bounding boxes if any are found; otherwise, None.
+            Tensor or None: A tensor of all valid bounding boxes if any are found; otherwise, None.
         """
         bboxes = []
         for seg_data in seg_data_one_img:
@@ -145,7 +145,7 @@ class YoloDataset(Dataset):
         indices = torch.randint(0, len(self), (num,))
         return [self.get_data(idx)[:2] for idx in indices]
 
-    def __getitem__(self, idx) -> Union[Image.Image, torch.Tensor]:
+    def __getitem__(self, idx) -> Tuple[Image.Image, Tensor, Tensor, List[str]]:
         img, bboxes, img_path = self.get_data(idx)
         img, bboxes, rev_tensor = self.transform(img, bboxes)
         return img, bboxes, rev_tensor, img_path
@@ -170,17 +170,17 @@ class YoloDataLoader(DataLoader):
             collate_fn=self.collate_fn,
         )
 
-    def collate_fn(self, batch: List[Tuple[torch.Tensor, torch.Tensor]]) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+    def collate_fn(self, batch: List[Tuple[Tensor, Tensor]]) -> Tuple[Tensor, List[Tensor]]:
         """
         A collate function to handle batching of images and their corresponding targets.
 
         Args:
             batch (list of tuples): Each tuple contains:
-                - image (torch.Tensor): The image tensor.
-                - labels (torch.Tensor): The tensor of labels for the image.
+                - image (Tensor): The image tensor.
+                - labels (Tensor): The tensor of labels for the image.
 
         Returns:
-            Tuple[torch.Tensor, List[torch.Tensor]]: A tuple containing:
+            Tuple[Tensor, List[Tensor]]: A tuple containing:
                 - A tensor of batched images.
                 - A list of tensors, each corresponding to bboxes for each image in the batch.
         """
@@ -213,7 +213,7 @@ def create_dataloader(data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: st
 
 class StreamDataLoader:
     def __init__(self, data_cfg: DataConfig):
-        self.source = Path(data_cfg.source)
+        self.source = data_cfg.source
         self.running = True
         self.is_stream = isinstance(self.source, int) or str(self.source).lower().startswith("rtmp://")
 
@@ -225,6 +225,7 @@ class StreamDataLoader:
 
             self.cap = cv2.VideoCapture(self.source)
         else:
+            self.source = Path(self.source)
             self.queue = Queue()
             self.thread = Thread(target=self.load_source)
             self.thread.start()

yolo/tools/dataset_preparation.py

@@ -82,7 +82,7 @@ def prepare_dataset(dataset_cfg: DatasetConfig, task: str):
                 logger.error(f"Error verifying the {dataset_type} dataset after extraction.")
 
 
-def prepare_weight(download_link: Optional[str] = None, weight_path: Path = "v9-c.pt"):
+def prepare_weight(download_link: Optional[str] = None, weight_path: Path = Path("v9-c.pt")):
     weight_name = weight_path.name
     if download_link is None:
         download_link = "https://github.com/WongKinYiu/yolov9mit/releases/download/v1.0-alpha/"
@@ -97,13 +97,3 @@ def prepare_weight(download_link: Optional[str] = None, weight_path: Path = "v9-
         download_file(weight_link, weight_path)
     except requests.exceptions.RequestException as e:
         logger.warning(f"Failed to download the weight file: {e}")
-
-
-if __name__ == "__main__":
-    import sys
-
-    sys.path.append("./")
-    from utils.logging_utils import custom_logger
-
-    custom_logger()
-    prepare_weight()

yolo/tools/drawer.py

@@ -7,6 +7,9 @@ from loguru import logger
 from PIL import Image, ImageDraw, ImageFont
 from torchvision.transforms.functional import to_pil_image
 
+from yolo.config.config import ModelConfig
+from yolo.model.yolo import YOLO
+
 
 def draw_bboxes(
     img: Union[Image.Image, torch.Tensor],
@@ -62,7 +65,7 @@ def draw_bboxes(
     return img
 
 
-def draw_model(*, model_cfg=None, model=None, v7_base=False):
+def draw_model(*, model_cfg: ModelConfig = None, model: YOLO = None, v7_base=False):
     from graphviz import Digraph
 
     if model_cfg:

yolo/utils/logging_utils.py

@@ -138,7 +138,8 @@ class ProgressLogger(Progress):
     def finish_train(self):
         self.remove_task(self.task_epoch)
         self.stop()
-        self.wandb.finish()
+        if self.use_wandb:
+            self.wandb.finish()
 
 
 def custom_wandb_log(string="", level=int, newline=True, repeat=True, prefix=True, silent=False):