🔀 [Merge] branch 'TEST'

.github/workflows/deploy.yaml

@@ -53,17 +53,17 @@ jobs:
 
     - name: Run Validation
       run: |
-        python yolo/lazy.py task=validation dataset=mock
-        python yolo/lazy.py task=validation dataset=mock model=v9-s
-        python yolo/lazy.py task=validation dataset=mock name=AnyNameYouWant
+        python yolo/lazy.py task=validation use_wandb=False dataset=mock
+        python yolo/lazy.py task=validation use_wandb=False dataset=mock model=v9-s
+        python yolo/lazy.py task=validation use_wandb=False dataset=mock name=AnyNameYouWant
 
     - name: Run Inference
       run: |
-        python yolo/lazy.py task=inference
-        python yolo/lazy.py task=inference model=v7
-        python yolo/lazy.py task=inference +quite=True
-        python yolo/lazy.py task=inference name=AnyNameYouWant
-        python yolo/lazy.py task=inference image_size=\[480,640]
-        python yolo/lazy.py task=inference task.nms.min_confidence=0.1
-        python yolo/lazy.py task=inference task.fast_inference=deploy
-        python yolo/lazy.py task=inference task.data.source=tests/data/images/val
+        python yolo/lazy.py task=inference use_wandb=False
+        python yolo/lazy.py task=inference use_wandb=False model=v7
+        python yolo/lazy.py task=inference use_wandb=False +quite=True
+        python yolo/lazy.py task=inference use_wandb=False name=AnyNameYouWant
+        python yolo/lazy.py task=inference use_wandb=False image_size=\[480,640]
+        python yolo/lazy.py task=inference use_wandb=False task.nms.min_confidence=0.1
+        python yolo/lazy.py task=inference use_wandb=False task.fast_inference=deploy
+        python yolo/lazy.py task=inference use_wandb=False task.data.source=tests/data/images/val

requirements.txt

@@ -1,6 +1,7 @@
 einops
 graphviz
 hydra-core
+lightning
 loguru
 numpy
 opencv-python

tests/conftest.py

@@ -4,15 +4,16 @@
 import pytest
 import torch
 from hydra import compose, initialize
+from lightning import Trainer
 
 project_root = Path(__file__).resolve().parent.parent
 sys.path.append(str(project_root))
 
 from yolo import Anc2Box, Config, Vec2Box, create_converter, create_model
 from yolo.model.yolo import YOLO
-from yolo.tools.data_loader import StreamDataLoader, YoloDataLoader
+from yolo.tools.data_loader import StreamDataLoader, create_dataloader
 from yolo.tools.dataset_preparation import prepare_dataset
-from yolo.utils.logging_utils import ProgressLogger, set_seed
+from yolo.utils.logging_utils import set_seed, setup
 
 
 def pytest_configure(config):
@@ -52,18 +53,6 @@ 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)
@@ -76,6 +65,24 @@ def model_v7(inference_v7_cfg: Config, device) -> YOLO:
     return model.to(device)
 
 
+@pytest.fixture(scope="session")
+def solver(train_cfg: Config) -> Trainer:
+    train_cfg.use_wandb = False
+    callbacks, loggers, save_path = setup(train_cfg)
+    trainer = Trainer(
+        accelerator="auto",
+        max_epochs=getattr(train_cfg.task, "epoch", None),
+        precision="16-mixed",
+        callbacks=callbacks,
+        logger=loggers,
+        log_every_n_steps=1,
+        gradient_clip_val=10,
+        deterministic=True,
+        default_root_dir=save_path,
+    )
+    return trainer
+
+
 @pytest.fixture(scope="session")
 def vec2box(train_cfg: Config, model: YOLO, device) -> Vec2Box:
     vec2box = create_converter(train_cfg.model.name, model, train_cfg.model.anchor, train_cfg.image_size, device)
@@ -93,13 +100,13 @@ def anc2box(inference_v7_cfg: Config, model: YOLO, device) -> Anc2Box:
 @pytest.fixture(scope="session")
 def train_dataloader(train_cfg: Config):
     prepare_dataset(train_cfg.dataset, task="train")
-    return YoloDataLoader(train_cfg.task.data, train_cfg.dataset, train_cfg.task.task)
+    return create_dataloader(train_cfg.task.data, train_cfg.dataset, train_cfg.task.task)
 
 
 @pytest.fixture(scope="session")
 def validation_dataloader(validation_cfg: Config):
     prepare_dataset(validation_cfg.dataset, task="val")
-    return YoloDataLoader(validation_cfg.task.data, validation_cfg.dataset, validation_cfg.task.task)
+    return create_dataloader(validation_cfg.task.data, validation_cfg.dataset, validation_cfg.task.task)
 
 
 @pytest.fixture(scope="session")

tests/test_tools/test_data_loader.py

@@ -1,11 +1,13 @@
 import sys
 from pathlib import Path
 
+from torch.utils.data import DataLoader
+
 project_root = Path(__file__).resolve().parent.parent.parent
 sys.path.append(str(project_root))
 
 from yolo.config.config import Config
-from yolo.tools.data_loader import StreamDataLoader, YoloDataLoader, create_dataloader
+from yolo.tools.data_loader import StreamDataLoader, create_dataloader
 
 
 def test_create_dataloader_cache(train_cfg: Config):
@@ -25,7 +27,7 @@ def test_create_dataloader_cache(train_cfg: Config):
     assert m_image_paths == l_image_paths
 
 
-def test_training_data_loader_correctness(train_dataloader: YoloDataLoader):
+def test_training_data_loader_correctness(train_dataloader: DataLoader):
     """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
@@ -38,7 +40,7 @@ def test_training_data_loader_correctness(train_dataloader: YoloDataLoader):
     assert list(image_paths) == list(expected_paths)
 
 
-def test_validation_data_loader_correctness(validation_dataloader: YoloDataLoader):
+def test_validation_data_loader_correctness(validation_dataloader: DataLoader):
     batch_size, images, targets, reverse_tensors, image_paths = next(iter(validation_dataloader))
     assert batch_size == 4
     assert images.shape == (4, 3, 640, 640)

tests/test_tools/test_loss_functions.py

@@ -1,4 +1,5 @@
 import sys
+from math import isinf, isnan
 from pathlib import Path
 
 import pytest
@@ -51,6 +52,6 @@ def test_yolo_loss(loss_function, data):
     predicts, targets = data
     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"])
+    assert isnan(loss_dict["Loss/BoxLoss"])
+    assert isnan(loss_dict["Loss/DFLLoss"])
+    assert isinf(loss_dict["Loss/BCELoss"])

tests/test_tools/test_solver.py

@@ -1,79 +1,81 @@
 import sys
+from math import isclose
 from pathlib import Path
 
 import pytest
-from torch import allclose, tensor
+from lightning.pytorch import Trainer
+from torch.utils.data import DataLoader
 
 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.data_loader import StreamDataLoader, YoloDataLoader
-from yolo.tools.solver import ModelTester, ModelTrainer, ModelValidator
+from yolo.tools.data_loader import StreamDataLoader
+from yolo.tools.solver import InferenceModel, TrainModel, ValidateModel
 from yolo.utils.bounding_box_utils import Anc2Box, Vec2Box
 
 
 @pytest.fixture
-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
-    )
+def model_validator(validation_cfg: Config):
+    validator = ValidateModel(validation_cfg)
     return validator
 
 
-def test_model_validator_initialization(model_validator: ModelValidator):
+def test_model_validator_initialization(solver: Trainer, model_validator: ValidateModel):
     assert isinstance(model_validator.model, YOLO)
-    assert hasattr(model_validator, "solve")
+    assert hasattr(solver, "validate")
 
 
-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=0.1)
-    print(mAPs)
-    assert allclose(mAPs["mAP.5:.95"], except_mAPs["mAP.5:.95"], rtol=0.1)
+def test_model_validator_solve_mock_dataset(
+    solver: Trainer, model_validator: ValidateModel, validation_dataloader: DataLoader
+):
+    mAPs = solver.validate(model_validator, dataloaders=validation_dataloader)[0]
+    except_mAPs = {"map_50": 0.7379, "map": 0.5617}
+    assert isclose(mAPs["map_50"], except_mAPs["map_50"], abs_tol=1e-4)
+    assert isclose(mAPs["map"], except_mAPs["map"], abs_tol=0.1)
 
 
 @pytest.fixture
-def model_tester(inference_cfg: Config, model: YOLO, vec2box: Vec2Box, validation_progress_logger, device):
-    tester = ModelTester(inference_cfg, model, vec2box, validation_progress_logger, device)
+def model_tester(inference_cfg: Config):
+    tester = InferenceModel(inference_cfg)
     return tester
 
 
 @pytest.fixture
-def modelv7_tester(inference_v7_cfg: Config, model_v7: YOLO, anc2box: Anc2Box, validation_progress_logger, device):
-    tester = ModelTester(inference_v7_cfg, model_v7, anc2box, validation_progress_logger, device)
+def modelv7_tester(inference_v7_cfg: Config):
+    tester = InferenceModel(inference_v7_cfg)
     return tester
 
 
-def test_model_tester_initialization(model_tester: ModelTester):
+def test_model_tester_initialization(solver: Trainer, model_tester: InferenceModel):
     assert isinstance(model_tester.model, YOLO)
-    assert hasattr(model_tester, "solve")
+    assert hasattr(solver, "predict")
 
 
-def test_model_tester_solve_single_image(model_tester: ModelTester, file_stream_data_loader: StreamDataLoader):
-    model_tester.solve(file_stream_data_loader)
+def test_model_tester_solve_single_image(
+    solver: Trainer, model_tester: InferenceModel, file_stream_data_loader: StreamDataLoader
+):
+    solver.predict(model_tester, file_stream_data_loader)
 
 
-def test_modelv7_tester_solve_single_image(modelv7_tester: ModelTester, file_stream_data_loader_v7: StreamDataLoader):
-    modelv7_tester.solve(file_stream_data_loader_v7)
+def test_modelv7_tester_solve_single_image(
+    solver: Trainer, modelv7_tester: InferenceModel, file_stream_data_loader_v7: StreamDataLoader
+):
+    solver.predict(modelv7_tester, file_stream_data_loader_v7)
 
 
 @pytest.fixture
-def model_trainer(train_cfg: Config, model: YOLO, vec2box: Vec2Box, train_progress_logger, device):
+def model_trainer(train_cfg: Config):
     train_cfg.task.epoch = 2
-    trainer = ModelTrainer(train_cfg, model, vec2box, train_progress_logger, device, use_ddp=False)
+    trainer = TrainModel(train_cfg)
     return trainer
 
 
-def test_model_trainer_initialization(model_trainer: ModelTrainer):
-
+def test_model_trainer_initialization(solver: Trainer, model_trainer: TrainModel):
     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
+    assert hasattr(solver, "fit")
+    assert solver.optimizers is not None
 
 
 # def test_model_trainer_solve_mock_dataset(model_trainer: ModelTrainer, train_dataloader: YoloDataLoader):

tests/test_utils/test_bounding_box_utils.py

@@ -216,9 +216,8 @@ def test_calculate_map():
     ground_truths = tensor([[0, 50, 50, 150, 150], [0, 30, 30, 100, 100]])  # [class, x1, y1, x2, y2]
 
     mAP = calculate_map(predictions, ground_truths)
+    expected_ap50 = tensor(0.5050)
+    expected_ap50_95 = tensor(0.2020)
 
-    expected_ap50 = tensor(0.5)
-    expected_ap50_95 = tensor(0.2)
-
-    assert isclose(mAP["mAP.5"], expected_ap50, atol=1e-5), f"AP50 mismatch"
-    assert isclose(mAP["mAP.5:.95"], expected_ap50_95, atol=1e-5), f"Mean AP mismatch"
+    assert isclose(mAP["map_50"], expected_ap50, atol=1e-4), f"AP50 mismatch"
+    assert isclose(mAP["map"], expected_ap50_95, atol=1e-4), f"Mean AP mismatch"

yolo/__init__.py

@@ -10,7 +10,7 @@
     YOLORichModelSummary,
     YOLORichProgressBar,
 )
-from yolo.utils.model_utils import PostProccess
+from yolo.utils.model_utils import PostProcess
 
 all = [
     "create_model",
@@ -29,5 +29,5 @@
     "create_dataloader",
     "FastModelLoader",
     "TrainModel",
-    "PostProccess",
+    "PostProcess",
 ]

yolo/lazy.py

@@ -14,29 +14,30 @@
 
 @hydra.main(config_path="config", config_name="config", version_base=None)
 def main(cfg: Config):
-    callbacks, loggers = setup(cfg)
+    callbacks, loggers, save_path = setup(cfg)
 
     trainer = Trainer(
-        accelerator="cuda",
+        accelerator="auto",
         max_epochs=getattr(cfg.task, "epoch", None),
         precision="16-mixed",
         callbacks=callbacks,
         logger=loggers,
         log_every_n_steps=1,
         gradient_clip_val=10,
         deterministic=True,
+        enable_progress_bar=not getattr(cfg, "quite", False),
+        default_root_dir=save_path,
     )
 
-    match cfg.task.task:
-        case "train":
-            model = TrainModel(cfg)
-            trainer.fit(model)
-        case "validation":
-            model = ValidateModel(cfg)
-            trainer.validate(model)
-        case "inference":
-            model = InferenceModel(cfg)
-            trainer.predict(model)
+    if cfg.task.task == "train":
+        model = TrainModel(cfg)
+        trainer.fit(model)
+    if cfg.task.task == "validation":
+        model = ValidateModel(cfg)
+        trainer.validate(model)
+    if cfg.task.task == "inference":
+        model = InferenceModel(cfg)
+        trainer.predict(model)
 
 
 if __name__ == "__main__":

yolo/tools/data_loader.py

@@ -170,7 +170,7 @@ def collate_fn(batch: List[Tuple[Tensor, Tensor]]) -> Tuple[Tensor, List[Tensor]
     """
     batch_size = len(batch)
     target_sizes = [item[1].size(0) for item in batch]
-    # TODO: Improve readability of these proccess
+    # TODO: Improve readability of these process
     # TODO: remove maxBbox or reduce loss function memory usage
     batch_targets = torch.zeros(batch_size, min(max(target_sizes), 100), 5)
     batch_targets[:, :, 0] = -1

yolo/tools/loss_functions.py

@@ -119,7 +119,7 @@ def __init__(self, cfg: Config, vec2box) -> None:
 
     def __call__(
         self, aux_predicts: List[Tensor], main_predicts: List[Tensor], targets: Tensor
-    ) -> Tuple[Tensor, Dict[str, Tensor]]:
+    ) -> Tuple[Tensor, Dict[str, float]]:
         # TODO: Need Refactor this region, make it flexible!
         aux_iou, aux_dfl, aux_cls = self.loss(aux_predicts, targets)
         main_iou, main_dfl, main_cls = self.loss(main_predicts, targets)