🔀 [Merge] branch 'Lightning'

Author: henrytsui000

tests/test_utils/test_bounding_box_utils.py

@@ -146,23 +146,64 @@ def test_anc2box_autoanchor(inference_v7_cfg: Config):
 
 
 def test_bbox_nms():
-    cls_dist = tensor(
-        [[[0.1, 0.7, 0.2], [0.6, 0.3, 0.1]], [[0.4, 0.4, 0.2], [0.5, 0.4, 0.1]]]  # Example class distribution
+    cls_dist = torch.tensor(
+        [
+            [
+                [0.7, 0.1, 0.2],  # High confidence, class 0
+                [0.3, 0.6, 0.1],  # High confidence, class 1
+                [-3.0, -2.0, -1.0],  # low confidence, class 2
+                [0.6, 0.2, 0.2],  # Medium confidence, class 0
+            ],
+            [
+                [0.55, 0.25, 0.2],  # Medium confidence, class 0
+                [-4.0, -0.5, -2.0],  # low confidence, class 1
+                [0.15, 0.2, 0.65],  # Medium confidence, class 2
+                [0.8, 0.1, 0.1],  # High confidence, class 0
+            ],
+        ],
+        dtype=float32,
     )
-    bbox = tensor(
-        [[[50, 50, 100, 100], [60, 60, 110, 110]], [[40, 40, 90, 90], [70, 70, 120, 120]]],  # Example bounding boxes
+
+    bbox = torch.tensor(
+        [
+            [
+                [0, 0, 160, 120],  # Overlaps with box 4
+                [160, 120, 320, 240],
+                [0, 120, 160, 240],
+                [16, 12, 176, 132],
+            ],
+            [
+                [0, 0, 160, 120],  # Overlaps with box 4
+                [160, 120, 320, 240],
+                [0, 120, 160, 240],
+                [16, 12, 176, 132],
+            ],
+        ],
         dtype=float32,
     )
+
     nms_cfg = NMSConfig(min_confidence=0.5, min_iou=0.5)
 
-    expected_output = [
-        tensor(
+    # Batch 1:
+    #  - box 1 is kept with class 0 as it has a higher confidence than box 4 i.e. box 4 is filtered out
+    #  - box 2 is kept with class 1
+    #  - box 3 is rejected by the confidence filter
+    # Batch 2:
+    #  - box 4 is kept with class 0 as it has a higher confidence than box 1 i.e. box 1 is filtered out
+    #  - box 2 is rejected by the confidence filter
+    #  - box 3 is kept with class 2
+    expected_output = torch.tensor(
+        [
             [
-                [1.0000, 50.0000, 50.0000, 100.0000, 100.0000, 0.6682],
-                [0.0000, 60.0000, 60.0000, 110.0000, 110.0000, 0.6457],
-            ]
-        )
-    ]
+                [0.0, 0.0, 0.0, 160.0, 120.0, 0.6682],
+                [1.0, 160.0, 120.0, 320.0, 240.0, 0.6457],
+            ],
+            [
+                [0.0, 16.0, 12.0, 176.0, 132.0, 0.6900],
+                [2.0, 0.0, 120.0, 160.0, 240.0, 0.6570],
+            ],
+        ]
+    )
 
     output = bbox_nms(cls_dist, bbox, nms_cfg)
 

yolo/__init__.py

@@ -2,29 +2,32 @@
 from yolo.model.yolo import create_model
 from yolo.tools.data_loader import AugmentationComposer, create_dataloader
 from yolo.tools.drawer import draw_bboxes
-from yolo.tools.solver import ModelTester, ModelTrainer, ModelValidator
+from yolo.tools.solver import TrainModel
 from yolo.utils.bounding_box_utils import Anc2Box, Vec2Box, bbox_nms, create_converter
 from yolo.utils.deploy_utils import FastModelLoader
-from yolo.utils.logging_utils import ProgressLogger, custom_logger
+from yolo.utils.logging_utils import (
+    ImageLogger,
+    YOLORichModelSummary,
+    YOLORichProgressBar,
+)
 from yolo.utils.model_utils import PostProccess
 
 all = [
     "create_model",
     "Config",
-    "ProgressLogger",
+    "YOLORichProgressBar",
     "NMSConfig",
-    "custom_logger",
+    "YOLORichModelSummary",
     "validate_log_directory",
     "draw_bboxes",
     "Vec2Box",
     "Anc2Box",
     "bbox_nms",
     "create_converter",
     "AugmentationComposer",
+    "ImageLogger",
     "create_dataloader",
     "FastModelLoader",
-    "ModelTester",
-    "ModelTrainer",
-    "ModelValidator",
+    "TrainModel",
     "PostProccess",
 ]

yolo/config/general.yaml

@@ -7,7 +7,7 @@ out_path: runs
 exist_ok: True
 
 lucky_number: 10
-use_wandb: False
+use_wandb: True
 use_tensorboard: False
 
 weight: True # Path to weight or True for auto, False for no pretrained weight

yolo/config/task/inference.yaml

@@ -8,4 +8,4 @@ data:
 nms:
   min_confidence: 0.5
   min_iou: 0.5
-# save_predict: True
+save_predict: True

yolo/config/task/validation.yaml

@@ -8,5 +8,5 @@ data:
   pin_memory: True
   data_augment: {}
 nms:
-  min_confidence: 0.05
-  min_iou: 0.9
+  min_confidence: 0.0001
+  min_iou: 0.7

yolo/lazy.py

@@ -2,41 +2,41 @@
 from pathlib import Path
 
 import hydra
+from lightning import Trainer
 
 project_root = Path(__file__).resolve().parent.parent
 sys.path.append(str(project_root))
 
 from yolo.config.config import Config
-from yolo.model.yolo import create_model
-from yolo.tools.data_loader import create_dataloader
-from yolo.tools.solver import ModelTester, ModelTrainer, ModelValidator
-from yolo.utils.bounding_box_utils import create_converter
-from yolo.utils.deploy_utils import FastModelLoader
-from yolo.utils.logging_utils import ProgressLogger
-from yolo.utils.model_utils import get_device
+from yolo.tools.solver import InferenceModel, TrainModel, ValidateModel
+from yolo.utils.logging_utils import setup
 
 
 @hydra.main(config_path="config", config_name="config", version_base=None)
 def main(cfg: Config):
-    progress = ProgressLogger(cfg, exp_name=cfg.name)
-    device, use_ddp = get_device(cfg.device)
-    dataloader = create_dataloader(cfg.task.data, cfg.dataset, cfg.task.task, use_ddp)
-    if getattr(cfg.task, "fast_inference", False):
-        model = FastModelLoader(cfg).load_model(device)
-    else:
-        model = create_model(cfg.model, class_num=cfg.dataset.class_num, weight_path=cfg.weight)
-        model = model.to(device)
-
-    converter = create_converter(cfg.model.name, model, cfg.model.anchor, cfg.image_size, device)
-
-    if cfg.task.task == "train":
-        solver = ModelTrainer(cfg, model, converter, progress, device, use_ddp)
-    if cfg.task.task == "validation":
-        solver = ModelValidator(cfg.task, cfg.dataset, model, converter, progress, device)
-    if cfg.task.task == "inference":
-        solver = ModelTester(cfg, model, converter, progress, device)
-    progress.start()
-    solver.solve(dataloader)
+    callbacks, loggers = setup(cfg)
+
+    trainer = Trainer(
+        accelerator="cuda",
+        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,
+    )
+
+    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 __name__ == "__main__":

yolo/model/module.py

@@ -3,10 +3,10 @@
 import torch
 import torch.nn.functional as F
 from einops import rearrange
-from loguru import logger
 from torch import Tensor, nn
 from torch.nn.common_types import _size_2_t
 
+from yolo.utils.logger import logger
 from yolo.utils.module_utils import auto_pad, create_activation_function, round_up
 
 

yolo/model/yolo.py

@@ -3,12 +3,12 @@
 from typing import Dict, List, Union
 
 import torch
-from loguru import logger
 from omegaconf import ListConfig, OmegaConf
 from torch import nn
 
 from yolo.config.config import ModelConfig, YOLOLayer
 from yolo.tools.dataset_preparation import prepare_weight
+from yolo.utils.logger import logger
 from yolo.utils.module_utils import get_layer_map
 
 
@@ -32,10 +32,10 @@ def __init__(self, model_cfg: ModelConfig, class_num: int = 80):
     def build_model(self, model_arch: Dict[str, List[Dict[str, Dict[str, Dict]]]]):
         self.layer_index = {}
         output_dim, layer_idx = [3], 1
-        logger.info(f"🚜 Building YOLO")
+        logger.info(f":tractor: Building YOLO")
         for arch_name in model_arch:
             if model_arch[arch_name]:
-                logger.info(f"  🏗️  Building {arch_name}")
+                logger.info(f"  :building_construction:  Building {arch_name}")
             for layer_idx, layer_spec in enumerate(model_arch[arch_name], start=layer_idx):
                 layer_type, layer_info = next(iter(layer_spec.items()))
                 layer_args = layer_info.get("args", {})
@@ -123,7 +123,7 @@ def save_load_weights(self, weights: Union[Path, OrderedDict]):
             weights: A OrderedDict containing the new weights.
         """
         if isinstance(weights, Path):
-            weights = torch.load(weights, map_location=torch.device("cpu"))
+            weights = torch.load(weights, map_location=torch.device("cpu"), weights_only=False)
         if "model_state_dict" in weights:
             weights = weights["model_state_dict"]
 
@@ -144,7 +144,7 @@ def save_load_weights(self, weights: Union[Path, OrderedDict]):
 
         for error_name, error_set in error_dict.items():
             for weight_name in error_set:
-                logger.warning(f"⚠️ Weight {error_name} for key: {'.'.join(weight_name)}")
+                logger.warning(f":warning: Weight {error_name} for key: {'.'.join(weight_name)}")
 
         self.model.load_state_dict(model_state_dict)
 
@@ -171,7 +171,7 @@ def create_model(model_cfg: ModelConfig, weight_path: Union[bool, Path] = True,
             prepare_weight(weight_path=weight_path)
         if weight_path.exists():
             model.save_load_weights(weight_path)
-            logger.info("✅ Success load model & weight")
+            logger.info(":white_check_mark: Success load model & weight")
     else:
-        logger.info("✅ Success load model")
+        logger.info(":white_check_mark: Success load model")
     return model