🔀 [Merge] branch 'MODELv2' into DEPLOY

yolo/config/model/v9-m.yaml

@@ -0,0 +1,133 @@
+anchor:
+  reg_max: 16
+
+model:
+  backbone:
+    - Conv:
+        args: {out_channels: 32, kernel_size: 3, stride: 2}
+        source: 0
+    - Conv:
+        args: {out_channels: 64, kernel_size: 3, stride: 2}
+    - RepNCSPELAN:
+        args: {out_channels: 128, part_channels: 128}
+
+    - AConv:
+        args: {out_channels: 240}
+    - RepNCSPELAN:
+        args: {out_channels: 240, part_channels: 240}
+        tags: B3
+
+    - AConv:
+        args: {out_channels: 360}
+    - RepNCSPELAN:
+        args: {out_channels: 360, part_channels: 360}
+        tags: B4
+
+    - AConv:
+        args: {out_channels: 480}
+    - RepNCSPELAN:
+        args: {out_channels: 480, part_channels: 480}
+        tags: B5
+
+  neck:
+    - SPPELAN:
+        args: {out_channels: 480}
+        tags: N3
+
+    - UpSample:
+        args: {scale_factor: 2, mode: nearest}
+    - Concat:
+        source: [-1, B4]
+    - RepNCSPELAN:
+        args: {out_channels: 360, part_channels: 360}
+        tags: N4
+
+    - UpSample:
+        args: {scale_factor: 2, mode: nearest}
+    - Concat:
+        source: [-1, B3]
+
+  head:
+    - RepNCSPELAN:
+        args: {out_channels: 240, part_channels: 240}
+        tags: P3
+
+    - AConv:
+        args: {out_channels: 184}
+    - Concat:
+        source: [-1, N4]
+    - RepNCSPELAN:
+        args: {out_channels: 360, part_channels: 360}
+        tags: P4
+
+    - AConv:
+        args: {out_channels: 240}
+    - Concat:
+        source: [-1, N3]
+    - RepNCSPELAN:
+        args: {out_channels: 480, part_channels: 480}
+        tags: P5
+
+  detection:
+    - MultiheadDetection:
+        source: [P3, P4, P5]
+        tags: Main
+        args:
+            reg_max: ${model.anchor.reg_max}
+        output: True
+
+  auxiliary:
+    - CBLinear:
+        source: B3
+        args: {out_channels: [240]}
+        tags: R3
+    - CBLinear:
+        source: B4
+        args: {out_channels: [240, 360]}
+        tags: R4
+    - CBLinear:
+        source: B5
+        args: {out_channels: [240, 360, 480]}
+        tags: R5
+
+    - Conv:
+        args: {out_channels: 32, kernel_size: 3, stride: 2}
+        source: 0
+    - Conv:
+        args: {out_channels: 64, kernel_size: 3, stride: 2}
+    - RepNCSPELAN:
+        args: {out_channels: 128, part_channels: 128}
+
+    - AConv:
+        args: {out_channels: 240}
+    - CBFuse:
+        source: [R3, R4, R5, -1]
+        args: {index: [0, 0, 0]}
+    - RepNCSPELAN:
+        args: {out_channels: 240, part_channels: 240}
+        tags: A3
+
+    - AConv:
+        args: {out_channels: 360}
+    - CBFuse:
+        source: [R4, R5, -1]
+        args: {index: [1, 1]}
+    - RepNCSPELAN:
+        args: {out_channels: 360, part_channels: 360}
+        tags: A4
+
+    - AConv:
+        args: {out_channels: 480}
+    - CBFuse:
+        source: [R5, -1]
+        args: {index: [2]}
+    - RepNCSPELAN:
+        args: {out_channels: 480, part_channels: 480}
+        tags: A5
+
+    - MultiheadDetection:
+        source: [A3, A4, A5]
+        tags: AUX
+        args:
+            reg_max: ${model.anchor.reg_max}
+        output: True

yolo/config/model/v9-s.yaml

@@ -0,0 +1,134 @@
+anchor:
+  reg_max: 16
+
+model:
+  backbone:
+    - Conv:
+        args: {out_channels: 32, kernel_size: 3, stride: 2}
+        source: 0
+    - Conv:
+        args: {out_channels: 64, kernel_size: 3, stride: 2}
+    - ELAN:
+        args: {out_channels: 64, part_channels: 64}
+
+    - AConv:
+        args: {out_channels: 128}
+    - RepNCSPELAN:
+        args:
+            out_channels: 128
+            part_channels: 128
+            csp_args: {repeat_num: 3}
+        tags: B3 # 18
+
+    - AConv:
+        args: {out_channels: 192}
+    - RepNCSPELAN:
+        args:
+            out_channels: 192
+            part_channels: 192
+            csp_args: {repeat_num: 3}
+        tags: B4
+
+    - AConv:
+        args: {out_channels: 256}
+    - RepNCSPELAN:
+        args:
+            out_channels: 256
+            part_channels: 256
+            csp_args: {repeat_num: 3}
+        tags: B5
+
+  neck:
+    - SPPELAN:
+        args: {out_channels: 256}
+        tags: N3
+
+    - UpSample:
+        args: {scale_factor: 2, mode: nearest}
+    - Concat:
+        source: [-1, B4]
+    - RepNCSPELAN:
+        args:
+            out_channels: 192
+            part_channels: 192
+            csp_args: {repeat_num: 3}
+        tags: N4
+
+    - UpSample:
+        args: {scale_factor: 2, mode: nearest}
+    - Concat:
+        source: [-1, B3]
+
+    - RepNCSPELAN:
+        args:
+            out_channels: 128
+            part_channels: 128
+            csp_args: {repeat_num: 3}
+        tags: P3
+    - AConv:
+        args: {out_channels: 96}
+    - Concat:
+        source: [-1, N4]
+
+    - RepNCSPELAN:
+        args:
+            out_channels: 192
+            part_channels: 192
+            csp_args: {repeat_num: 3}
+        tags: P4
+    - AConv:
+        args: {out_channels: 128}
+    - Concat:
+        source: [-1, N3]
+
+    - RepNCSPELAN:
+        args:
+            out_channels: 256
+            part_channels: 256
+            csp_args: {repeat_num: 3}
+        tags: P5
+
+  detection:
+    - MultiheadDetection:
+        source: [P3, P4, P5]
+        tags: Main
+        args:
+            reg_max: ${model.anchor.reg_max}
+        output: True
+
+  head:
+    - SPPELAN:
+        source: B5
+        args: {out_channels: 256}
+        tags: A5
+
+    - UpSample:
+        args: {scale_factor: 2, mode: nearest}
+    - Concat:
+        source: [-1, B4]
+
+    - RepNCSPELAN:
+        args:
+            out_channels: 192
+            part_channels: 192
+            csp_args: {repeat_num: 3}
+        tags: A4
+
+    - UpSample:
+        args: {scale_factor: 2, mode: nearest}
+    - Concat:
+        source: [-1, B3]
+
+    - RepNCSPELAN:
+        args:
+            out_channels: 128
+            part_channels: 128
+            csp_args: {repeat_num: 3}
+        tags: A3
+
+    - MultiheadDetection:
+        source: [A3, A4, A5]
+        tags: AUX
+        args:
+            reg_max: ${model.anchor.reg_max}
+        output: True

yolo/lazy.py

@@ -13,17 +13,16 @@ from yolo.tools.data_loader import create_dataloader
 from yolo.tools.solver import ModelTester, ModelTrainer
 from yolo.utils.bounding_box_utils import Vec2Box
 from yolo.utils.deploy_utils import FastModelLoader
-from yolo.utils.logging_utils import custom_logger, validate_log_directory
+from yolo.utils.logging_utils import ProgressLogger
 
 
 @hydra.main(config_path="config", config_name="config", version_base=None)
 def main(cfg: Config):
-    custom_logger()
-    save_path = validate_log_directory(cfg, exp_name=cfg.name)
+    progress = ProgressLogger(cfg, exp_name=cfg.name)
     dataloader = create_dataloader(cfg.task.data, cfg.dataset, cfg.task.task)
     device = torch.device(cfg.device)
     if getattr(cfg.task, "fast_inference", False):
-        model = FastModelLoader(cfg).load_model()
+        model = FastModelLoader(cfg, device).load_model()
         device = torch.device(cfg.device)
     else:
         model = create_model(cfg.model, class_num=cfg.class_num, weight_path=cfg.weight, device=device)
@@ -31,11 +30,11 @@ def main(cfg: Config):
     vec2box = Vec2Box(model, cfg.image_size, device)
 
     if cfg.task.task == "train":
-        trainer = ModelTrainer(cfg, model, vec2box, save_path, device)
+        trainer = ModelTrainer(cfg, model, vec2box, progress, device)
         trainer.solve(dataloader)
 
     if cfg.task.task == "inference":
-        tester = ModelTester(cfg, model, vec2box, save_path, device)
+        tester = ModelTester(cfg, model, vec2box, progress, device)
         tester.solve(dataloader)
 
 

yolo/model/module.py

@@ -192,6 +192,36 @@ class RepNCSP(nn.Module):
         return self.conv3(torch.cat((x1, x2), dim=1))
 
 
+class ELAN(nn.Module):
+    """ELAN  structure."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        part_channels: int,
+        *,
+        process_channels: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__()
+
+        if process_channels is None:
+            process_channels = part_channels // 2
+
+        self.conv1 = Conv(in_channels, part_channels, 1, **kwargs)
+        self.conv2 = Conv(part_channels // 2, process_channels, 3, padding=1, **kwargs)
+        self.conv3 = Conv(process_channels, process_channels, 3, padding=1, **kwargs)
+        self.conv4 = Conv(part_channels + 2 * process_channels, out_channels, 1, **kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x1, x2 = self.conv1(x).chunk(2, 1)
+        x3 = self.conv2(x2)
+        x4 = self.conv3(x3)
+        x5 = self.conv4(torch.cat([x1, x2, x3, x4], dim=1))
+        return x5
+
+
 class RepNCSPELAN(nn.Module):
     """RepNCSPELAN block combining RepNCSP blocks with ELAN structure."""
 
@@ -230,6 +260,21 @@ class RepNCSPELAN(nn.Module):
         return x5
 
 
+class AConv(nn.Module):
+    """Downsampling module combining average and max pooling with convolution for feature reduction."""
+
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        mid_layer = {"kernel_size": 3, "stride": 2}
+        self.avg_pool = Pool("avg", kernel_size=2, stride=1)
+        self.conv = Conv(in_channels, out_channels, **mid_layer)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.avg_pool(x)
+        x = self.conv(x)
+        return x
+
+
 class ADown(nn.Module):
     """Downsampling module combining average and max pooling with convolution for feature reduction."""
 
@@ -498,26 +543,6 @@ class CSPDark(nn.Module):
         return self.cv2(torch.cat((self.cb(y[0]), y[1]), 1))
 
 
-# ELAN
-class ELAN(nn.Module):
-    # ELAN
-    def __init__(self, in_channels, out_channels, med_channels, elan_repeat=2, cb_repeat=2, ratio=1.0):
-
-        super().__init__()
-
-        h_channels = med_channels // 2
-        self.cv1 = Conv(in_channels, med_channels, 1, 1)
-        self.cb = nn.ModuleList(ConvBlock(h_channels, repeat=cb_repeat, ratio=ratio) for _ in range(elan_repeat))
-        self.cv2 = Conv((2 + elan_repeat) * h_channels, out_channels, 1, 1)
-
-    def forward(self, x):
-
-        y = list(self.cv1(x).chunk(2, 1))
-        y.extend((m(y[-1])) for m in self.cb)
-
-        return self.cv2(torch.cat(y, 1))
-
-
 class CSPELAN(nn.Module):
     # ELAN
     def __init__(self, in_channels, out_channels, med_channels, elan_repeat=2, cb_repeat=2, ratio=1.0):

yolo/model/yolo.py

@@ -1,5 +1,5 @@
 import os
-from typing import Any, Dict, List, Union
+from typing import Dict, List, Optional, Union
 
 import torch
 from loguru import logger
@@ -43,7 +43,7 @@ class YOLO(nn.Module):
                 source = self.get_source_idx(layer_info.get("source", -1), layer_idx)
 
                 # Find in channels
-                if any(module in layer_type for module in ["Conv", "ELAN", "ADown", "CBLinear"]):
+                if any(module in layer_type for module in ["Conv", "ELAN", "ADown", "AConv", "CBLinear"]):
                     layer_args["in_channels"] = output_dim[source]
                 if "Detection" in layer_type:
                     layer_args["in_channels"] = [output_dim[idx] for idx in source]
@@ -81,7 +81,7 @@ class YOLO(nn.Module):
         return output
 
     def get_out_channels(self, layer_type: str, layer_args: dict, output_dim: list, source: Union[int, list]):
-        if any(module in layer_type for module in ["Conv", "ELAN", "ADown"]):
+        if any(module in layer_type for module in ["Conv", "ELAN", "ADown", "AConv"]):
             return layer_args["out_channels"]
         if layer_type == "CBFuse":
             return output_dim[source[-1]]
@@ -117,9 +117,7 @@ class YOLO(nn.Module):
             raise ValueError(f"Unsupported layer type: {layer_type}")
 
 
-def create_model(
-    model_cfg: ModelConfig, class_num: int = 80, weight_path: str = "weights/v9-c.pt", device: device = device("cuda")
-) -> YOLO:
+def create_model(model_cfg: ModelConfig, weight_path: Optional[str], device: device, class_num: int = 80) -> YOLO:
     """Constructs and returns a model from a Dictionary configuration file.
 
     Args:
@@ -135,8 +133,9 @@ def create_model(
         if not os.path.exists(weight_path):
             logger.info(f"🌐 Weight {weight_path} not found, try downloading")
             prepare_weight(weight_path=weight_path)
-        model.model.load_state_dict(torch.load(weight_path, map_location=device))
-        logger.info("✅ Success load model weight")
+        if os.path.exists(weight_path):
+            model.model.load_state_dict(torch.load(weight_path, map_location=device), strict=False)
+            logger.info("✅ Success load model weight")
 
     log_model_structure(model.model)
     draw_model(model=model)

yolo/tools/data_loader.py

@@ -23,7 +23,6 @@ from yolo.tools.data_augmentation import (
     VerticalFlip,
 )
 from yolo.tools.dataset_preparation import prepare_dataset
-from yolo.tools.drawer import draw_bboxes
 from yolo.utils.dataset_utils import (
     create_image_metadata,
     locate_label_paths,
@@ -204,7 +203,7 @@ def create_dataloader(data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: st
         return StreamDataLoader(data_cfg)
 
     if dataset_cfg.auto_download:
-        prepare_dataset(dataset_cfg)
+        prepare_dataset(dataset_cfg, task)
 
     return YoloDataLoader(data_cfg, dataset_cfg, task)
 

yolo/tools/dataset_preparation.py

@@ -52,7 +52,7 @@ def check_files(directory, expected_count=None):
     return len(files) == expected_count if expected_count is not None else bool(files)
 
 
-def prepare_dataset(dataset_cfg: DatasetConfig):
+def prepare_dataset(dataset_cfg: DatasetConfig, task: str):
     """
     Prepares dataset by downloading and unzipping if necessary.
     """
@@ -60,8 +60,8 @@ def prepare_dataset(dataset_cfg: DatasetConfig):
     for data_type, settings in dataset_cfg.auto_download.items():
         base_url = settings["base_url"]
         for dataset_type, dataset_args in settings.items():
-            if dataset_type == "base_url":
-                continue  # Skip the base_url entry
+            if dataset_type != "annotations" and dataset_cfg.get(task, task) != dataset_type:
+                continue
             file_name = f"{dataset_args.get('file_name', dataset_type)}.zip"
             url = f"{base_url}{file_name}"
             local_zip_path = os.path.join(data_dir, file_name)

yolo/tools/solver.py

@@ -14,7 +14,7 @@ from yolo.tools.data_loader import StreamDataLoader, 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, bbox_nms, calculate_map
-from yolo.utils.logging_utils import ProgressTracker
+from yolo.utils.logging_utils import ProgressLogger
 from yolo.utils.model_utils import (
     ExponentialMovingAverage,
     create_optimizer,
@@ -23,7 +23,7 @@ from yolo.utils.model_utils import (
 
 
 class ModelTrainer:
-    def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, save_path: str, device):
+    def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, progress: ProgressLogger, device):
         train_cfg: TrainConfig = cfg.task
         self.model = model
         self.vec2box = vec2box
@@ -31,11 +31,11 @@ class ModelTrainer:
         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 = ProgressTracker(cfg.name, save_path, cfg.use_wandb)
+        self.progress = progress
         self.num_epochs = cfg.task.epoch
 
         self.validation_dataloader = create_dataloader(cfg.task.validation.data, cfg.dataset, cfg.task.validation.task)
-        self.validator = ModelValidator(cfg.task.validation, model, vec2box, save_path, device, self.progress)
+        self.validator = ModelValidator(cfg.task.validation, model, vec2box, progress, device, self.progress)
 
         if getattr(train_cfg.ema, "enabled", False):
             self.ema = ExponentialMovingAverage(model, decay=train_cfg.ema.decay)
@@ -102,14 +102,15 @@ class ModelTrainer:
 
 
 class ModelTester:
-    def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, save_path: str, device):
+    def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, progress: ProgressLogger, device):
         self.model = model
         self.device = device
         self.vec2box = vec2box
-        self.progress = ProgressTracker(cfg, save_path, cfg.use_wandb)
+        self.progress = progress
 
         self.nms = cfg.task.nms
-        self.save_path = save_path
+        self.save_path = os.path.join(progress.save_path, "images")
+        os.makedirs(self.save_path, exist_ok=True)
         self.save_predict = getattr(cfg.task, "save_predict", None)
         self.idx2label = cfg.class_list
 
@@ -142,8 +143,7 @@ class ModelTester:
                         break
                     if not self.save_predict:
                         continue
-
-                if self.save_predict == False:
+                if self.save_predict != False:
                     save_image_path = os.path.join(self.save_path, f"frame{idx:03d}.png")
                     img.save(save_image_path)
                     logger.info(f"💾 Saved visualize image at {save_image_path}")
@@ -164,16 +164,13 @@ class ModelValidator:
         validation_cfg: ValidationConfig,
         model: YOLO,
         vec2box: Vec2Box,
-        save_path: str,
         device,
-        # TODO: think Progress?
-        progress: ProgressTracker,
+        progress: ProgressLogger,
     ):
         self.model = model
         self.vec2box = vec2box
         self.device = device
         self.progress = progress
-        self.save_path = save_path
 
         self.nms = validation_cfg.nms
 

yolo/utils/logging_utils.py

@@ -38,15 +38,18 @@ def custom_logger(quite: bool = False):
     )
 
 
-class ProgressTracker:
-    def __init__(self, exp_name: str, save_path: str, use_wandb: bool = False):
+class ProgressLogger:
+    def __init__(self, cfg: Config, exp_name: str):
+        custom_logger(getattr(cfg, "quite", False))
+        self.save_path = validate_log_directory(cfg, exp_name=cfg.name)
+
         self.progress = Progress(
             TextColumn("[progress.description]{task.description}"),
             BarColumn(bar_width=None),
             TextColumn("{task.completed:.0f}/{task.total:.0f}"),
             TimeRemainingColumn(),
         )
-        self.use_wandb = use_wandb
+        self.use_wandb = cfg.use_wandb
         if self.use_wandb:
             wandb.errors.term._log = custom_wandb_log
             self.wandb = wandb.init(