🧱 [Update] config files struct, make it clearly

yolo/config/README.md

@@ -1 +0,0 @@
-model configuration

yolo/config/config.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Dict, List, Union
+from typing import Any, Dict, List, Optional, Union
 
 from torch import nn
 
@@ -10,25 +10,47 @@ class AnchorConfig:
     strides: List[int]
 
 
+@dataclass
+class LayerConfg:
+    args: Dict
+    source: Union[int, str, List[int]]
+    tags: str
+
+
+@dataclass
+class BlockConfig:
+    block: List[Dict[str, LayerConfg]]
+
+
 @dataclass
 class Model:
     anchor: AnchorConfig
-    model: Dict[str, List[Dict[str, Union[Dict, List, int]]]]
+    model: Dict[str, BlockConfig]
 
 
 @dataclass
-class Download:
-    auto: bool
+class DownloadDetail:
+    url: str
+    file_size: int
+
+
+@dataclass
+class DownloadOptions:
+    details: Dict[str, DownloadDetail]
+
+
+@dataclass
+class DatasetConfig:
     path: str
+    auto_download: Optional[DownloadOptions]
 
 
 @dataclass
-class DataLoaderConfig:
-    batch_size: int
-    class_num: int
-    image_size: List[int]
+class DataConfig:
     shuffle: bool
+    batch_size: int
     pin_memory: bool
+    data_augment: Dict[str, int]
 
 
 @dataclass
@@ -44,16 +66,24 @@ class OptimizerConfig:
 
 
 @dataclass
-class SchedulerArgs:
-    step_size: int
-    gamma: float
+class MatcherConfig:
+    iou: str
+    topk: int
+    factor: Dict[str, int]
+
+
+@dataclass
+class LossConfig:
+    objective: Dict[str, int]
+    aux: Union[bool, float]
+    matcher: MatcherConfig
 
 
 @dataclass
 class SchedulerConfig:
     type: str
-    args: SchedulerArgs
-    warmup: Dict[str, Union[str, int, float]]
+    warmup: Dict[str, Union[int, float]]
+    args: Dict[str, Any]
 
 
 @dataclass
@@ -62,66 +92,48 @@ class EMAConfig:
     decay: float
 
 
-@dataclass
-class MatcherConfig:
-    iou: str
-    topk: int
-    factor: Dict[str, int]
-
-
-@dataclass
-class LossConfig:
-    objective: List[List]
-    aux: Union[bool, float]
-    matcher: MatcherConfig
-
-
 @dataclass
 class TrainConfig:
+    task: str
+    dataset: DatasetConfig
     epoch: int
+    data: DataConfig
     optimizer: OptimizerConfig
+    loss: LossConfig
     scheduler: SchedulerConfig
     ema: EMAConfig
-    loss: LossConfig
 
 
 @dataclass
-class GeneralConfig:
-    out_path: str
-    task: str
-    device: Union[str, int, List[int]]
-    cpu_num: int
-    use_wandb: bool
-    lucky_number: 10
-    exist_ok: bool
-    resume_train: bool
-    use_TensorBoard: bool
+class NMSConfig:
+    min_confidence: int
+    min_iou: int
 
 
 @dataclass
-class HyperConfig:
-    general: GeneralConfig
-    data: DataLoaderConfig
-    train: TrainConfig
+class InferenceConfig:
+    task: str
+    nms: NMSConfig
 
 
 @dataclass
-class Dataset:
-    file_name: str
-    num_files: int
+class Config:
+    task: Union[TrainConfig, InferenceConfig]
+    model: Model
+    name: str
 
+    device: Union[str, int, List[int]]
+    cpu_num: int
 
-@dataclass
-class Datasets:
-    base_url: str
-    images: Dict[str, Dataset]
+    class_num: int
+    image_size: List[int]
 
+    out_path: str
+    exist_ok: bool
 
-@dataclass
-class Download:
-    auto: bool
-    save_path: str
-    datasets: Datasets
+    lucky_number: 10
+    use_wandb: bool
+    use_TensorBoard: bool
 
 
 @dataclass
@@ -134,11 +146,3 @@ class YOLOLayer(nn.Module):
 
     def __post_init__(self):
         super().__init__()
-
-
-@dataclass
-class Config:
-    model: Model
-    download: Download
-    hyper: HyperConfig
-    name: str

yolo/config/config.yaml

@@ -1,13 +1,12 @@
 hydra:
   run:
-    dir: ./runs
+    dir: runs
 
-defaults:
- - data: coco
- - download: ../data/download
- - augmentation: ../data/augmentation
- - model: v9-c
- - hyper: default
- - _self_
- 
 name: v9-dev
+
+defaults:
+  - _self_
+  - task: train
+  - model: v9-c
+  - general
+ 

yolo/config/data/augmentation.yaml

@@ -1,3 +0,0 @@
-Mosaic: 1
-# MixUp: 1
-HorizontalFlip: 0.5

yolo/config/data/coco.yaml

@@ -1 +0,0 @@
-path: data/coco

yolo/config/general.yaml

@@ -0,0 +1,12 @@
+deivce: [0]
+cpu_num: 16
+
+class_num: 80
+image_size: [640, 640]
+
+out_path: runs
+exist_ok: True
+
+lucky_number: 10
+use_wandb: False
+use_TensorBoard: False

yolo/config/hyper/default.yaml

@@ -1,48 +0,0 @@
-general:
-  out_path: runs
-  task: train
-  deivce: [0]
-  cpu_num: 16
-  use_wandb: False
-  lucky_number: 10
-  exist_ok: True
-  resume_train: False
-  use_TensorBoard: False
-data:
-  batch_size: 16 
-  class_num: 80
-  image_size: [640, 640]
-  shuffle: True
-  pin_memory: True
-train:
-  epoch: 500
-  optimizer:
-    type: SGD
-    args:
-      lr: 0.01
-      weight_decay: 0.0005
-      momentum: 0.937
-  loss:
-    objective:
-      BCELoss: 0.5
-      BoxLoss: 7.5  
-      DFLoss: 1.5
-    aux:
-      0.25
-    matcher:
-      iou: CIoU
-      topk: 10
-      factor:
-        iou: 6.0
-        cls: 0.5
-  scheduler:
-    type: LinearLR
-    warmup:
-      epochs: 3.0
-    args:
-      total_iters: ${hyper.train.epoch}
-      start_factor: 1
-      end_factor: 0.01
-  ema:
-      enabled: true
-      decay: 0.995

yolo/config/{data/download.yaml → task/dataset/coco.yaml}

@@ -1,6 +1,6 @@
-auto: True
-save_path: data/coco
-datasets:
+path: data/coco
+
+auto_download:
   images:
     base_url: http://images.cocodataset.org/zips/
     train2017:
@@ -15,7 +15,4 @@ datasets:
   annotations:
     base_url: http://images.cocodataset.org/annotations/
     annotations:
-      file_name: annotations_trainval2017
-hydra:
-  run:
-    dir: ./runs
+      file_name: annotations_trainval2017

yolo/config/task/train.yaml

@@ -0,0 +1,48 @@
+task: train
+defaults:
+  - dataset: coco
+
+epoch: 500
+
+data:
+  batch_size: 16 
+  shuffle: True
+  pin_memory: True
+  data_augment:
+    Mosaic: 1
+    # MixUp: 1
+    HorizontalFlip: 0.5
+
+optimizer:
+  type: SGD
+  args:
+    lr: 0.01
+    weight_decay: 0.0005
+    momentum: 0.937
+
+loss:
+  objective:
+    BCELoss: 0.5
+    BoxLoss: 7.5  
+    DFLoss: 1.5
+  aux:
+    0.25
+  matcher:
+    iou: CIoU
+    topk: 10
+    factor:
+      iou: 6.0
+      cls: 0.5
+
+scheduler:
+  type: LinearLR
+  warmup:
+    epochs: 3.0
+  args:
+    total_iters: ${task.epoch}
+    start_factor: 1
+    end_factor: 0.01
+
+ema:
+  enabled: true
+  decay: 0.995

yolo/model/yolo.py

@@ -123,7 +123,7 @@ def get_model(cfg: Config) -> YOLO:
         YOLO: An instance of the model defined by the given configuration.
     """
     OmegaConf.set_struct(cfg.model, False)
-    model = YOLO(cfg.model, cfg.hyper.data.class_num)
+    model = YOLO(cfg.model, cfg.class_num)
     logger.info("✅ Success load model")
     log_model_structure(model.model)
     draw_model(model=model)

yolo/tools/data_loader.py

@@ -12,7 +12,7 @@ from torch.utils.data import DataLoader, Dataset
 from torchvision.transforms import functional as TF
 from tqdm.rich import tqdm
 
-from yolo.config.config import Config
+from yolo.config.config import Config, TrainConfig
 from yolo.tools.data_augmentation import (
     AugmentationComposer,
     HorizontalFlip,
@@ -20,6 +20,7 @@ from yolo.tools.data_augmentation import (
     Mosaic,
     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,
@@ -29,16 +30,16 @@ from yolo.utils.dataset_utils import (
 
 
 class YoloDataset(Dataset):
-    def __init__(self, config: dict, phase: str = "train2017", image_size: int = 640):
-        dataset_cfg = config.data
-        augment_cfg = config.augmentation
-        phase_name = dataset_cfg.get(phase, phase)
+    def __init__(self, config: TrainConfig, phase: str = "train2017", image_size: int = 640):
+        augment_cfg = config.data.data_augment
+        # TODO: add yaml -> train: train2017
+        phase_name = config.dataset.auto_download.get(phase, phase)
         self.image_size = image_size
 
         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(dataset_cfg.path, phase_name)
+        self.data = self.load_data(config.dataset.path, phase_name)
 
     def load_data(self, dataset_path, phase_name):
         """
@@ -159,15 +160,15 @@ class YoloDataset(Dataset):
 class YoloDataLoader(DataLoader):
     def __init__(self, config: Config):
         """Initializes the YoloDataLoader with hydra-config files."""
-        hyper = config.hyper.data
-        dataset = YoloDataset(config)
+        data_cfg = config.task.data
+        dataset = YoloDataset(config.task)
 
         super().__init__(
             dataset,
-            batch_size=hyper.batch_size,
-            shuffle=hyper.shuffle,
-            num_workers=config.hyper.general.cpu_num,
-            pin_memory=hyper.pin_memory,
+            batch_size=data_cfg.batch_size,
+            shuffle=data_cfg.shuffle,
+            num_workers=config.cpu_num,
+            pin_memory=data_cfg.pin_memory,
             collate_fn=self.collate_fn,
         )
 
@@ -197,7 +198,10 @@ class YoloDataLoader(DataLoader):
         return batch_images, batch_targets
 
 
-def create_dataloader(config):
+def create_dataloader(config: Config):
+    if config.task.dataset.auto_download:
+        prepare_dataset(config.task.dataset)
+
     return YoloDataLoader(config)
 
 
@@ -211,7 +215,7 @@ if __name__ == "__main__":
     import sys
 
     sys.path.append("./")
-    from tools.logging_utils import custom_logger
+    from utils.logging_utils import custom_logger
 
     custom_logger()
     main()

yolo/tools/dataset_preparation.py

@@ -6,6 +6,8 @@ from hydra import main
 from loguru import logger
 from tqdm import tqdm
 
+from yolo.config.config import DatasetConfig
+
 
 def download_file(url, destination):
     """
@@ -45,12 +47,12 @@ def check_files(directory, expected_count=None):
 
 
 @main(config_path="../config/data", config_name="download", version_base=None)
-def prepare_dataset(cfg):
+def prepare_dataset(cfg: DatasetConfig):
     """
     Prepares dataset by downloading and unzipping if necessary.
     """
-    data_dir = cfg.save_path
-    for data_type, settings in cfg.datasets.items():
+    data_dir = cfg.path
+    for data_type, settings in cfg.auto_download.items():
         base_url = settings["base_url"]
         for dataset_type, dataset_args in settings.items():
             if dataset_type == "base_url":

yolo/tools/format_converters.py

@@ -0,0 +1,29 @@
+def convert_weight(old_state_dict, new_state_dict, model_size: int = 38):
+    # TODO: need to refactor
+    for idx in range(model_size):
+        new_list, old_list = [], []
+        for weight_name, weight_value in new_state_dict.items():
+            if weight_name.split(".")[0] == str(idx):
+                new_list.append((weight_name, None))
+        for weight_name, weight_value in old_state_dict.items():
+            if f"model.{idx+1}." in weight_name:
+                old_list.append((weight_name, weight_value))
+        if len(new_list) == len(old_list):
+            for (weight_name, _), (_, weight_value) in zip(new_list, old_list):
+                new_state_dict[weight_name] = weight_value
+        else:
+            for weight_name, weight_value in old_list:
+                if "dfl" in weight_name:
+                    continue
+                _, _, conv_name, conv_idx, *details = weight_name.split(".")
+                if conv_name == "cv4" or conv_name == "cv5":
+                    conv_idx = str(int(conv_idx) + 3)
+
+                if conv_name == "cv2" or conv_name == "cv4":
+                    conv_task = "anchor_conv"
+                if conv_name == "cv3" or conv_name == "cv5":
+                    conv_task = "class_conv"
+
+                weight_name = ".".join(["37", "heads", conv_idx, conv_task, *details])
+                new_state_dict[weight_name] = weight_value
+    return new_state_dict

yolo/tools/loss_functions.py

@@ -75,8 +75,8 @@ class DFLoss(nn.Module):
 class YOLOLoss:
     def __init__(self, cfg: Config) -> None:
         self.reg_max = cfg.model.anchor.reg_max
-        self.class_num = cfg.hyper.data.class_num
-        self.image_size = list(cfg.hyper.data.image_size)
+        self.class_num = cfg.class_num
+        self.image_size = list(cfg.image_size)
         self.strides = cfg.model.anchor.strides
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -89,7 +89,7 @@ class YOLOLoss:
         self.dfl = DFLoss(self.anchors, self.scaler, self.reg_max)
         self.iou = BoxLoss()
 
-        self.matcher = BoxMatcher(cfg.hyper.train.loss.matcher, self.class_num, self.anchors)
+        self.matcher = BoxMatcher(cfg.task.loss.matcher, self.class_num, self.anchors)
         self.box_converter = AnchorBoxConverter(cfg, device)
 
     def separate_anchor(self, anchors):
@@ -127,11 +127,11 @@ class YOLOLoss:
 class DualLoss:
     def __init__(self, cfg: Config) -> None:
         self.loss = YOLOLoss(cfg)
-        self.aux_rate = cfg.hyper.train.loss.aux
+        self.aux_rate = cfg.task.loss.aux
 
-        self.iou_rate = cfg.hyper.train.loss.objective["BoxLoss"]
-        self.dfl_rate = cfg.hyper.train.loss.objective["DFLoss"]
-        self.cls_rate = cfg.hyper.train.loss.objective["BCELoss"]
+        self.iou_rate = cfg.task.loss.objective["BoxLoss"]
+        self.dfl_rate = cfg.task.loss.objective["DFLoss"]
+        self.cls_rate = cfg.task.loss.objective["BCELoss"]
 
     def __call__(self, predicts: List[Tensor], targets: Tensor) -> Tuple[Tensor, Dict[str, Tensor]]:
         targets[:, :, 1:] = targets[:, :, 1:] * self.loss.scale_up

yolo/tools/trainer.py

@@ -18,7 +18,7 @@ from yolo.utils.model_utils import (
 
 class ModelTrainer:
     def __init__(self, cfg: Config, save_path: str, device):
-        train_cfg: TrainConfig = cfg.hyper.train
+        train_cfg: TrainConfig = cfg.task
         model = get_model(cfg)
 
         self.model = model.to(device)

yolo/utils/bounding_box_utils.py

@@ -7,7 +7,7 @@ from einops import rearrange
 from torch import Tensor
 from torchvision.ops import batched_nms
 
-from yolo.config.config import Config, MatcherConfig
+from yolo.config.config import Config, MatcherConfig, NMSConfig
 
 
 def calculate_iou(bbox1, bbox2, metrics="iou") -> Tensor:
@@ -127,8 +127,8 @@ def generate_anchors(image_size: List[int], strides: List[int], device):
 class AnchorBoxConverter:
     def __init__(self, cfg: Config, device: torch.device) -> None:
         self.reg_max = cfg.model.anchor.reg_max
-        self.class_num = cfg.hyper.data.class_num
-        self.image_size = list(cfg.hyper.data.image_size)
+        self.class_num = cfg.class_num
+        self.image_size = list(cfg.image_size)
         self.strides = cfg.model.anchor.strides
 
         self.scale_up = torch.tensor(self.image_size * 2, device=device)
@@ -291,17 +291,17 @@ class BoxMatcher:
         return torch.cat([align_cls, align_bbox], dim=-1), valid_mask.bool()
 
 
-def bbox_nms(predicts: Tensor, min_conf: float = 0, min_iou: float = 0.5):
+def bbox_nms(predicts: Tensor, nms_cfg: NMSConfig):
     cls_dist, bbox = predicts.split([80, 4], dim=-1)
 
     # filter class by confidence
     cls_val, cls_idx = cls_dist.max(dim=-1, keepdim=True)
-    valid_mask = cls_val > min_conf
+    valid_mask = cls_val > nms_cfg.min_confidence
     valid_cls = cls_idx[valid_mask]
     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, min_iou)
+    nms_idx = batched_nms(valid_box, valid_cls, batch_idx, nms_cfg.min_iou)
     predicts_nms = []
     for idx in range(batch_idx.max() + 1):
         instance_idx = nms_idx[idx == batch_idx[nms_idx]]

yolo/utils/logging_utils.py

@@ -24,7 +24,7 @@ from rich.table import Table
 from torch import Tensor
 from torch.optim import Optimizer
 
-from yolo.config.config import Config, GeneralConfig, YOLOLayer
+from yolo.config.config import Config, YOLOLayer
 
 
 def custom_logger():
@@ -110,11 +110,11 @@ def log_model_structure(model: List[YOLOLayer]):
     console.print(table)
 
 
-def validate_log_directory(general_cfg: GeneralConfig, exp_name):
-    base_path = os.path.join(general_cfg.out_path, general_cfg.task)
+def validate_log_directory(cfg: Config, exp_name: str):
+    base_path = os.path.join(cfg.out_path, cfg.task.task)
     save_path = os.path.join(base_path, exp_name)
 
-    if not general_cfg.exist_ok:
+    if not cfg.exist_ok:
         index = 1
         old_exp_name = exp_name
         while os.path.isdir(save_path):