🔀 [Merge] branch 'Lightning'

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,18 +2,22 @@ from yolo.config.config import Config, NMSConfig
 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",
@@ -21,10 +25,9 @@ all = [
     "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 @@ import sys
 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 @@ from typing import Any, Dict, List, Optional, Tuple
 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 pathlib import Path
 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 @@ class YOLO(nn.Module):
     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 @@ class YOLO(nn.Module):
             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 @@ class YOLO(nn.Module):
 
         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

yolo/tools/data_loader.py

@@ -5,12 +5,10 @@ from typing import Generator, List, Tuple, Union
 
 import numpy as np
 import torch
-from loguru import logger
 from PIL import Image
 from rich.progress import track
 from torch import Tensor
 from torch.utils.data import DataLoader, Dataset
-from torch.utils.data.distributed import DistributedSampler
 
 from yolo.config.config import DataConfig, DatasetConfig
 from yolo.tools.data_augmentation import *
@@ -20,7 +18,9 @@ from yolo.utils.dataset_utils import (
     create_image_metadata,
     locate_label_paths,
     scale_segmentation,
+    tensorlize,
 )
+from yolo.utils.logger import logger
 
 
 class YoloDataset(Dataset):
@@ -32,7 +32,8 @@ class YoloDataset(Dataset):
         transforms = [eval(aug)(prob) for aug, prob in augment_cfg.items()]
         self.transform = AugmentationComposer(transforms, self.image_size)
         self.transform.get_more_data = self.get_more_data
-        self.data = self.load_data(Path(dataset_cfg.path), phase_name)
+        img_paths, bboxes = tensorlize(self.load_data(Path(dataset_cfg.path), phase_name))
+        self.img_paths, self.bboxes = img_paths, bboxes
 
     def load_data(self, dataset_path: Path, phase_name: str):
         """
@@ -45,15 +46,15 @@ class YoloDataset(Dataset):
         Returns:
             dict: The loaded data from the cache for the specified phase.
         """
-        cache_path = dataset_path / f"{phase_name}.cache"
+        cache_path = dataset_path / f"{phase_name}.cache1"
 
         if not cache_path.exists():
-            logger.info("🏭 Generating {} cache", phase_name)
+            logger.info(f":factory: Generating {phase_name} cache")
             data = self.filter_data(dataset_path, phase_name)
             torch.save(data, cache_path)
         else:
             data = torch.load(cache_path, weights_only=False)
-            logger.info("📦 Loaded {} cache", phase_name)
+            logger.info(f":package: Loaded {phase_name} cache")
         return data
 
     def filter_data(self, dataset_path: Path, phase_name: str) -> list:
@@ -103,7 +104,7 @@ class YoloDataset(Dataset):
             img_path = images_path / image_name
             data.append((img_path, labels))
             valid_inputs += 1
-        logger.info("Recorded {}/{} valid inputs", valid_inputs, len(images_list))
+        logger.info(f"Recorded {valid_inputs}/{len(images_list)} valid inputs")
         return data
 
     def load_valid_labels(self, label_path: str, seg_data_one_img: list) -> Union[Tensor, None]:
@@ -132,9 +133,11 @@ class YoloDataset(Dataset):
             return torch.zeros((0, 5))
 
     def get_data(self, idx):
-        img_path, bboxes = self.data[idx]
-        img = Image.open(img_path).convert("RGB")
-        return img, bboxes, img_path
+        img_path, bboxes = self.img_paths[idx], self.bboxes[idx]
+        valid_mask = bboxes[:, 0] != -1
+        with Image.open(img_path) as img:
+            img = img.convert("RGB")
+        return img, torch.from_numpy(bboxes[valid_mask]), img_path
 
     def get_more_data(self, num: int = 1):
         indices = torch.randint(0, len(self), (num,))
@@ -143,67 +146,59 @@ class YoloDataset(Dataset):
     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)
+        bboxes[:, [1, 3]] *= self.image_size[0]
+        bboxes[:, [2, 4]] *= self.image_size[1]
         return img, bboxes, rev_tensor, img_path
 
     def __len__(self) -> int:
-        return len(self.data)
-
-
-class YoloDataLoader(DataLoader):
-    def __init__(self, data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: str = "train", use_ddp: bool = False):
-        """Initializes the YoloDataLoader with hydra-config files."""
-        dataset = YoloDataset(data_cfg, dataset_cfg, task)
-        sampler = DistributedSampler(dataset, shuffle=data_cfg.shuffle) if use_ddp else None
-        self.image_size = data_cfg.image_size[0]
-        super().__init__(
-            dataset,
-            batch_size=data_cfg.batch_size,
-            sampler=sampler,
-            shuffle=data_cfg.shuffle and not use_ddp,
-            num_workers=data_cfg.cpu_num,
-            pin_memory=data_cfg.pin_memory,
-            collate_fn=self.collate_fn,
-        )
-
-    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.
+        return len(self.bboxes)
 
-        Args:
-            batch (list of tuples): Each tuple contains:
-                - image (Tensor): The image tensor.
-                - labels (Tensor): The tensor of labels for the image.
 
-        Returns:
-            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.
-        """
-        batch_size = len(batch)
-        target_sizes = [item[1].size(0) for item in batch]
-        # TODO: Improve readability of these proccess
-        # 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
-        for idx, target_size in enumerate(target_sizes):
-            batch_targets[idx, : min(target_size, 100)] = batch[idx][1][:100]
-        batch_targets[:, :, 1:] *= self.image_size
+def collate_fn(batch: List[Tuple[Tensor, Tensor]]) -> Tuple[Tensor, List[Tensor]]:
+    """
+    A collate function to handle batching of images and their corresponding targets.
 
-        batch_images, _, batch_reverse, batch_path = zip(*batch)
-        batch_images = torch.stack(batch_images)
-        batch_reverse = torch.stack(batch_reverse)
+    Args:
+        batch (list of tuples): Each tuple contains:
+            - image (Tensor): The image tensor.
+            - labels (Tensor): The tensor of labels for the image.
 
-        return batch_size, batch_images, batch_targets, batch_reverse, batch_path
+    Returns:
+        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.
+    """
+    batch_size = len(batch)
+    target_sizes = [item[1].size(0) for item in batch]
+    # TODO: Improve readability of these proccess
+    # 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
+    for idx, target_size in enumerate(target_sizes):
+        batch_targets[idx, : min(target_size, 100)] = batch[idx][1][:100]
 
+    batch_images, _, batch_reverse, batch_path = zip(*batch)
+    batch_images = torch.stack(batch_images)
+    batch_reverse = torch.stack(batch_reverse)
 
-def create_dataloader(data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: str = "train", use_ddp: bool = False):
+    return batch_size, batch_images, batch_targets, batch_reverse, batch_path
+
+
+def create_dataloader(data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: str = "train"):
     if task == "inference":
         return StreamDataLoader(data_cfg)
 
     if dataset_cfg.auto_download:
         prepare_dataset(dataset_cfg, task)
-
-    return YoloDataLoader(data_cfg, dataset_cfg, task, use_ddp)
+    dataset = YoloDataset(data_cfg, dataset_cfg, task)
+
+    return DataLoader(
+        dataset,
+        batch_size=data_cfg.batch_size,
+        num_workers=data_cfg.cpu_num,
+        pin_memory=data_cfg.pin_memory,
+        collate_fn=collate_fn,
+    )
 
 
 class StreamDataLoader:

yolo/tools/dataset_preparation.py

@@ -3,10 +3,10 @@ from pathlib import Path
 from typing import Optional
 
 import requests
-from loguru import logger
 from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
 
 from yolo.config.config import DatasetConfig
+from yolo.utils.logger import logger
 
 
 def download_file(url, destination: Path):
@@ -30,7 +30,7 @@ def download_file(url, destination: Path):
                 for data in response.iter_content(chunk_size=1024 * 1024):  # 1 MB chunks
                     file.write(data)
                     progress.update(task, advance=len(data))
-    logger.info("✅ Download completed.")
+    logger.info(":white_check_mark: Download completed.")
 
 
 def unzip_file(source: Path, destination: Path):
@@ -71,7 +71,7 @@ def prepare_dataset(dataset_cfg: DatasetConfig, task: str):
 
             final_place.mkdir(parents=True, exist_ok=True)
             if check_files(final_place, dataset_args.get("file_num")):
-                logger.info(f"✅ Dataset {dataset_type: <12} already verified.")
+                logger.info(f":white_check_mark: Dataset {dataset_type: <12} already verified.")
                 continue
 
             if not local_zip_path.exists():

yolo/tools/drawer.py

@@ -3,12 +3,12 @@ from typing import List, Optional, Union
 
 import numpy as np
 import torch
-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
+from yolo.utils.logger import logger
 
 
 def draw_bboxes(
@@ -121,6 +121,6 @@ def draw_model(*, model_cfg: ModelConfig = None, model: YOLO = None, v7_base=Fal
                 dot.edge(str(idx), str(jdx))
     try:
         dot.render("Model-arch", format="png", cleanup=True)
-        logger.info("🎨 Drawing Model Architecture at Model-arch.png")
+        logger.info(":artist_palette: Drawing Model Architecture at Model-arch.png")
     except:
-        logger.warning("⚠️ Could not find graphviz backend, continue without drawing the model architecture")
+        logger.warning(":warning: Could not find graphviz backend, continue without drawing the model architecture")

yolo/tools/loss_functions.py

@@ -2,12 +2,12 @@ from typing import Any, Dict, List, Tuple
 
 import torch
 import torch.nn.functional as F
-from loguru import logger
 from torch import Tensor, nn
 from torch.nn import BCEWithLogitsLoss
 
 from yolo.config.config import Config, LossConfig
 from yolo.utils.bounding_box_utils import BoxMatcher, Vec2Box, calculate_iou
+from yolo.utils.logger import logger
 
 
 class BCELoss(nn.Module):
@@ -124,17 +124,19 @@ class DualLoss:
         aux_iou, aux_dfl, aux_cls = self.loss(aux_predicts, targets)
         main_iou, main_dfl, main_cls = self.loss(main_predicts, targets)
 
+        total_loss = [
+            self.iou_rate * (aux_iou * self.aux_rate + main_iou),
+            self.dfl_rate * (aux_dfl * self.aux_rate + main_dfl),
+            self.cls_rate * (aux_cls * self.aux_rate + main_cls),
+        ]
         loss_dict = {
-            "BoxLoss": self.iou_rate * (aux_iou * self.aux_rate + main_iou),
-            "DFLoss": self.dfl_rate * (aux_dfl * self.aux_rate + main_dfl),
-            "BCELoss": self.cls_rate * (aux_cls * self.aux_rate + main_cls),
+            f"Loss/{name}Loss": value.detach().item() for name, value in zip(["Box", "DFL", "BCE"], total_loss)
         }
-        loss_sum = sum(list(loss_dict.values())) / len(loss_dict)
-        return loss_sum, loss_dict
+        return sum(total_loss), loss_dict
 
 
 def create_loss_function(cfg: Config, vec2box) -> DualLoss:
     # TODO: make it flexible, if cfg doesn't contain aux, only use SingleLoss
     loss_function = DualLoss(cfg, vec2box)
-    logger.info("✅ Success load loss function")
+    logger.info(":white_check_mark: Success load loss function")
     return loss_function

yolo/tools/solver.py

@@ -1,267 +1,154 @@
-import contextlib
-import io
-import json
-import os
 import time
-from collections import defaultdict
 from pathlib import Path
-from typing import Dict, Optional
 
-import torch
-from loguru import logger
-from pycocotools.coco import COCO
-from torch import Tensor, distributed
-from torch.cuda.amp import GradScaler, autocast
-from torch.nn.parallel import DistributedDataParallel as DDP
-from torch.utils.data import DataLoader
+import cv2
+import numpy as np
+from lightning import LightningModule
+from torchmetrics.detection import MeanAveragePrecision
 
-from yolo.config.config import Config, DatasetConfig, TrainConfig, ValidationConfig
-from yolo.model.yolo import YOLO
-from yolo.tools.data_loader import StreamDataLoader, create_dataloader
-from yolo.tools.drawer import draw_bboxes, draw_model
+from yolo.config.config import Config
+from yolo.model.yolo import create_model
+from yolo.tools.data_loader import create_dataloader
+from yolo.tools.drawer import draw_bboxes
 from yolo.tools.loss_functions import create_loss_function
-from yolo.utils.bounding_box_utils import Vec2Box, calculate_map
-from yolo.utils.dataset_utils import locate_label_paths
-from yolo.utils.logging_utils import ProgressLogger, log_model_structure
-from yolo.utils.model_utils import (
-    ExponentialMovingAverage,
-    PostProccess,
-    collect_prediction,
-    create_optimizer,
-    create_scheduler,
-    predicts_to_json,
-)
-from yolo.utils.solver_utils import calculate_ap
+from yolo.utils.bounding_box_utils import create_converter, to_metrics_format
+from yolo.utils.model_utils import PostProccess, create_optimizer, create_scheduler
 
 
-class ModelTrainer:
-    def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, progress: ProgressLogger, device, use_ddp: bool):
-        train_cfg: TrainConfig = cfg.task
-        self.model = model if not use_ddp else DDP(model, device_ids=[device])
-        self.use_ddp = use_ddp
-        self.vec2box = vec2box
-        self.device = device
-        self.optimizer = create_optimizer(model, train_cfg.optimizer)
-        self.scheduler = create_scheduler(self.optimizer, train_cfg.scheduler)
-        self.loss_fn = create_loss_function(cfg, vec2box)
-        self.progress = progress
-        self.num_epochs = cfg.task.epoch
-        self.mAPs_dict = defaultdict(list)
+class BaseModel(LightningModule):
+    def __init__(self, cfg: Config):
+        super().__init__()
+        self.model = create_model(cfg.model, class_num=cfg.dataset.class_num, weight_path=cfg.weight)
 
-        self.weights_dir = self.progress.save_path / "weights"
-        self.weights_dir.mkdir(exist_ok=True)
+    def forward(self, x):
+        return self.model(x)
 
-        if not progress.quite_mode:
-            log_model_structure(model.model)
-            draw_model(model=model)
 
-        self.validation_dataloader = create_dataloader(
-            cfg.task.validation.data, cfg.dataset, cfg.task.validation.task, use_ddp
-        )
-        self.validator = ModelValidator(cfg.task.validation, cfg.dataset, model, vec2box, progress, device)
-
-        if getattr(train_cfg.ema, "enabled", False):
-            self.ema = ExponentialMovingAverage(model, decay=train_cfg.ema.decay)
+class ValidateModel(BaseModel):
+    def __init__(self, cfg: Config):
+        super().__init__(cfg)
+        self.cfg = cfg
+        if self.cfg.task.task == "validation":
+            self.validation_cfg = self.cfg.task
         else:
-            self.ema = None
-        self.scaler = GradScaler()
-
-    def train_one_batch(self, images: Tensor, targets: Tensor):
-        images, targets = images.to(self.device), targets.to(self.device)
-        self.optimizer.zero_grad()
-
-        with autocast():
-            predicts = self.model(images)
-            aux_predicts = self.vec2box(predicts["AUX"])
-            main_predicts = self.vec2box(predicts["Main"])
-            loss, loss_item = self.loss_fn(aux_predicts, main_predicts, targets)
-
-        self.scaler.scale(loss).backward()
-        self.scaler.unscale_(self.optimizer)
-        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=10.0)
-        self.scaler.step(self.optimizer)
-        self.scaler.update()
-
-        return loss_item
-
-    def train_one_epoch(self, dataloader):
-        self.model.train()
-        total_loss = defaultdict(float)
-        total_samples = 0
-        self.optimizer.next_epoch(len(dataloader))
-        for batch_size, images, targets, *_ in dataloader:
-            self.optimizer.next_batch()
-            loss_each = self.train_one_batch(images, targets)
-
-            for loss_name, loss_val in loss_each.items():
-                if self.use_ddp:  # collecting loss for each batch
-                    distributed.all_reduce(loss_val, op=distributed.ReduceOp.AVG)
-                total_loss[loss_name] += loss_val.item() * batch_size
-            total_samples += batch_size
-            self.progress.one_batch(loss_each)
-
-        for loss_val in total_loss.values():
-            loss_val /= total_samples
-
-        if self.scheduler:
-            self.scheduler.step()
-
-        return total_loss
-
-    def save_checkpoint(self, epoch_idx: int, file_name: Optional[str] = None):
-        file_name = file_name or f"E{epoch_idx:03d}.pt"
-        file_path = self.weights_dir / file_name
-
-        checkpoint = {
-            "epoch": epoch_idx,
-            "model_state_dict": self.model.state_dict(),
-            "optimizer_state_dict": self.optimizer.state_dict(),
-        }
-        if self.ema:
-            self.ema.apply_shadow()
-            checkpoint["model_state_dict_ema"] = self.model.state_dict()
-            self.ema.restore()
-
-        logger.info(f"💾 success save at {file_path}")
-        torch.save(checkpoint, file_path)
-
-    def good_epoch(self, mAPs: Dict[str, Tensor]) -> bool:
-        save_flag = True
-        for mAP_key, mAP_val in mAPs.items():
-            self.mAPs_dict[mAP_key].append(mAP_val)
-            if mAP_val < max(self.mAPs_dict[mAP_key]):
-                save_flag = False
-        return save_flag
-
-    def solve(self, dataloader: DataLoader):
-        logger.info("🚄 Start Training!")
-        num_epochs = self.num_epochs
-
-        self.progress.start_train(num_epochs)
-        for epoch_idx in range(num_epochs):
-            if self.use_ddp:
-                dataloader.sampler.set_epoch(epoch_idx)
-
-            self.progress.start_one_epoch(len(dataloader), "Train", self.optimizer, epoch_idx)
-            epoch_loss = self.train_one_epoch(dataloader)
-            self.progress.finish_one_epoch(epoch_loss, epoch_idx=epoch_idx)
-
-            mAPs = self.validator.solve(self.validation_dataloader, epoch_idx=epoch_idx)
-            if mAPs is not None and self.good_epoch(mAPs):
-                self.save_checkpoint(epoch_idx=epoch_idx)
-            # TODO: save model if result are better than before
-        self.progress.finish_train()
-
-
-class ModelTester:
-    def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, progress: ProgressLogger, device):
-        self.model = model
-        self.device = device
-        self.progress = progress
-
-        self.post_proccess = PostProccess(vec2box, cfg.task.nms)
-        self.save_path = progress.save_path / "images"
-        os.makedirs(self.save_path, exist_ok=True)
-        self.save_predict = getattr(cfg.task, "save_predict", None)
-        self.idx2label = cfg.dataset.class_list
-
-    def solve(self, dataloader: StreamDataLoader):
-        logger.info("👀 Start Inference!")
-        if isinstance(self.model, torch.nn.Module):
-            self.model.eval()
-
-        if dataloader.is_stream:
-            import cv2
-            import numpy as np
-
-            last_time = time.time()
-        try:
-            for idx, (images, rev_tensor, origin_frame) in enumerate(dataloader):
-                images = images.to(self.device)
-                rev_tensor = rev_tensor.to(self.device)
-                with torch.no_grad():
-                    predicts = self.model(images)
-                    predicts = self.post_proccess(predicts, rev_tensor)
-                img = draw_bboxes(origin_frame, predicts, idx2label=self.idx2label)
-
-                if dataloader.is_stream:
-                    img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
-                    fps = 1 / (time.time() - last_time)
-                    cv2.putText(img, f"FPS: {fps:.2f}", (0, 15), 0, 0.5, (100, 255, 0), 1, cv2.LINE_AA)
-                    last_time = time.time()
-                    cv2.imshow("Prediction", img)
-                    if cv2.waitKey(1) & 0xFF == ord("q"):
-                        break
-                    if not self.save_predict:
-                        continue
-                if self.save_predict != False:
-                    save_image_path = self.save_path / f"frame{idx:03d}.png"
-                    img.save(save_image_path)
-                    logger.info(f"💾 Saved visualize image at {save_image_path}")
-
-        except (KeyboardInterrupt, Exception) as e:
-            dataloader.stop_event.set()
-            dataloader.stop()
-            if isinstance(e, KeyboardInterrupt):
-                logger.error("User Keyboard Interrupt")
-            else:
-                raise e
-        dataloader.stop()
+            self.validation_cfg = self.cfg.task.validation
+        self.metric = MeanAveragePrecision(iou_type="bbox", box_format="xyxy")
+        self.metric.warn_on_many_detections = False
+        self.val_loader = create_dataloader(self.validation_cfg.data, self.cfg.dataset, self.validation_cfg.task)
+
+    def setup(self, stage):
+        self.vec2box = create_converter(
+            self.cfg.model.name, self.model, self.cfg.model.anchor, self.cfg.image_size, self.device
+        )
+        self.post_proccess = PostProccess(self.vec2box, self.validation_cfg.nms)
 
+    def val_dataloader(self):
+        return self.val_loader
 
-class ModelValidator:
-    def __init__(
-        self,
-        validation_cfg: ValidationConfig,
-        dataset_cfg: DatasetConfig,
-        model: YOLO,
-        vec2box: Vec2Box,
-        progress: ProgressLogger,
-        device,
-    ):
-        self.model = model
-        self.device = device
-        self.progress = progress
+    def validation_step(self, batch, batch_idx):
+        batch_size, images, targets, rev_tensor, img_paths = batch
+        predicts = self.post_proccess(self(images))
+        batch_metrics = self.metric(
+            [to_metrics_format(predict) for predict in predicts], [to_metrics_format(target) for target in targets]
+        )
 
-        self.post_proccess = PostProccess(vec2box, validation_cfg.nms)
-        self.json_path = self.progress.save_path / "predict.json"
+        self.log_dict(
+            {
+                "map": batch_metrics["map"],
+                "map_50": batch_metrics["map_50"],
+            },
+            on_step=True,
+            batch_size=batch_size,
+        )
+        return predicts
+
+    def on_validation_epoch_end(self):
+        epoch_metrics = self.metric.compute()
+        del epoch_metrics["classes"]
+        self.log_dict(epoch_metrics, prog_bar=True, rank_zero_only=True)
+        self.log_dict(
+            {"PyCOCO/AP @ .5:.95": epoch_metrics["map"], "PyCOCO/AP @ .5": epoch_metrics["map_50"]}, rank_zero_only=True
+        )
+        self.metric.reset()
+
+
+class TrainModel(ValidateModel):
+    def __init__(self, cfg: Config):
+        super().__init__(cfg)
+        self.cfg = cfg
+        self.train_loader = create_dataloader(self.cfg.task.data, self.cfg.dataset, self.cfg.task.task)
+
+    def setup(self, stage):
+        super().setup(stage)
+        self.loss_fn = create_loss_function(self.cfg, self.vec2box)
+
+    def train_dataloader(self):
+        return self.train_loader
+
+    def on_train_epoch_start(self):
+        self.trainer.optimizers[0].next_epoch(len(self.train_loader))
+
+    def training_step(self, batch, batch_idx):
+        lr_dict = self.trainer.optimizers[0].next_batch()
+        batch_size, images, targets, *_ = batch
+        predicts = self(images)
+        aux_predicts = self.vec2box(predicts["AUX"])
+        main_predicts = self.vec2box(predicts["Main"])
+        loss, loss_item = self.loss_fn(aux_predicts, main_predicts, targets)
+        self.log_dict(
+            loss_item,
+            prog_bar=True,
+            on_epoch=True,
+            batch_size=batch_size,
+            rank_zero_only=True,
+        )
+        self.log_dict(lr_dict, prog_bar=False, logger=True, on_epoch=False, rank_zero_only=True)
+        return loss * batch_size
 
-        with contextlib.redirect_stdout(io.StringIO()):
-            # TODO: load with config file
-            json_path, _ = locate_label_paths(Path(dataset_cfg.path), dataset_cfg.get("validation", "val"))
-            if json_path:
-                self.coco_gt = COCO(json_path)
+    def configure_optimizers(self):
+        optimizer = create_optimizer(self.model, self.cfg.task.optimizer)
+        scheduler = create_scheduler(optimizer, self.cfg.task.scheduler)
+        return [optimizer], [scheduler]
 
-    def solve(self, dataloader, epoch_idx=1):
-        # logger.info("🧪 Start Validation!")
-        self.model.eval()
-        predict_json, mAPs = [], defaultdict(list)
-        self.progress.start_one_epoch(len(dataloader), task="Validate")
-        for batch_size, images, targets, rev_tensor, img_paths in dataloader:
-            images, targets, rev_tensor = images.to(self.device), targets.to(self.device), rev_tensor.to(self.device)
-            with torch.no_grad():
-                predicts = self.model(images)
-                predicts = self.post_proccess(predicts)
-                for idx, predict in enumerate(predicts):
-                    mAP = calculate_map(predict, targets[idx])
-                    for mAP_key, mAP_val in mAP.items():
-                        mAPs[mAP_key].append(mAP_val)
 
-            avg_mAPs = {key: 100 * torch.mean(torch.stack(val)) for key, val in mAPs.items()}
-            self.progress.one_batch(avg_mAPs)
+class InferenceModel(BaseModel):
+    def __init__(self, cfg: Config):
+        super().__init__(cfg)
+        self.cfg = cfg
+        # TODO: Add FastModel
+        self.predict_loader = create_dataloader(cfg.task.data, cfg.dataset, cfg.task.task)
 
-            predict_json.extend(predicts_to_json(img_paths, predicts, rev_tensor))
-        self.progress.finish_one_epoch(avg_mAPs, epoch_idx=epoch_idx)
-        self.progress.visualize_image(images, targets, predicts, epoch_idx=epoch_idx)
+    def setup(self, stage):
+        self.vec2box = create_converter(
+            self.cfg.model.name, self.model, self.cfg.model.anchor, self.cfg.image_size, self.device
+        )
+        self.post_process = PostProcess(self.vec2box, self.cfg.task.nms)
 
-        with open(self.json_path, "w") as f:
-            predict_json = collect_prediction(predict_json, self.progress.local_rank)
-            if self.progress.local_rank != 0:
-                return
-            json.dump(predict_json, f)
-        if hasattr(self, "coco_gt"):
-            self.progress.start_pycocotools()
-            result = calculate_ap(self.coco_gt, predict_json)
-            self.progress.finish_pycocotools(result, epoch_idx)
+    def predict_dataloader(self):
+        return self.predict_loader
 
-        return avg_mAPs
+    def predict_step(self, batch, batch_idx):
+        images, rev_tensor, origin_frame = batch
+        predicts = self.post_process(self(images), rev_tensor)
+        img = draw_bboxes(origin_frame, predicts, idx2label=self.cfg.dataset.class_list)
+        if getattr(self.predict_loader, "is_stream", None):
+            fps = self._display_stream(img)
+        else:
+            fps = None
+        if getattr(self.cfg.task, "save_predict", None):
+            self._save_image(img, batch_idx)
+        return img, fps
+
+    def _display_stream(self, img):
+        img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+        fps = 1 / (time.time() - self.trainer.current_epoch_start_time)
+        cv2.putText(img, f"FPS: {fps:.2f}", (0, 15), 0, 0.5, (100, 255, 0), 1, cv2.LINE_AA)
+        cv2.imshow("Prediction", img)
+        if cv2.waitKey(1) & 0xFF == ord("q"):
+            self.trainer.should_stop = True
+        return fps
+
+    def _save_image(self, img, batch_idx):
+        save_image_path = Path(self.logger.save_dir) / f"frame{batch_idx:03d}.png"
+        img.save(save_image_path)
+        print(f"💾 Saved visualize image at {save_image_path}")

yolo/utils/bounding_box_utils.py

@@ -4,17 +4,17 @@ from typing import Dict, List, Optional, Tuple, Union
 import torch
 import torch.nn.functional as F
 from einops import rearrange
-from loguru import logger
 from torch import Tensor, arange, tensor
 from torchvision.ops import batched_nms
 
 from yolo.config.config import AnchorConfig, MatcherConfig, ModelConfig, NMSConfig
 from yolo.model.yolo import YOLO
+from yolo.utils.logger import logger
 
 
 def calculate_iou(bbox1, bbox2, metrics="iou") -> Tensor:
     metrics = metrics.lower()
-    EPS = 1e-9
+    EPS = 1e-7
     dtype = bbox1.dtype
     bbox1 = bbox1.to(torch.float32)
     bbox2 = bbox2.to(torch.float32)
@@ -69,7 +69,8 @@ def calculate_iou(bbox1, bbox2, metrics="iou") -> Tensor:
         (bbox2[..., 2] - bbox2[..., 0]) / (bbox2[..., 3] - bbox2[..., 1] + EPS)
     )
     v = (4 / (math.pi**2)) * (arctan**2)
-    alpha = v / (v - iou + 1 + EPS)
+    with torch.no_grad():
+        alpha = v / (v - iou + 1 + EPS)
     # Compute CIoU
     ciou = diou - alpha * v
     return ciou.to(dtype)
@@ -129,7 +130,10 @@ def generate_anchors(image_size: List[int], strides: List[int]):
         shift = stride // 2
         h = torch.arange(0, H, stride) + shift
         w = torch.arange(0, W, stride) + shift
-        anchor_h, anchor_w = torch.meshgrid(h, w, indexing="ij")
+        if torch.__version__ >= "2.3.0":
+            anchor_h, anchor_w = torch.meshgrid(h, w, indexing="ij")
+        else:
+            anchor_h, anchor_w = torch.meshgrid(h, w)
         anchor = torch.stack([anchor_w.flatten(), anchor_h.flatten()], dim=-1)
         anchors.append(anchor)
     all_anchors = torch.cat(anchors, dim=0)
@@ -207,7 +211,7 @@ class BoxMatcher:
         topk_masks = topk_targets > 0
         return topk_targets, topk_masks
 
-    def filter_duplicates(self, target_matrix: Tensor):
+    def filter_duplicates(self, target_matrix: Tensor, topk_mask: Tensor):
         """
         Filter the maximum suitability target index of each anchor.
 
@@ -217,9 +221,11 @@ class BoxMatcher:
         Returns:
             unique_indices [batch x anchors x 1]: The index of the best targets for each anchors
         """
-        # TODO: add a assert for no target on the image
-        unique_indices = target_matrix.argmax(dim=1)
-        return unique_indices[..., None]
+        duplicates = (topk_mask.sum(1, keepdim=True) > 1).repeat([1, topk_mask.size(1), 1])
+        max_idx = F.one_hot(target_matrix.argmax(1), topk_mask.size(1)).permute(0, 2, 1)
+        topk_mask = torch.where(duplicates, max_idx, topk_mask)
+        unique_indices = topk_mask.argmax(dim=1)
+        return unique_indices[..., None], topk_mask.sum(1), topk_mask
 
     def __call__(self, target: Tensor, predict: Tuple[Tensor]) -> Tuple[Tensor, Tensor]:
         """Matches each target to the most suitable anchor.
@@ -271,16 +277,15 @@ class BoxMatcher:
         topk_targets, topk_mask = self.filter_topk(target_matrix, topk=self.topk)
 
         # delete one anchor pred assign to mutliple gts
-        unique_indices = self.filter_duplicates(topk_targets)
-
-        # TODO: do we need grid_mask? Filter the valid groud truth
-        valid_mask = (grid_mask.sum(dim=-2) * topk_mask.sum(dim=-2)).bool()
+        unique_indices, valid_mask, topk_mask = self.filter_duplicates(iou_mat, topk_mask)
 
         align_bbox = torch.gather(target_bbox, 1, unique_indices.repeat(1, 1, 4))
         align_cls = torch.gather(target_cls, 1, unique_indices).squeeze(-1)
         align_cls = F.one_hot(align_cls, self.class_num)
 
         # normalize class ditribution
+        iou_mat *= topk_mask
+        target_matrix *= topk_mask
         max_target = target_matrix.amax(dim=-1, keepdim=True)
         max_iou = iou_mat.amax(dim=-1, keepdim=True)
         normalize_term = (target_matrix / (max_target + 1e-9)) * max_iou
@@ -295,7 +300,7 @@ class Vec2Box:
         self.device = device
 
         if hasattr(anchor_cfg, "strides"):
-            logger.info(f"🈶 Found stride of model {anchor_cfg.strides}")
+            logger.info(f":japanese_not_free_of_charge_button: Found stride of model {anchor_cfg.strides}")
             self.strides = anchor_cfg.strides
         else:
             logger.info("🧸 Found no stride of model, performed a dummy test for auto-anchor size")
@@ -339,7 +344,7 @@ class Anc2Box:
         self.device = device
 
         if hasattr(anchor_cfg, "strides"):
-            logger.info(f"🈶 Found stride of model {anchor_cfg.strides}")
+            logger.info(f":japanese_not_free_of_charge_button: Found stride of model {anchor_cfg.strides}")
             self.strides = anchor_cfg.strides
         else:
             logger.info("🧸 Found no stride of model, performed a dummy test for auto-anchor size")
@@ -413,7 +418,7 @@ def bbox_nms(cls_dist: Tensor, bbox: Tensor, nms_cfg: NMSConfig, confidence: Opt
     valid_box = bbox[valid_mask.repeat(1, 1, 4)].view(-1, 4)
 
     batch_idx, *_ = torch.where(valid_mask)
-    nms_idx = batched_nms(valid_box, valid_cls, batch_idx, nms_cfg.min_iou)
+    nms_idx = batched_nms(valid_box, valid_con, batch_idx, nms_cfg.min_iou)
     predicts_nms = []
     for idx in range(cls_dist.size(0)):
         instance_idx = nms_idx[idx == batch_idx[nms_idx]]
@@ -471,3 +476,10 @@ def calculate_map(predictions, ground_truths, iou_thresholds=arange(0.5, 1, 0.05
         "mAP.5:.95": torch.mean(torch.stack(aps)),
     }
     return mAP
+
+
+def to_metrics_format(prediction: Tensor) -> Dict[str, Union[float, Tensor]]:
+    bbox = {"boxes": prediction[:, 1:5], "labels": prediction[:, 0].int()}
+    if prediction.size(1) == 6:
+        bbox["scores"] = prediction[:, 5]
+    return bbox

yolo/utils/dataset_utils.py

@@ -5,9 +5,10 @@ from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
-from loguru import logger
+import torch
 
 from yolo.tools.data_conversion import discretize_categories
+from yolo.utils.logger import logger
 
 
 def locate_label_paths(dataset_path: Path, phase_name: Path) -> Tuple[Path, Path]:
@@ -111,3 +112,16 @@ def scale_segmentation(
         seg_array_with_cat.append(scaled_flat_seg_data)
 
     return seg_array_with_cat
+
+
+def tensorlize(data):
+    img_paths, bboxes = zip(*data)
+    max_box = max(bbox.size(0) for bbox in bboxes)
+    padded_bbox_list = []
+    for bbox in bboxes:
+        padding = torch.full((max_box, 5), -1, dtype=torch.float32)
+        padding[: bbox.size(0)] = bbox
+        padded_bbox_list.append(padding)
+    bboxes = np.stack(padded_bbox_list)
+    img_paths = np.array(img_paths)
+    return img_paths, bboxes

yolo/utils/deploy_utils.py

@@ -1,11 +1,11 @@
 from pathlib import Path
 
 import torch
-from loguru import logger
 from torch import Tensor
 
 from yolo.config.config import Config
 from yolo.model.yolo import create_model
+from yolo.utils.logger import logger
 
 
 class FastModelLoader:
@@ -21,10 +21,10 @@ class FastModelLoader:
 
     def _validate_compiler(self):
         if self.compiler not in ["onnx", "trt", "deploy"]:
-            logger.warning(f"⚠️ Compiler '{self.compiler}' is not supported. Using original model.")
+            logger.warning(f":warning: Compiler '{self.compiler}' is not supported. Using original model.")
             self.compiler = None
         if self.cfg.device == "mps" and self.compiler == "trt":
-            logger.warning("🍎 TensorRT does not support MPS devices. Using original model.")
+            logger.warning(":red_apple: TensorRT does not support MPS devices. Using original model.")
             self.compiler = None
 
     def load_model(self, device):
@@ -59,7 +59,7 @@ class FastModelLoader:
             providers = ["CUDAExecutionProvider"]
         try:
             ort_session = InferenceSession(self.model_path, providers=providers)
-            logger.info("🚀 Using ONNX as MODEL frameworks!")
+            logger.info(":rocket: Using ONNX as MODEL frameworks!")
         except Exception as e:
             logger.warning(f"🈳 Error loading ONNX model: {e}")
             ort_session = self._create_onnx_model(providers)
@@ -79,7 +79,7 @@ class FastModelLoader:
             output_names=["output"],
             dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
         )
-        logger.info(f"📥 ONNX model saved to {self.model_path}")
+        logger.info(f":inbox_tray: ONNX model saved to {self.model_path}")
         return InferenceSession(self.model_path, providers=providers)
 
     def _load_trt_model(self):
@@ -88,7 +88,7 @@ class FastModelLoader:
         try:
             model_trt = TRTModule()
             model_trt.load_state_dict(torch.load(self.model_path))
-            logger.info("🚀 Using TensorRT as MODEL frameworks!")
+            logger.info(":rocket: Using TensorRT as MODEL frameworks!")
         except FileNotFoundError:
             logger.warning(f"🈳 No found model weight at {self.model_path}")
             model_trt = self._create_trt_model()
@@ -102,5 +102,5 @@ class FastModelLoader:
         logger.info(f"♻️ Creating TensorRT model")
         model_trt = torch2trt(model.cuda(), [dummy_input])
         torch.save(model_trt.state_dict(), self.model_path)
-        logger.info(f"📥 TensorRT model saved to {self.model_path}")
+        logger.info(f":inbox_tray: TensorRT model saved to {self.model_path}")
         return model_trt

yolo/utils/logger.py

@@ -0,0 +1,11 @@
+import logging
+
+from lightning.pytorch.utilities.rank_zero import rank_zero_only
+from rich.console import Console
+from rich.logging import RichHandler
+
+logger = logging.getLogger("yolo")
+logger.setLevel(logging.DEBUG)
+logger.propagate = False
+if rank_zero_only.rank == 0 and not logger.hasHandlers():
+    logger.addHandler(RichHandler(console=Console(), show_level=True, show_path=True, show_time=True, markup=True))

yolo/utils/logging_utils.py

@@ -11,55 +11,39 @@ Example:
     custom_logger()
 """
 
-import os
-import random
-import sys
+import logging
 from collections import deque
+from logging import FileHandler
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
 import wandb
-import wandb.errors.term
-from loguru import logger
+from lightning import LightningModule, Trainer, seed_everything
+from lightning.pytorch.callbacks import Callback, RichModelSummary, RichProgressBar
+from lightning.pytorch.callbacks.progress.rich_progress import CustomProgress
+from lightning.pytorch.loggers import TensorBoardLogger, WandbLogger
+from lightning.pytorch.utilities import rank_zero_only
 from omegaconf import ListConfig
+from rich import get_console, reconfigure
 from rich.console import Console, Group
-from rich.progress import (
-    BarColumn,
-    Progress,
-    SpinnerColumn,
-    TextColumn,
-    TimeRemainingColumn,
-)
+from rich.logging import RichHandler
 from rich.table import Table
+from rich.text import Text
 from torch import Tensor
 from torch.nn import ModuleList
-from torch.optim import Optimizer
-from torchvision.transforms.functional import pil_to_tensor
+from typing_extensions import override
 
 from yolo.config.config import Config, YOLOLayer
 from yolo.model.yolo import YOLO
-from yolo.tools.drawer import draw_bboxes
+from yolo.utils.logger import logger
 from yolo.utils.solver_utils import make_ap_table
 
 
-def custom_logger(quite: bool = False):
-    logger.remove()
-    if quite:
-        return
-    logger.add(
-        sys.stderr,
-        colorize=True,
-        format="<fg #003385>[{time:MM/DD HH:mm:ss}]</> <level>{level: ^8}</level>| <level>{message}</level>",
-    )
-
-
 # TODO: should be moved to correct position
 def set_seed(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
+    seed_everything(seed)
     if torch.cuda.is_available():
         torch.cuda.manual_seed(seed)
         torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
@@ -67,189 +51,220 @@ def set_seed(seed):
     torch.backends.cudnn.benchmark = False
 
 
-class ProgressLogger(Progress):
-    def __init__(self, cfg: Config, exp_name: str, *args, **kwargs):
-        set_seed(cfg.lucky_number)
-        self.local_rank = int(os.getenv("LOCAL_RANK", "0"))
-        self.quite_mode = self.local_rank or getattr(cfg, "quite", False)
-        custom_logger(self.quite_mode)
-        self.save_path = validate_log_directory(cfg, exp_name=cfg.name)
-
-        progress_bar = (
-            SpinnerColumn(),
-            TextColumn("[progress.description]{task.description}"),
-            BarColumn(bar_width=None),
-            TextColumn("{task.completed:.0f}/{task.total:.0f}"),
-            TimeRemainingColumn(),
-        )
-        self.ap_table = Table()
-        # TODO: load maxlen by config files
-        self.ap_past_list = deque(maxlen=5)
-        self.last_result = 0
-        super().__init__(*args, *progress_bar, **kwargs)
-
-        self.use_wandb = cfg.use_wandb
-        if self.use_wandb and self.local_rank == 0:
-            wandb.errors.term._log = custom_wandb_log
-            self.wandb = wandb.init(
-                project="YOLO", resume="allow", mode="online", dir=self.save_path, id=None, name=exp_name
-            )
+class YOLOCustomProgress(CustomProgress):
+    def get_renderable(self):
+        renderable = Group(*self.get_renderables())
+        if hasattr(self, "table"):
+            renderable = Group(*self.get_renderables(), self.table)
+        return renderable
 
-        self.use_tensorboard = cfg.use_tensorboard
-        if self.use_tensorboard and self.local_rank == 0:
-            from torch.utils.tensorboard import SummaryWriter
 
-            self.tb_writer = SummaryWriter(log_dir=self.save_path / "tensorboard")
-            logger.opt(colors=True).info(f"📍 Enable TensorBoard locally at <blue><u>http://localhost:6006</></>")
+class YOLORichProgressBar(RichProgressBar):
+    @override
+    @rank_zero_only
+    def _init_progress(self, trainer: "Trainer") -> None:
+        if self.is_enabled and (self.progress is None or self._progress_stopped):
+            self._reset_progress_bar_ids()
+            reconfigure(**self._console_kwargs)
+            self._console = Console()
+            self._console.clear_live()
+            self.progress = YOLOCustomProgress(
+                *self.configure_columns(trainer),
+                auto_refresh=False,
+                disable=self.is_disabled,
+                console=self._console,
+            )
+            self.progress.start()
+
+            self._progress_stopped = False
+
+            self.max_result = 0
+            self.past_results = deque(maxlen=5)
+            self.progress.table = Table()
+
+    @override
+    def _get_train_description(self, current_epoch: int) -> str:
+        return Text("[cyan]Train [white]|")
+
+    @override
+    @rank_zero_only
+    def on_train_start(self, trainer, pl_module):
+        self._init_progress(trainer)
+        num_epochs = trainer.max_epochs - 1
+        self.task_epoch = self._add_task(
+            total_batches=num_epochs,
+            description=f"[cyan]Start Training {num_epochs} epochs",
+        )
+        self.max_result = 0
+        self.past_results.clear()
+        self.progress.update(self.task_epoch, advance=-0.5)
+
+    @override
+    @rank_zero_only
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch: Any, batch_idx: int):
+        self._update(self.train_progress_bar_id, batch_idx + 1)
+        self._update_metrics(trainer, pl_module)
+        epoch_descript = "[cyan]Train [white]|"
+        batch_descript = "[green]Train [white]|"
+        metrics = self.get_metrics(trainer, pl_module)
+        metrics.pop("v_num")
+        for metrics_name, metrics_val in metrics.items():
+            if "Loss_step" in metrics_name:
+                epoch_descript += f"{metrics_name.removesuffix('_step').split('/')[1]: ^9}|"
+                batch_descript += f"   {metrics_val:2.2f}  |"
+
+        self.progress.update(self.task_epoch, advance=1 / self.total_train_batches, description=epoch_descript)
+        self.progress.update(self.train_progress_bar_id, description=batch_descript)
+        self.refresh()
 
-    def rank_check(logging_function):
-        def wrapper(self, *args, **kwargs):
-            if getattr(self, "local_rank", 0) != 0:
-                return
-            return logging_function(self, *args, **kwargs)
+    @override
+    @rank_zero_only
+    def on_train_end(self, trainer: "Trainer", pl_module: "LightningModule") -> None:
+        self._update_metrics(trainer, pl_module)
+        self.progress.remove_task(self.train_progress_bar_id)
+        self.train_progress_bar_id = None
+
+    @override
+    @rank_zero_only
+    def on_validation_end(self, trainer: "Trainer", pl_module: "LightningModule") -> None:
+        if trainer.state.fn == "fit":
+            self._update_metrics(trainer, pl_module)
+        self.reset_dataloader_idx_tracker()
+        all_metrics = self.get_metrics(trainer, pl_module)
+
+        ap_ar_list = [
+            key
+            for key in all_metrics.keys()
+            if key.startswith(("map", "mar")) and not key.endswith(("_step", "_epoch"))
+        ]
+        score = np.array([all_metrics[key] for key in ap_ar_list]) * 100
+
+        self.progress.table, ap_main = make_ap_table(score, self.past_results, self.max_result, trainer.current_epoch)
+        self.max_result = np.maximum(score, self.max_result)
+        self.past_results.append((trainer.current_epoch, ap_main))
+
+    @override
+    def refresh(self) -> None:
+        if self.progress:
+            self.progress.refresh()
+
+    @property
+    def validation_description(self) -> str:
+        return "[green]Validation"
+
+
+class YOLORichModelSummary(RichModelSummary):
+    @staticmethod
+    @override
+    def summarize(
+        summary_data: List[Tuple[str, List[str]]],
+        total_parameters: int,
+        trainable_parameters: int,
+        model_size: float,
+        total_training_modes: Dict[str, int],
+        **summarize_kwargs: Any,
+    ) -> None:
+        from lightning.pytorch.utilities.model_summary import get_human_readable_count
+
+        console = get_console()
+
+        header_style: str = summarize_kwargs.get("header_style", "bold magenta")
+        table = Table(header_style=header_style)
+        table.add_column(" ", style="dim")
+        table.add_column("Name", justify="left", no_wrap=True)
+        table.add_column("Type")
+        table.add_column("Params", justify="right")
+        table.add_column("Mode")
+
+        column_names = list(zip(*summary_data))[0]
+
+        for column_name in ["In sizes", "Out sizes"]:
+            if column_name in column_names:
+                table.add_column(column_name, justify="right", style="white")
+
+        rows = list(zip(*(arr[1] for arr in summary_data)))
+        for row in rows:
+            table.add_row(*row)
+
+        console.print(table)
+
+        parameters = []
+        for param in [trainable_parameters, total_parameters - trainable_parameters, total_parameters, model_size]:
+            parameters.append("{:<{}}".format(get_human_readable_count(int(param)), 10))
+
+        grid = Table(header_style=header_style)
+        table.add_column(" ", style="dim")
+        grid.add_column("[bold]Attributes[/]")
+        grid.add_column("Value")
+
+        grid.add_row("[bold]Trainable params[/]", f"{parameters[0]}")
+        grid.add_row("[bold]Non-trainable params[/]", f"{parameters[1]}")
+        grid.add_row("[bold]Total params[/]", f"{parameters[2]}")
+        grid.add_row("[bold]Total estimated model params size (MB)[/]", f"{parameters[3]}")
+        grid.add_row("[bold]Modules in train mode[/]", f"{total_training_modes['train']}")
+        grid.add_row("[bold]Modules in eval mode[/]", f"{total_training_modes['eval']}")
+
+        console.print(grid)
+
+
+class ImageLogger(Callback):
+    def on_validation_batch_end(self, trainer: Trainer, pl_module, outputs, batch, batch_idx) -> None:
+        if batch_idx != 0:
+            return
+        batch_size, images, targets, rev_tensor, img_paths = batch
+        gt_boxes = targets[0] if targets.ndim == 3 else targets
+        pred_boxes = outputs[0] if isinstance(outputs, list) else outputs
+        images = [images[0]]
+        step = trainer.current_epoch
+        for logger in trainer.loggers:
+            if isinstance(logger, WandbLogger):
+                logger.log_image("Input Image", images, step=step)
+                logger.log_image("Ground Truth", images, step=step, boxes=[log_bbox(gt_boxes)])
+                logger.log_image("Prediction", images, step=step, boxes=[log_bbox(pred_boxes)])
+
+
+def setup_logger(logger_name):
+    class EmojiFormatter(logging.Formatter):
+        def format(self, record, emoji=":high_voltage:"):
+            return f"{emoji} {super().format(record)}"
+
+    rich_handler = RichHandler(markup=True)
+    rich_handler.setFormatter(EmojiFormatter("%(message)s"))
+    rich_logger = logging.getLogger(logger_name)
+    if rich_logger:
+        rich_logger.handlers.clear()
+        rich_logger.addHandler(rich_handler)
+
+
+def setup(cfg: Config):
+    # seed_everything(cfg.lucky_number)
+    if hasattr(cfg, "quite"):
+        logger.removeHandler("YOLO_logger")
+        return
 
-        return wrapper
+    setup_logger("lightning.fabric")
+    setup_logger("lightning.pytorch")
 
-    def get_renderable(self):
-        renderable = Group(*self.get_renderables(), self.ap_table)
-        return renderable
+    def custom_wandb_log(string="", level=int, newline=True, repeat=True, prefix=True, silent=False):
+        if silent:
+            return
+        for line in string.split("\n"):
+            logger.info(Text.from_ansi(":globe_with_meridians: " + line))
 
-    @rank_check
-    def start_train(self, num_epochs: int):
-        self.task_epoch = self.add_task(f"[cyan]Start Training {num_epochs} epochs", total=num_epochs)
-        self.update(self.task_epoch, advance=-0.5)
-
-    @rank_check
-    def start_one_epoch(
-        self, num_batches: int, task: str = "Train", optimizer: Optimizer = None, epoch_idx: int = None
-    ):
-        self.num_batches = num_batches
-        self.task = task
-        if hasattr(self, "task_epoch"):
-            self.update(self.task_epoch, description=f"[cyan] Preparing Data")
-
-        if optimizer is not None:
-            lr_values = [params["lr"] for params in optimizer.param_groups]
-            lr_names = ["Learning Rate/bias", "Learning Rate/norm", "Learning Rate/conv"]
-            if self.use_wandb:
-                for lr_name, lr_value in zip(lr_names, lr_values):
-                    self.wandb.log({lr_name: lr_value}, step=epoch_idx)
-
-            if self.use_tensorboard:
-                for lr_name, lr_value in zip(lr_names, lr_values):
-                    self.tb_writer.add_scalar(lr_name, lr_value, global_step=epoch_idx)
-
-        self.batch_task = self.add_task(f"[green] Phase: {task}", total=num_batches)
-
-    @rank_check
-    def one_batch(self, batch_info: Dict[str, Tensor] = None):
-        epoch_descript = "[cyan]" + self.task + "[white] |"
-        batch_descript = "|"
-        if self.task == "Train":
-            self.update(self.task_epoch, advance=1 / self.num_batches)
-        for info_name, info_val in batch_info.items():
-            epoch_descript += f"{info_name: ^9}|"
-            batch_descript += f"   {info_val:2.2f}  |"
-        self.update(self.batch_task, advance=1, description=f"[green]{self.task} [white]{batch_descript}")
-        if hasattr(self, "task_epoch"):
-            self.update(self.task_epoch, description=epoch_descript)
-
-    @rank_check
-    def finish_one_epoch(self, batch_info: Dict[str, Any] = None, epoch_idx: int = -1):
-        if self.task == "Train":
-            prefix = "Loss"
-        elif self.task == "Validate":
-            prefix = "Metrics"
-        batch_info = {f"{prefix}/{key}": value for key, value in batch_info.items()}
-        if self.use_wandb:
-            self.wandb.log(batch_info, step=epoch_idx)
-        if self.use_tensorboard:
-            for key, value in batch_info.items():
-                self.tb_writer.add_scalar(key, value, epoch_idx)
-
-        self.remove_task(self.batch_task)
-
-    @rank_check
-    def visualize_image(
-        self,
-        images: Optional[Tensor] = None,
-        ground_truth: Optional[Tensor] = None,
-        prediction: Optional[Union[List[Tensor], Tensor]] = None,
-        epoch_idx: int = 0,
-    ) -> None:
-        """
-        Upload the ground truth bounding boxes, predicted bounding boxes, and the original image to wandb or TensorBoard.
-
-        Args:
-            images (Optional[Tensor]): Tensor of images with shape (BZ, 3, 640, 640).
-            ground_truth (Optional[Tensor]): Ground truth bounding boxes with shape (BZ, N, 5) or (N, 5). Defaults to None.
-            prediction (prediction: Optional[Union[List[Tensor], Tensor]]): List of predicted bounding boxes with shape (N, 6) or (N, 6). Defaults to None.
-            epoch_idx (int): Current epoch index. Defaults to 0.
-        """
-        if images is not None:
-            images = images[0] if images.ndim == 4 else images
-            if self.use_wandb:
-                wandb.log({"Input Image": wandb.Image(images)}, step=epoch_idx)
-            if self.use_tensorboard:
-                self.tb_writer.add_image("Media/Input Image", images, 1)
-
-        if ground_truth is not None:
-            gt_boxes = ground_truth[0] if ground_truth.ndim == 3 else ground_truth
-            if self.use_wandb:
-                wandb.log(
-                    {"Ground Truth": wandb.Image(images, boxes={"predictions": {"box_data": log_bbox(gt_boxes)}})},
-                    step=epoch_idx,
-                )
-            if self.use_tensorboard:
-                self.tb_writer.add_image("Media/Ground Truth", pil_to_tensor(draw_bboxes(images, gt_boxes)), epoch_idx)
-
-        if prediction is not None:
-            pred_boxes = prediction[0] if isinstance(prediction, list) else prediction
-            if self.use_wandb:
-                wandb.log(
-                    {"Prediction": wandb.Image(images, boxes={"predictions": {"box_data": log_bbox(pred_boxes)}})},
-                    step=epoch_idx,
-                )
-            if self.use_tensorboard:
-                self.tb_writer.add_image("Media/Prediction", pil_to_tensor(draw_bboxes(images, pred_boxes)), epoch_idx)
-
-    @rank_check
-    def start_pycocotools(self):
-        self.batch_task = self.add_task("[green]Run pycocotools", total=1)
-
-    @rank_check
-    def finish_pycocotools(self, result, epoch_idx=-1):
-        ap_table, ap_main = make_ap_table(result * 100, self.ap_past_list, self.last_result, epoch_idx)
-        self.last_result = np.maximum(result, self.last_result)
-        self.ap_past_list.append((epoch_idx, ap_main))
-        self.ap_table = ap_table
-
-        if self.use_wandb:
-            self.wandb.log({"PyCOCO/AP @ .5:.95": ap_main[2], "PyCOCO/AP @ .5": ap_main[5]})
-        if self.use_tensorboard:
-            # TODO: waiting torch bugs fix, https://github.com/pytorch/pytorch/issues/32651
-            self.tb_writer.add_scalar("PyCOCO/AP @ .5:.95", ap_main[2], epoch_idx)
-            self.tb_writer.add_scalar("PyCOCO/AP @ .5", ap_main[5], epoch_idx)
-
-        self.update(self.batch_task, advance=1)
-        self.refresh()
-        self.remove_task(self.batch_task)
+    wandb.errors.term._log = custom_wandb_log
 
-    @rank_check
-    def finish_train(self):
-        self.remove_task(self.task_epoch)
-        self.stop()
-        if self.use_wandb:
-            self.wandb.finish()
-        if self.use_tensorboard:
-            self.tb_writer.close()
+    save_path = validate_log_directory(cfg, cfg.name)
 
+    progress, loggers = [], []
+    progress.append(YOLORichProgressBar())
+    progress.append(YOLORichModelSummary())
+    progress.append(ImageLogger())
+    if cfg.use_tensorboard:
+        loggers.append(TensorBoardLogger(log_graph="all", save_dir=save_path))
+    if cfg.use_wandb:
+        loggers.append(WandbLogger(project="YOLO", name=cfg.name, save_dir=save_path, id=None))
 
-def custom_wandb_log(string="", level=int, newline=True, repeat=True, prefix=True, silent=False):
-    if silent:
-        return
-    for line in string.split("\n"):
-        logger.opt(raw=not newline, colors=True).info("🌐 " + line)
+    return progress, loggers
 
 
 def log_model_structure(model: Union[ModuleList, YOLOLayer, YOLO]):
@@ -279,6 +294,7 @@ def log_model_structure(model: Union[ModuleList, YOLOLayer, YOLO]):
     console.print(table)
 
 
+@rank_zero_only
 def validate_log_directory(cfg: Config, exp_name: str) -> Path:
     base_path = Path(cfg.out_path, cfg.task.task)
     save_path = base_path / exp_name
@@ -296,8 +312,8 @@ def validate_log_directory(cfg: Config, exp_name: str) -> Path:
             )
 
     save_path.mkdir(parents=True, exist_ok=True)
-    logger.opt(colors=True).info(f"📄 Created log folder: <u><fg #808080>{save_path}</></>")
-    logger.add(save_path / "output.log", mode="w", backtrace=True, diagnose=True)
+    logger.info(f"📄 Created log folder: [blue b u]{save_path}[/]")
+    logger.addHandler(FileHandler(save_path / "output.log"))
     return save_path
 
 
@@ -332,4 +348,4 @@ def log_bbox(
             bbox_entry["scores"] = {"confidence": conf[0]}
         bbox_list.append(bbox_entry)
 
-    return bbox_list
+    return {"predictions": {"box_data": bbox_list}}

yolo/utils/model_utils.py

@@ -4,7 +4,6 @@ from typing import List, Optional, Type, Union
 
 import torch
 import torch.distributed as dist
-from loguru import logger
 from omegaconf import ListConfig
 from torch import Tensor
 from torch.optim import Optimizer
@@ -13,6 +12,7 @@ from torch.optim.lr_scheduler import LambdaLR, SequentialLR, _LRScheduler
 from yolo.config.config import IDX_TO_ID, NMSConfig, OptimizerConfig, SchedulerConfig
 from yolo.model.yolo import YOLO
 from yolo.utils.bounding_box_utils import bbox_nms, transform_bbox
+from yolo.utils.logger import logger
 
 
 class ExponentialMovingAverage:
@@ -52,9 +52,9 @@ def create_optimizer(model: YOLO, optim_cfg: OptimizerConfig) -> Optimizer:
     conv_params = [p for name, p in model.named_parameters() if "weight" in name and "bn" not in name]
 
     model_parameters = [
-        {"params": bias_params, "weight_decay": 0},
-        {"params": conv_params},
-        {"params": norm_params, "weight_decay": 0},
+        {"params": bias_params, "momentum": 0.8, "weight_decay": 0},
+        {"params": conv_params, "momentum": 0.8},
+        {"params": norm_params, "momentum": 0.8, "weight_decay": 0},
     ]
 
     def next_epoch(self, batch_num):
@@ -65,12 +65,16 @@ def create_optimizer(model: YOLO, optim_cfg: OptimizerConfig) -> Optimizer:
 
     def next_batch(self):
         self.batch_idx += 1
+        lr_dict = dict()
         for lr_idx, param_group in enumerate(self.param_groups):
             min_lr, max_lr = self.min_lr[lr_idx], self.max_lr[lr_idx]
             param_group["lr"] = min_lr + (self.batch_idx) * (max_lr - min_lr) / self.batch_num
+            lr_dict[f"LR/{lr_idx}"] = param_group["lr"]
+        return lr_dict
 
     optimizer_class.next_batch = next_batch
     optimizer_class.next_epoch = next_epoch
+
     optimizer = optimizer_class(model_parameters, **optim_cfg.args)
     optimizer.max_lr = [0.1, 0, 0]
     return optimizer
@@ -168,6 +172,7 @@ def predicts_to_json(img_paths, predicts, rev_tensor):
     batch_json = []
     for img_path, bboxes, box_reverse in zip(img_paths, predicts, rev_tensor):
         scale, shift = box_reverse.split([1, 4])
+        bboxes = bboxes.clone()
         bboxes[:, 1:5] = (bboxes[:, 1:5] - shift[None]) / scale[None]
         bboxes[:, 1:5] = transform_bbox(bboxes[:, 1:5], "xyxy -> xywh")
         for cls, *pos, conf in bboxes:

yolo/utils/solver_utils.py

@@ -1,5 +1,6 @@
 import contextlib
 import io
+from typing import Dict
 
 import numpy as np
 from pycocotools.coco import COCO
@@ -17,7 +18,7 @@ def calculate_ap(coco_gt: COCO, pd_path):
     return coco_eval.stats
 
 
-def make_ap_table(score, past_result=[], last_score=None, epoch=-1):
+def make_ap_table(score: Dict[str, float], past_result=[], max_result=None, epoch=-1):
     ap_table = Table()
     ap_table.add_column("Epoch", justify="center", style="white", width=5)
     ap_table.add_column("Avg. Precision", justify="left", style="cyan")
@@ -30,7 +31,7 @@ def make_ap_table(score, past_result=[], last_score=None, epoch=-1):
     if past_result:
         ap_table.add_row()
 
-    color = np.where(last_score <= score, "[green]", "[red]")
+    color = np.where(max_result <= score, "[green]", "[red]")
 
     this_ap = ("AP @ .5:.95", color[0], score[0], "AP @        .5", color[1], score[1])
     metrics = [