✨ [New] load v9 model via yolo.py

yolo/model/yolo.py

@@ -2,8 +2,9 @@ from typing import Any, Dict, List, Union
 
 import torch.nn as nn
 from loguru import logger
-from omegaconf import OmegaConf
+from omegaconf import ListConfig, OmegaConf
 
+from yolo.config.config import Model
 from yolo.tools.layer_helper import get_layer_map
 
 
@@ -18,7 +19,7 @@ class YOLO(nn.Module):
 
     def __init__(self, model_cfg: Model):
         super(YOLO, self).__init__()
-        self.nc = model_cfg["nc"]
+        self.num_classes = model_cfg["num_classes"]
         self.layer_map = get_layer_map()  # Get the map Dict[str: Module]
         self.build_model(model_cfg.model)
 
@@ -26,24 +27,30 @@ class YOLO(nn.Module):
         model_list = nn.ModuleList()
         output_dim = [3]
         layer_indices_by_tag = {}
+        layer_idx = 1
         logger.info(f"🚜 Building YOLO")
         for arch_name in model_arch:
             logger.info(f"  🏗️  Building {arch_name}")
-            for layer_idx, layer_spec in enumerate(model_arch[arch_name], start=1):
+            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", {})
-                source = layer_info.get("source", -1)
-                output = layer_info.get("output", False)
 
+                # Get input source
+                source = layer_info.get("source", -1)
                 if isinstance(source, str):
                     source = layer_indices_by_tag[source]
-                if "Conv" in layer_type:
+                elif isinstance(source, ListConfig):
+                    source = [layer_indices_by_tag[idx] if isinstance(idx, str) else idx for idx in source]
+
+                # Find in channels
+                if any(module in layer_type for module in ["Conv", "ELAN", "ADown", "CBLinear"]):
                     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]
+                if "Detection" in layer_type:
+                    layer_args["in_channels"] = [output_dim[idx] for idx in source]
+                    layer_args["num_classes"] = self.num_classes
 
-                layer = self.create_layer(layer_type, source, output, **layer_args)
+                # create layers
+                layer = self.create_layer(layer_type, source, layer_info, **layer_args)
                 model_list.append(layer)
 
                 if "tags" in layer_info:
@@ -53,13 +60,15 @@ class YOLO(nn.Module):
 
                 out_channels = self.get_out_channels(layer_type, layer_args, output_dim, source)
                 output_dim.append(out_channels)
+            layer_idx += 1
+
         self.model = model_list
 
     def forward(self, x):
         y = [x]
         output = []
         for layer in self.model:
-            if OmegaConf.is_list(layer.source):
+            if isinstance(layer.source, list):
                 model_input = [y[idx] for idx in layer.source]
             else:
                 model_input = y[layer.source]
@@ -70,8 +79,10 @@ 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 "Conv" in layer_type:
+        if any(module in layer_type for module in ["Conv", "ELAN", "ADown"]):
             return layer_args["out_channels"]
+        if layer_type == "CBFuse":
+            return output_dim[source[-1]]
         if layer_type in ["Pool", "UpSample"]:
             return output_dim[source]
         if layer_type == "Concat":
@@ -79,11 +90,12 @@ class YOLO(nn.Module):
         if layer_type == "IDetect":
             return None
 
-    def create_layer(self, layer_type: str, source: Union[int, list], output=False, **kwargs):
+    def create_layer(self, layer_type: str, source: Union[int, list], layer_info, **kwargs):
         if layer_type in self.layer_map:
             layer = self.layer_map[layer_type](**kwargs)
-            layer.source = source
-            layer.output = output
+            setattr(layer, "source", source)
+            setattr(layer, "output", layer_info.get("output", False))
+            setattr(layer, "tags", layer_info.get("tags", None))
             return layer
         else:
             raise ValueError(f"Unsupported layer type: {layer_type}")