🔀 [Merge] branch 'SETUP' into DATASET

config/config.py

@@ -0,0 +1,13 @@
+from dataclasses import dataclass
+from typing import List, Dict, Union
+
+
+@dataclass
+class Model:
+    anchor: List[List[int]]
+    model: Dict[str, List[Dict[str, Union[Dict, List, int]]]]
+
+
+@dataclass
+class Config:
+    model: Model

config/config.yaml

@@ -0,0 +1,8 @@
+hydra:
+  run:
+    dir: ./runs
+
+defaults:
+ - data_config: coco.yaml
+ - model: v7-base.yaml
+ - _self_

model/module.py

@@ -297,7 +297,7 @@ class CSPELAN(nn.Module):
 
 
 class Concat(nn.Module):
-    def __init__(self, dim=-1):
+    def __init__(self, dim=1):
         super(Concat, self).__init__()
         self.dim = dim
 

model/yolo.py

@@ -1,10 +1,13 @@
+import inspect
+from typing import Any, Dict, List, Union
+
+import torch
 import torch.nn as nn
 from loguru import logger
-from typing import Dict, Any, List
-import inspect
-from utils.tools import load_model_cfg
+from omegaconf import OmegaConf
 
 from model import module
+from utils.tools import load_model_cfg
 
 
 def get_layer_map():
@@ -23,7 +26,6 @@ def get_layer_map():
 class YOLO(nn.Module):
     """
     A preliminary YOLO (You Only Look Once) model class still under development.
-    #TODO: Next: Finish forward proccess
 
     Parameters:
         model_cfg: Configuration for the YOLO model. Expected to define the layers,
@@ -33,19 +35,16 @@ class YOLO(nn.Module):
     def __init__(self, model_cfg: Dict[str, Any]):
         super(YOLO, self).__init__()
         self.nc = model_cfg["nc"]
-        self.layer_map = get_layer_map()  # Dynamically get the mapping
-        self.build_model(model_cfg["model"])
-        print(self.model)
-        # raise NotImplementedError("Constructor not implemented.")
+        self.layer_map = get_layer_map()  # Get the map Dict[str: Module]
+        self.build_model(model_cfg.model)
 
     def build_model(self, model_arch: Dict[str, List[Dict[str, Dict[str, Dict]]]]):
         model_list = nn.ModuleList()
         output_dim = [3]
         layer_indices_by_tag = {}
-
-        for arch_name, arch in model_arch.items():
-            logger.info(f"Building model-{arch_name}")
-            for layer_idx, layer_spec in enumerate(arch, start=1):
+        for arch_name in model_arch:
+            logger.info(f"🏗️  Building model-{arch_name}")
+            for layer_idx, layer_spec in enumerate(model_arch[arch_name], start=1):
                 layer_type, layer_info = next(iter(layer_spec.items()))
                 layer_args = layer_info.get("args", {})
                 source = layer_info.get("source", -1)
@@ -56,8 +55,9 @@ class YOLO(nn.Module):
                     layer_args["in_channels"] = output_dim[source]
                 if "Detect" in layer_type:
                     layer_args["nc"] = self.nc
+                    layer_args["ch"] = [output_dim[idx] for idx in source]
 
-                layer = self.create_layer(layer_type, **layer_args)
+                layer = self.create_layer(layer_type, source, **layer_args)
                 model_list.append(layer)
 
                 if "tags" in layer_info:
@@ -69,22 +69,32 @@ class YOLO(nn.Module):
                 output_dim.append(out_channels)
         self.model = model_list
 
-    def get_out_channels(self, layer_type, layer_args, output_dim, source):
+    def forward(self, x):
+        y = [x]
+        for layer in self.model:
+            if OmegaConf.is_list(layer.source):
+                model_input = [y[idx] for idx in layer.source]
+            else:
+                model_input = y[layer.source]
+            x = layer(model_input)
+            y.append(x)
+        return x
+
+    def get_out_channels(self, layer_type: str, layer_args: dict, output_dim: list, source: Union[int, list]):
         if "Conv" in layer_type:
             return layer_args["out_channels"]
+        if layer_type in ["MaxPool", "UpSample"]:
+            return output_dim[source]
         if layer_type == "Concat":
             return sum(output_dim[idx] for idx in source)
-        if "Pool" in layer_type:
-            return output_dim[source] // 2
-        if layer_type == "UpSample":
-            return output_dim[source] * 2
         if layer_type == "IDetect":
             return None
 
-    def create_layer(self, layer_type: str, **kwargs):
-        # Dictionary mapping layer names to actual layer classes
+    def create_layer(self, layer_type: str, source: Union[int, list], **kwargs):
         if layer_type in self.layer_map:
-            return self.layer_map[layer_type](**kwargs)
+            layer = self.layer_map[layer_type](**kwargs)
+            layer.source = source
+            return layer
         else:
             raise ValueError(f"Unsupported layer type: {layer_type}")
 
@@ -99,6 +109,7 @@ def get_model(model_cfg: dict) -> YOLO:
         YOLO: An instance of the model defined by the given configuration.
     """
     model = YOLO(model_cfg)
+    logger.info("✅ Success load model")
     return model
 
 

train.py

@@ -2,25 +2,18 @@ import argparse
 from loguru import logger
 from model.yolo import get_model
 from utils.tools import load_model_cfg, custom_logger
+import hydra
+from config.config import Config
+from omegaconf import OmegaConf
 
 
-def parse_arguments() -> argparse.Namespace:
-    """
-    Parse command-line arguments to get the model configuration file.
-
-    Returns:
-        argparse.Namespace: The command-line arguments object with 'config' attribute.
-    """
-    parser = argparse.ArgumentParser(description="Load a YOLO model configuration and display the model.")
-    parser.add_argument(
-        "--model-config", type=str, default="v7-base", help="Name or path to the model configuration file."
-    )
-    return parser.parse_args()
+@hydra.main(config_path="config", config_name="config", version_base=None)
+def main(cfg: Config):
+    OmegaConf.set_struct(cfg, False)
+    model = get_model(cfg.model)
+    logger.info("Success load model")
 
 
 if __name__ == "__main__":
     custom_logger()
-    args = parse_arguments()
-    model_cfg = load_model_cfg(args.model_config)
-    model = get_model(model_cfg)
-    logger.info("Success load model: {}", model)
+    main()