🔧 [Update] all structure of yaml, more readability

yolo/config/config.py

@@ -23,8 +23,9 @@ class BlockConfig:
 
 
 @dataclass
-class Model:
+class ModelConfig:
     anchor: AnchorConfig
+    class_num: int
     model: Dict[str, BlockConfig]
 
 
@@ -50,7 +51,10 @@ class DataConfig:
     shuffle: bool
     batch_size: int
     pin_memory: bool
+    cpu_num: int
+    image_size: List[int]
     data_augment: Dict[str, int]
+    source: Optional[Union[str, int]]
 
 
 @dataclass
@@ -92,18 +96,6 @@ class EMAConfig:
     decay: float
 
 
-@dataclass
-class TrainConfig:
-    task: str
-    dataset: DatasetConfig
-    epoch: int
-    data: DataConfig
-    optimizer: OptimizerConfig
-    loss: LossConfig
-    scheduler: SchedulerConfig
-    ema: EMAConfig
-
-
 @dataclass
 class NMSConfig:
     min_confidence: int
@@ -113,15 +105,35 @@ class NMSConfig:
 @dataclass
 class InferenceConfig:
     task: str
-    source: Union[str, int]
     nms: NMSConfig
+    data: DataConfig
     fast_inference: Optional[None]
 
 
+@dataclass
+class ValidationConfig:
+    task: str
+    nms: NMSConfig
+    data: DataConfig
+
+
+@dataclass
+class TrainConfig:
+    task: str
+    epoch: int
+    data: DataConfig
+    optimizer: OptimizerConfig
+    loss: LossConfig
+    scheduler: SchedulerConfig
+    ema: EMAConfig
+    validation: ValidationConfig
+
+
 @dataclass
 class Config:
-    task: Union[TrainConfig, InferenceConfig]
-    model: Model
+    task: Union[TrainConfig, InferenceConfig, ValidationConfig]
+    dataset: DatasetConfig
+    model: ModelConfig
     name: str
 
     device: Union[str, int, List[int]]

yolo/config/config.yaml

@@ -7,6 +7,7 @@ name: v9-dev
 defaults:
   - _self_
   - task: train
+  - dataset: coco
   - model: v9-c
   - general
  

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

@@ -1,6 +1,6 @@
 path: data/coco
 train: train2017
-
+validation: val2017
 
 auto_download:
   images:

yolo/config/dataset/dev.yaml

@@ -0,0 +1,5 @@
+path: data/dev
+train: train
+validation: test
+
+auto_download: 

yolo/config/model/v9-c.yaml

@@ -2,6 +2,8 @@ anchor:
   reg_max: 16
   strides: [8, 16, 32]
 
+class_num: ${class_num}
+
 model:
   backbone:
     - Conv:

yolo/config/task/inference.yaml

@@ -1,10 +1,9 @@
 task: inference
-source: demo/images/inference/image.png
+
 fast_inference:  # onnx, trt or Empty
 data:
-  batch_size: 16
-  shuffle: False
-  pin_memory: True
+  source: demo/images/inference/image.png
+  image_size: ${image_size}
   data_augment: {}
 nms:
   min_confidence: 0.5

yolo/config/task/train.yaml

@@ -1,11 +1,14 @@
 task: train
+
 defaults:
-  - dataset: coco
+  - validation: ../validation
 
 epoch: 500
 
 data:
   batch_size: 16 
+  image_size: ${image_size}
+  cpu_num: ${cpu_num}
   shuffle: True
   pin_memory: True
   data_augment:

yolo/config/task/validation.yaml

@@ -0,0 +1,12 @@
+task: validation
+
+data:
+  batch_size: 16
+  image_size: ${image_size}
+  cpu_num: ${cpu_num}
+  shuffle: False
+  pin_memory: True
+  data_augment: {}
+nms:
+  min_confidence: 0.001
+  min_iou: 0.7

yolo/lazy.py

@@ -18,10 +18,10 @@ from yolo.utils.logging_utils import custom_logger, validate_log_directory
 @hydra.main(config_path="config", config_name="config", version_base=None)
 def main(cfg: Config):
     custom_logger()
-    save_path = validate_log_directory(cfg, cfg.name)
-    dataloader = create_dataloader(cfg)
+    save_path = validate_log_directory(cfg, exp_name=cfg.name)
+    dataloader = create_dataloader(cfg.task.data, cfg.dataset, cfg.task.task)
     device = torch.device(cfg.device)
-    if cfg.task.fast_inference:
+    if getattr(cfg.task, "fast_inference", False):
         model = FastModelLoader(cfg).load_model()
         device = torch.device(cfg.device)
     else:

yolo/model/yolo.py

@@ -6,7 +6,7 @@ import torch.nn as nn
 from loguru import logger
 from omegaconf import ListConfig, OmegaConf
 
-from yolo.config.config import Config, Model, YOLOLayer
+from yolo.config.config import Config, ModelConfig, YOLOLayer
 from yolo.tools.dataset_preparation import prepare_weight
 from yolo.tools.drawer import draw_model
 from yolo.utils.logging_utils import log_model_structure
@@ -22,9 +22,9 @@ class YOLO(nn.Module):
                    parameters, and any other relevant configuration details.
     """
 
-    def __init__(self, model_cfg: Model, num_classes: int):
+    def __init__(self, model_cfg: ModelConfig):
         super(YOLO, self).__init__()
-        self.num_classes = num_classes
+        self.num_classes = model_cfg.class_num
         self.layer_map = get_layer_map()  # Get the map Dict[str: Module]
         self.model: List[YOLOLayer] = nn.ModuleList()
         self.build_model(model_cfg.model)
@@ -126,7 +126,7 @@ def create_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.class_num)
+    model = YOLO(cfg.model)
     logger.info("✅ Success load model")
     if cfg.weight:
         if os.path.exists(cfg.weight):

yolo/tools/data_loader.py

@@ -2,7 +2,7 @@ import os
 from os import path
 from queue import Empty, Queue
 from threading import Event, Thread
-from typing import Generator, List, Optional, Tuple, Union
+from typing import Generator, List, Tuple, Union
 
 import cv2
 import hydra
@@ -15,7 +15,7 @@ from torch import Tensor
 from torch.utils.data import DataLoader, Dataset
 from torchvision.transforms import functional as TF
 
-from yolo.config.config import Config, TrainConfig
+from yolo.config.config import DataConfig, DatasetConfig
 from yolo.tools.data_augmentation import (
     AugmentationComposer,
     HorizontalFlip,
@@ -33,15 +33,15 @@ from yolo.utils.dataset_utils import (
 
 
 class YoloDataset(Dataset):
-    def __init__(self, config: TrainConfig, phase: str = "train2017", image_size: int = 640):
-        augment_cfg = config.data.data_augment
-        phase_name = config.dataset.get(phase, phase)
-        self.image_size = image_size
+    def __init__(self, data_cfg: DataConfig, dataset_cfg: DatasetConfig, phase: str = "train2017"):
+        augment_cfg = data_cfg.data_augment
+        self.image_size = data_cfg.image_size
+        phase_name = dataset_cfg.get(phase, phase)
 
         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(config.dataset.path, phase_name)
+        self.data = self.load_data(dataset_cfg.path, phase_name)
 
     def load_data(self, dataset_path, phase_name):
         """
@@ -161,16 +161,15 @@ class YoloDataset(Dataset):
 
 
 class YoloDataLoader(DataLoader):
-    def __init__(self, config: Config):
+    def __init__(self, data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: str = "train"):
         """Initializes the YoloDataLoader with hydra-config files."""
-        data_cfg = config.task.data
-        dataset = YoloDataset(config.task, config.task.task)
-
+        dataset = YoloDataset(data_cfg, dataset_cfg, task)
+        self.image_size = data_cfg.image_size[0]
         super().__init__(
             dataset,
             batch_size=data_cfg.batch_size,
             shuffle=data_cfg.shuffle,
-            num_workers=config.cpu_num,
+            num_workers=data_cfg.cpu_num,
             pin_memory=data_cfg.pin_memory,
             collate_fn=self.collate_fn,
         )
@@ -193,31 +192,33 @@ class YoloDataLoader(DataLoader):
         target_sizes = [item[1].size(0) for item in batch]
         # TODO: Improve readability of these proccess
         batch_targets = torch.zeros(batch_size, max(target_sizes), 5)
+        batch_targets[:, :, 0] = -1
         for idx, target_size in enumerate(target_sizes):
             batch_targets[idx, :target_size] = batch[idx][1]
+        batch_targets[:, :, 1:] *= self.image_size
 
         batch_images = torch.stack([item[0] for item in batch])
 
         return batch_images, batch_targets
 
 
-def create_dataloader(config: Config):
-    if config.task.task == "inference":
-        return StreamDataLoader(config)
+def create_dataloader(data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: str = "train"):
+    if task == "inference":
+        return StreamDataLoader(data_cfg)
 
-    if config.task.dataset.auto_download:
-        prepare_dataset(config.task.dataset)
+    if dataset_cfg.auto_download:
+        prepare_dataset(dataset_cfg)
 
-    return YoloDataLoader(config)
+    return YoloDataLoader(data_cfg, dataset_cfg, task)
 
 
 class StreamDataLoader:
-    def __init__(self, config: Config):
-        self.source = config.task.source
+    def __init__(self, data_cfg: DataConfig):
+        self.source = data_cfg.source
         self.running = True
         self.is_stream = isinstance(self.source, int) or self.source.lower().startswith("rtmp://")
 
-        self.transform = AugmentationComposer([], config.image_size[0])
+        self.transform = AugmentationComposer([], data_cfg.image_size)
         self.stop_event = Event()
 
         if self.is_stream:
@@ -301,19 +302,3 @@ class StreamDataLoader:
 
     def __len__(self):
         return self.queue.qsize() if not self.is_stream else 0
-
-
-@hydra.main(config_path="../config", config_name="config", version_base=None)
-def main(cfg):
-    dataloader = create_dataloader(cfg)
-    draw_bboxes(*next(iter(dataloader)))
-
-
-if __name__ == "__main__":
-    import sys
-
-    sys.path.append("./")
-    from utils.logging_utils import custom_logger
-
-    custom_logger()
-    main()

yolo/tools/dataset_preparation.py

@@ -52,12 +52,12 @@ def check_files(directory, expected_count=None):
     return len(files) == expected_count if expected_count is not None else bool(files)
 
 
-def prepare_dataset(cfg: DatasetConfig):
+def prepare_dataset(dataset_cfg: DatasetConfig):
     """
     Prepares dataset by downloading and unzipping if necessary.
     """
-    data_dir = cfg.path
-    for data_type, settings in cfg.auto_download.items():
+    data_dir = dataset_cfg.path
+    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":

yolo/tools/loss_functions.py

@@ -75,14 +75,11 @@ class DFLoss(nn.Module):
 class YOLOLoss:
     def __init__(self, cfg: Config) -> None:
         self.reg_max = cfg.model.anchor.reg_max
-        self.class_num = cfg.class_num
+        self.class_num = cfg.model.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")
 
-        self.reverse_reg = torch.arange(self.reg_max, dtype=torch.float32, device=device)
-        self.scale_up = torch.tensor(self.image_size * 2, device=device)
-
         self.anchors, self.scaler = generate_anchors(self.image_size, self.strides, device)
 
         self.cls = BCELoss()
@@ -90,7 +87,7 @@ class YOLOLoss:
         self.iou = BoxLoss()
 
         self.matcher = BoxMatcher(cfg.task.loss.matcher, self.class_num, self.anchors)
-        self.box_converter = AnchorBoxConverter(cfg, device)
+        self.box_converter = AnchorBoxConverter(cfg.model, self.image_size, device)
 
     def separate_anchor(self, anchors):
         """
@@ -134,7 +131,6 @@ class DualLoss:
         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
 
         # TODO: Need Refactor this region, make it flexible!
         predicts = divide_into_chunks(predicts[0], 2)

yolo/tools/solver.py

@@ -5,12 +5,12 @@ from torch import Tensor
 # TODO: We may can't use CUDA?
 from torch.cuda.amp import GradScaler, autocast
 
-from yolo.config.config import Config, TrainConfig
+from yolo.config.config import Config, TrainConfig, ValidationConfig
 from yolo.model.yolo import YOLO
-from yolo.tools.data_loader import StreamDataLoader
+from yolo.tools.data_loader import StreamDataLoader, create_dataloader
 from yolo.tools.drawer import draw_bboxes
 from yolo.tools.loss_functions import get_loss_function
-from yolo.utils.bounding_box_utils import AnchorBoxConverter, bbox_nms
+from yolo.utils.bounding_box_utils import AnchorBoxConverter, bbox_nms, calculate_map
 from yolo.utils.logging_utils import ProgressTracker
 from yolo.utils.model_utils import (
     ExponentialMovingAverage,
@@ -27,7 +27,7 @@ class ModelTrainer:
         self.optimizer = create_optimizer(model, train_cfg.optimizer)
         self.scheduler = create_scheduler(self.optimizer, train_cfg.scheduler)
         self.loss_fn = get_loss_function(cfg)
-        self.progress = ProgressTracker(cfg, save_path, cfg.use_wandb)
+        self.progress = ProgressTracker(cfg.name, save_path, cfg.use_wandb)
         self.num_epochs = cfg.task.epoch
 
         if getattr(train_cfg.ema, "enabled", False):

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, NMSConfig
+from yolo.config.config import MatcherConfig, ModelConfig, NMSConfig
 
 
 def calculate_iou(bbox1, bbox2, metrics="iou") -> Tensor:
@@ -125,14 +125,12 @@ 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.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)
-        self.anchors, self.scaler = generate_anchors(self.image_size, self.strides, device)
+    def __init__(self, model_cfg: ModelConfig, image_size: List[int], device: torch.device) -> None:
+        self.reg_max = model_cfg.anchor.reg_max
+        self.class_num = model_cfg.class_num
+        self.strides = model_cfg.anchor.strides
+
+        self.anchors, self.scaler = generate_anchors(image_size, self.strides, device)
         self.reverse_reg = torch.arange(self.reg_max, dtype=torch.float32, device=device)
 
     def __call__(self, predicts: List[Tensor], with_logits=False) -> Tensor:
@@ -255,7 +253,7 @@ class BoxMatcher:
         """
         predict_cls, predict_bbox = predict.split(self.class_num, dim=-1)  # B, HW x (C B) -> B x HW x C, B x HW x B
         target_cls, target_bbox = target.split([1, 4], dim=-1)  # B x N x (C B) -> B x N x C, B x N x B
-        target_cls = target_cls.long()
+        target_cls = target_cls.long().clamp(0)
 
         # get valid matrix (each gt appear in which anchor grid)
         grid_mask = self.get_valid_matrix(target_bbox)

yolo/utils/logging_utils.py

@@ -39,7 +39,7 @@ def custom_logger(quite: bool = False):
 
 
 class ProgressTracker:
-    def __init__(self, cfg: Config, save_path: str, use_wandb: bool = False):
+    def __init__(self, exp_name: str, save_path: str, use_wandb: bool = False):
         self.progress = Progress(
             TextColumn("[progress.description]{task.description}"),
             BarColumn(bar_width=None),

yolo/utils/model_utils.py

@@ -63,7 +63,7 @@ def create_scheduler(optimizer: Optimizer, schedule_cfg: SchedulerConfig) -> _LR
     if hasattr(schedule_cfg, "warmup"):
         wepoch = schedule_cfg.warmup.epochs
         lambda1 = lambda epoch: 0.1 + 0.9 * (epoch + 1 / wepoch) if epoch < wepoch else 1
-        lambda2 = lambda epoch: 10 - 9 * (epoch + 1 / wepoch) if epoch < wepoch else 1
+        lambda2 = lambda epoch: 10 - 9 * (epoch / wepoch) if epoch < wepoch else 1
         warmup_schedule = LambdaLR(optimizer, lr_lambda=[lambda1, lambda2, lambda1])
         schedule = SequentialLR(optimizer, schedulers=[warmup_schedule, schedule], milestones=[2])
     return schedule