🔧 [Add] TypeHint for yolo.model

yolo/config/config.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass
 from typing import Dict, List, Union
 
+from torch import nn
+
 
 @dataclass
 class AnchorConfig:
@@ -100,6 +102,17 @@ class Download:
     datasets: Datasets
 
 
+@dataclass
+class YOLOLayer(nn.Module):
+    source: Union[int, str, List[int]]
+    output: bool
+    tags: str
+    layer_type: str
+
+    def __post_init__(self):
+        super().__init__()
+
+
 @dataclass
 class Config:
     model: Model

yolo/model/yolo.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 from loguru import logger
 from omegaconf import ListConfig, OmegaConf
 
-from yolo.config.config import Config, Model
+from yolo.config.config import Config, Model, YOLOLayer
 from yolo.tools.layer_helper import get_layer_map
 
 
@@ -21,6 +21,7 @@ class YOLO(nn.Module):
         super(YOLO, self).__init__()
         self.num_classes = num_classes
         self.layer_map = get_layer_map()  # Get the map Dict[str: Module]
+        self.model: List[YOLOLayer] = nn.ModuleList()
         self.build_model(model_cfg.model)
 
     def build_model(self, model_arch: Dict[str, List[Dict[str, Dict[str, Dict]]]]):