🚚 [Rename] All utils, tools, higher readability

docs/HOWTO.md

@@ -88,3 +88,61 @@ Custom transformations should be designed to accept an image and its bounding bo
     # ... (Other Transform)
     CustomTransform: 0.5
     ```
+
+
+- **Utils**
+    - **bbox_utils**
+        - `class` Anchor2Box: transform predicted anchor to bounding box
+        - `class` Matcher: given prediction and groudtruth, find the groundtruth for each prediction
+        - `func` calculate_iou: calculate iou for given two list of bbox
+        - `func` transform_bbox: transform bbox from {xywh, xyxy, xcycwh} to {xywh, xyxy, xcycwh}
+        - `func` generate_anchors: given image size, make the anchor point for the given size
+    - **dataset_utils**
+        - `func` locate_label_paths:
+        - `func` create_image_metadata:
+        - `func` organize_annotations_by_image:
+        - `func` scale_segmentation:
+    - **logging_utils**
+        - `func` custom_log: custom loguru, overiding the origin logger
+        - `class` ProgressTracker: A class to handle output for each batch, epoch
+        - `func` log_model_structure: give a torch model, print it as a table
+        - `func` validate_log_directory: for given experiemnt, check if the log folder already existed
+    - **model_utils**
+        - `class` ExponentialMovingAverage: a mirror of model, do ema on model
+        - `func` create_optimizer: return a optimzer, for example SDG, ADAM
+        - `func` create_scheduler: return a scheduler, for example Step, Lambda
+    - **module_utils**
+        - `func` get_layer_map:
+        - `func` auto_pad: given a convolution block, return how many pixel should conv padding
+        - `func` create_activation_function: given a `func` name, return a activation `func`tion
+        - `func` round_up: given number and divider, return a number is mutliplcation of divider
+        - `func` divide_into_chunks: for a given list and n, seperate list to n sub list
+    - **trainer**
+        - `class` Trainer: a class can automatic train the model
+- **Tools**
+    - **converter_json2txt**
+        - `func` discretize_categories: given the dictionary class, turn id from 1: class
+        - `func` process_annotations: handle the whole dataset annotations
+        - `func` process_annotation: handle a annotation(a list of bounding box)
+        - `func` normalize_segmentation: normalize segmentation position to 0~1
+        - `func` convert_annotations: convert json annotations to txt file structure
+    - **data_augment**
+        - `class` AugmentationComposer: Compose a list of data augmentation strategy
+        - `class` VerticalFlip: a custom data augmentation, Random Vertical Flip
+        - `class` Mosaic: a data augmentation strategy, follow YOLOv5
+    - **dataloader**
+        - `class` YoloDataset: a custom dataset for training yolo's model
+        - `class` YoloDataLoader: a dataloader base on torch's dataloader, with custom allocate function
+        - `func` create_dataloader: given a config file, return a YOLO dataloader
+    - **drawer**
+        - `func` draw_bboxes: given a image and list of bbox, draw bbox on the image
+        - `func` draw_model: visualize the given model
+    - **get_dataset**
+        - `func` download_file: for a given link, downlaod the file
+        - `func` unzip_file: unzip the downlaoded zip to data/
+        - `func` check_files: check if the dataset file numbers is correct
+        - `func` prepare_dataset: automatic downlaod the dataset and check if it is correct
+    - **loss**
+        - `class` BoxLoss: a Custom Loss for bounding box
+        - `class` YOLOLoss: a implementation of yolov9 loss
+        - `class` DualLoss: a implementation of yolov9 loss with auxiliary detection head

examples/example_train.py

@@ -9,23 +9,23 @@ project_root = Path(__file__).resolve().parent.parent
 sys.path.append(str(project_root))
 
 from yolo.config.config import Config
-from yolo.tools.log_helper import custom_logger, get_valid_folder
-from yolo.tools.trainer import Trainer
-from yolo.utils.dataloader import get_dataloader
-from yolo.utils.get_dataset import prepare_dataset
+from yolo.tools.data_loader import create_dataloader
+from yolo.tools.dataset_preparation import prepare_dataset
+from yolo.tools.trainer import ModelTrainer
+from yolo.utils.logging_utils import custom_logger, validate_log_directory
 
 
 @hydra.main(config_path="../yolo/config", config_name="config", version_base=None)
 def main(cfg: Config):
     custom_logger()
-    save_path = get_valid_folder(cfg.hyper.general, cfg.name)
+    save_path = validate_log_directory(cfg.hyper.general, cfg.name)
     if cfg.download.auto:
         prepare_dataset(cfg.download)
 
-    dataloader = get_dataloader(cfg)
+    dataloader = create_dataloader(cfg)
     # TODO: get_device or rank, for DDP mode
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    trainer = Trainer(cfg, save_path, device)
+    trainer = ModelTrainer(cfg, save_path, device)
     trainer.train(dataloader, cfg.hyper.train.epoch)
 
 

tests/test_utils/test_dataaugment.py

@@ -9,7 +9,7 @@ from torchvision.transforms import functional as TF
 project_root = Path(__file__).resolve().parent.parent.parent
 sys.path.append(str(project_root))
 
-from yolo.utils.data_augment import Compose, HorizontalFlip, Mosaic, VerticalFlip
+from yolo.utils.data_augmentation import Compose, HorizontalFlip, Mosaic, VerticalFlip
 
 
 def test_horizontal_flip():

tests/test_utils/test_loss.py

@@ -8,7 +8,7 @@ from hydra import compose, initialize
 project_root = Path(__file__).resolve().parent.parent.parent
 sys.path.append(str(project_root))
 
-from yolo.utils.loss import YOLOLoss
+from yolo.utils.loss_functions import YOLOLoss
 
 
 @pytest.fixture

yolo/model/module.py

@@ -6,7 +6,7 @@ from loguru import logger
 from torch import Tensor, nn
 from torch.nn.common_types import _size_2_t
 
-from yolo.tools.module_helper import auto_pad, get_activation, round_up
+from yolo.utils.module_utils import auto_pad, create_activation_function, round_up
 
 
 # ----------- Basic Class ----------- #
@@ -26,7 +26,7 @@ class Conv(nn.Module):
         kwargs.setdefault("padding", auto_pad(kernel_size, **kwargs))
         self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, bias=False, **kwargs)
         self.bn = nn.BatchNorm2d(out_channels, eps=1e-3, momentum=3e-2)
-        self.act = get_activation(activation)
+        self.act = create_activation_function(activation)
 
     def forward(self, x: Tensor) -> Tensor:
         return self.act(self.bn(self.conv(x)))
@@ -109,7 +109,7 @@ class RepConv(nn.Module):
         **kwargs
     ):
         super().__init__()
-        self.act = get_activation(activation)
+        self.act = create_activation_function(activation)
         self.conv1 = Conv(in_channels, out_channels, kernel_size, activation=False, **kwargs)
         self.conv2 = Conv(in_channels, out_channels, 1, activation=False, **kwargs)
 

yolo/model/yolo.py

@@ -5,9 +5,9 @@ from loguru import logger
 from omegaconf import ListConfig, OmegaConf
 
 from yolo.config.config import Config, Model, YOLOLayer
-from yolo.tools.log_helper import log_model
-from yolo.tools.module_helper import get_layer_map
-from yolo.utils.drawer import draw_model
+from yolo.tools.drawer import draw_model
+from yolo.utils.logging_utils import log_model_structure
+from yolo.utils.module_utils import get_layer_map
 
 
 class YOLO(nn.Module):
@@ -125,6 +125,6 @@ def get_model(cfg: Config) -> YOLO:
     OmegaConf.set_struct(cfg.model, False)
     model = YOLO(cfg.model, cfg.hyper.data.class_num)
     logger.info("✅ Success load model")
-    log_model(model.model)
+    log_model_structure(model.model)
     # draw_model(model=model)
     return model

yolo/{utils/data_augment.py → tools/data_augmentation.py}

@@ -4,7 +4,7 @@ from PIL import Image
 from torchvision.transforms import functional as TF
 
 
-class Compose:
+class AugmentationComposer:
     """Composes several transforms together."""
 
     def __init__(self, transforms, image_size: int = 640):

yolo/{utils/dataloader.py → tools/data_loader.py}

@@ -13,13 +13,19 @@ from torchvision.transforms import functional as TF
 from tqdm.rich import tqdm
 
 from yolo.config.config import Config
-from yolo.tools.dataset_helper import (
-    create_image_info_dict,
-    find_labels_path,
-    get_scaled_segmentation,
+from yolo.tools.data_augmentation import (
+    AugmentationComposer,
+    HorizontalFlip,
+    MixUp,
+    Mosaic,
+    VerticalFlip,
+)
+from yolo.tools.drawer import draw_bboxes
+from yolo.utils.dataset_utils import (
+    create_image_metadata,
+    locate_label_paths,
+    scale_segmentation,
 )
-from yolo.utils.data_augment import Compose, HorizontalFlip, MixUp, Mosaic, VerticalFlip
-from yolo.utils.drawer import draw_bboxes
 
 
 class YoloDataset(Dataset):
@@ -30,7 +36,7 @@ class YoloDataset(Dataset):
         self.image_size = image_size
 
         transforms = [eval(aug)(prob) for aug, prob in augment_cfg.items()]
-        self.transform = Compose(transforms, self.image_size)
+        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)
 
@@ -68,10 +74,10 @@ class YoloDataset(Dataset):
             list: A list of tuples, each containing the path to an image file and its associated segmentation as a tensor.
         """
         images_path = path.join(dataset_path, "images", phase_name)
-        labels_path, data_type = find_labels_path(dataset_path, phase_name)
+        labels_path, data_type = locate_label_paths(dataset_path, phase_name)
         images_list = sorted(os.listdir(images_path))
         if data_type == "json":
-            annotations_index, image_info_dict = create_image_info_dict(labels_path)
+            annotations_index, image_info_dict = create_image_metadata(labels_path)
 
         data = []
         valid_inputs = 0
@@ -85,7 +91,7 @@ class YoloDataset(Dataset):
                 if image_info is None:
                     continue
                 annotations = annotations_index.get(image_info["id"], [])
-                image_seg_annotations = get_scaled_segmentation(annotations, image_info)
+                image_seg_annotations = scale_segmentation(annotations, image_info)
                 if not image_seg_annotations:
                     continue
 
@@ -191,13 +197,13 @@ class YoloDataLoader(DataLoader):
         return batch_images, batch_targets
 
 
-def get_dataloader(config):
+def create_dataloader(config):
     return YoloDataLoader(config)
 
 
 @hydra.main(config_path="../config", config_name="config", version_base=None)
 def main(cfg):
-    dataloader = get_dataloader(cfg)
+    dataloader = create_dataloader(cfg)
     draw_bboxes(*next(iter(dataloader)))
 
 
@@ -205,7 +211,7 @@ if __name__ == "__main__":
     import sys
 
     sys.path.append("./")
-    from tools.log_helper import custom_logger
+    from tools.logging_utils import custom_logger
 
     custom_logger()
     main()

yolo/tools/layer_helper.py

@@ -1,5 +0,0 @@
-import inspect
-
-import torch.nn as nn
-
-from yolo.model import module

yolo/{utils/loss.py → tools/loss_functions.py}

@@ -8,8 +8,13 @@ from torch import Tensor, nn
 from torch.nn import BCEWithLogitsLoss
 
 from yolo.config.config import Config
-from yolo.tools.bbox_helper import Anchor2Box, BoxMatcher, calculate_iou, make_anchor
-from yolo.tools.module_helper import make_chunk
+from yolo.utils.bounding_box_utils import (
+    AnchorBoxConverter,
+    BoxMatcher,
+    calculate_iou,
+    generate_anchors,
+)
+from yolo.utils.module_utils import divide_into_chunks
 
 
 class BCELoss(nn.Module):
@@ -78,14 +83,14 @@ class YOLOLoss:
         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 = make_anchor(self.image_size, self.strides, device)
+        self.anchors, self.scaler = generate_anchors(self.image_size, self.strides, device)
 
         self.cls = BCELoss()
         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.box_converter = Anchor2Box(cfg, device)
+        self.box_converter = AnchorBoxConverter(cfg, device)
 
     def separate_anchor(self, anchors):
         """
@@ -132,7 +137,7 @@ class DualLoss:
         targets[:, :, 1:] = targets[:, :, 1:] * self.loss.scale_up
 
         # TODO: Need Refactor this region, make it flexible!
-        predicts = make_chunk(predicts[0], 2)
+        predicts = divide_into_chunks(predicts[0], 2)
         aux_iou, aux_dfl, aux_cls = self.loss(predicts[0], targets)
         main_iou, main_dfl, main_cls = self.loss(predicts[1], targets)
 

yolo/tools/trainer.py

@@ -7,25 +7,29 @@ from torch.cuda.amp import GradScaler, autocast
 
 from yolo.config.config import Config, TrainConfig
 from yolo.model.yolo import get_model
-from yolo.tools.log_helper import CustomProgress
-from yolo.tools.model_helper import EMA, get_optimizer, get_scheduler
-from yolo.utils.loss import get_loss_function
+from yolo.tools.loss_functions import get_loss_function
+from yolo.utils.logging_utils import ProgressTracker
+from yolo.utils.model_utils import (
+    ExponentialMovingAverage,
+    create_optimizer,
+    create_scheduler,
+)
 
 
-class Trainer:
+class ModelTrainer:
     def __init__(self, cfg: Config, save_path: str, device):
         train_cfg: TrainConfig = cfg.hyper.train
         model = get_model(cfg)
 
         self.model = model.to(device)
         self.device = device
-        self.optimizer = get_optimizer(model, train_cfg.optimizer)
-        self.scheduler = get_scheduler(self.optimizer, train_cfg.scheduler)
+        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 = CustomProgress(cfg, save_path, use_wandb=True)
+        self.progress = ProgressTracker(cfg, save_path, use_wandb=True)
 
         if getattr(train_cfg.ema, "enabled", False):
-            self.ema = EMA(model, decay=train_cfg.ema.decay)
+            self.ema = ExponentialMovingAverage(model, decay=train_cfg.ema.decay)
         else:
             self.ema = None
         self.scaler = GradScaler()

yolo/utils/README.md

@@ -1,7 +0,0 @@
-task/train.py
-
-task/validate.py
-
-task/inference.py
-
-etc.

yolo/{tools/bbox_helper.py → utils/bounding_box_utils.py}

@@ -106,7 +106,7 @@ def transform_bbox(bbox: Tensor, indicator="xywh -> xyxy"):
     return bbox.to(dtype=data_type)
 
 
-def make_anchor(image_size: List[int], strides: List[int], device):
+def generate_anchors(image_size: List[int], strides: List[int], device):
     W, H = image_size
     anchors = []
     scaler = []
@@ -124,7 +124,7 @@ def make_anchor(image_size: List[int], strides: List[int], device):
     return all_anchors, all_scalers
 
 
-class Anchor2Box:
+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
@@ -132,7 +132,7 @@ class Anchor2Box:
         self.strides = cfg.model.anchor.strides
 
         self.scale_up = torch.tensor(self.image_size * 2, device=device)
-        self.anchors, self.scaler = make_anchor(self.image_size, self.strides, device)
+        self.anchors, self.scaler = generate_anchors(self.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:

yolo/{tools/dataset_helper.py → utils/dataset_utils.py}

@@ -6,10 +6,10 @@ from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
 
-from yolo.utils.converter_json2txt import discretize_categories
+from yolo.tools.data_conversion import discretize_categories
 
 
-def find_labels_path(dataset_path: str, phase_name: str):
+def locate_label_paths(dataset_path: str, phase_name: str):
     """
     Find the path to label files for a specified dataset and phase(e.g. training).
 
@@ -35,7 +35,7 @@ def find_labels_path(dataset_path: str, phase_name: str):
     raise FileNotFoundError("No labels found in the specified dataset path and phase name.")
 
 
-def create_image_info_dict(labels_path: str) -> Tuple[Dict[str, List], Dict[str, Dict]]:
+def create_image_metadata(labels_path: str) -> Tuple[Dict[str, List], Dict[str, Dict]]:
     """
     Create a dictionary containing image information and annotations indexed by image ID.
 
@@ -49,12 +49,12 @@ def create_image_info_dict(labels_path: str) -> Tuple[Dict[str, List], Dict[str,
     with open(labels_path, "r") as file:
         labels_data = json.load(file)
         id_to_idx = discretize_categories(labels_data.get("categories", [])) if "categories" in labels_data else None
-        annotations_index = index_annotations_by_image(labels_data, id_to_idx)  # check lookup is a good name?
+        annotations_index = organize_annotations_by_image(labels_data, id_to_idx)  # check lookup is a good name?
         image_info_dict = {path.splitext(img["file_name"])[0]: img for img in labels_data["images"]}
         return annotations_index, image_info_dict
 
 
-def index_annotations_by_image(data: Dict[str, Any], id_to_idx: Optional[Dict[int, int]]):
+def organize_annotations_by_image(data: Dict[str, Any], id_to_idx: Optional[Dict[int, int]]):
     """
     Use image index to lookup every annotations
     Args:
@@ -78,7 +78,7 @@ def index_annotations_by_image(data: Dict[str, Any], id_to_idx: Optional[Dict[in
     return annotation_lookup
 
 
-def get_scaled_segmentation(
+def scale_segmentation(
     annotations: List[Dict[str, Any]], image_dimensions: Dict[str, int]
 ) -> Optional[List[List[float]]]:
     """

yolo/{tools/log_helper.py → utils/logging_utils.py}

@@ -35,7 +35,7 @@ def custom_logger():
     )
 
 
-class CustomProgress:
+class ProgressTracker:
     def __init__(self, cfg: Config, save_path: str, use_wandb: bool = False):
         self.progress = Progress(
             TextColumn("[progress.description]{task.description}"),
@@ -87,7 +87,7 @@ def custom_wandb_log(string="", level=int, newline=True, repeat=True, prefix=Tru
         logger.opt(raw=not newline, colors=True).info("🌐 " + line)
 
 
-def log_model(model: List[YOLOLayer]):
+def log_model_structure(model: List[YOLOLayer]):
     console = Console()
     table = Table(title="Model Layers")
 
@@ -108,7 +108,7 @@ def log_model(model: List[YOLOLayer]):
     console.print(table)
 
 
-def get_valid_folder(general_cfg: GeneralConfig, exp_name):
+def validate_log_directory(general_cfg: GeneralConfig, exp_name):
     base_path = os.path.join(general_cfg.out_path, general_cfg.task)
     save_path = os.path.join(base_path, exp_name)
 

yolo/{tools/model_helper.py → utils/model_utils.py}

@@ -8,7 +8,7 @@ from yolo.config.config import OptimizerConfig, SchedulerConfig
 from yolo.model.yolo import YOLO
 
 
-class EMA:
+class ExponentialMovingAverage:
     def __init__(self, model: torch.nn.Module, decay: float):
         self.model = model
         self.decay = decay
@@ -32,7 +32,7 @@ class EMA:
             self.shadow[name].copy_(param.data)
 
 
-def get_optimizer(model: YOLO, optim_cfg: OptimizerConfig) -> Optimizer:
+def create_optimizer(model: YOLO, optim_cfg: OptimizerConfig) -> Optimizer:
     """Create an optimizer for the given model parameters based on the configuration.
 
     Returns:
@@ -52,7 +52,7 @@ def get_optimizer(model: YOLO, optim_cfg: OptimizerConfig) -> Optimizer:
     return optimizer_class(model_parameters, **optim_cfg.args)
 
 
-def get_scheduler(optimizer: Optimizer, schedule_cfg: SchedulerConfig) -> _LRScheduler:
+def create_scheduler(optimizer: Optimizer, schedule_cfg: SchedulerConfig) -> _LRScheduler:
     """Create a learning rate scheduler for the given optimizer based on the configuration.
 
     Returns:

yolo/{tools/module_helper.py → utils/module_utils.py}

@@ -5,20 +5,6 @@ from torch import Tensor, nn
 from torch.nn.common_types import _size_2_t
 
 
-def auto_pad(kernel_size: _size_2_t, dilation: _size_2_t = 1, **kwargs) -> Tuple[int, int]:
-    """
-    Auto Padding for the convolution blocks
-    """
-    if isinstance(kernel_size, int):
-        kernel_size = (kernel_size, kernel_size)
-    if isinstance(dilation, int):
-        dilation = (dilation, dilation)
-
-    pad_h = ((kernel_size[0] - 1) * dilation[0]) // 2
-    pad_w = ((kernel_size[1] - 1) * dilation[1]) // 2
-    return (pad_h, pad_w)
-
-
 def get_layer_map():
     """
     Dynamically generates a dictionary mapping class names to classes,
@@ -34,7 +20,21 @@ def get_layer_map():
     return layer_map
 
 
-def get_activation(activation: str) -> nn.Module:
+def auto_pad(kernel_size: _size_2_t, dilation: _size_2_t = 1, **kwargs) -> Tuple[int, int]:
+    """
+    Auto Padding for the convolution blocks
+    """
+    if isinstance(kernel_size, int):
+        kernel_size = (kernel_size, kernel_size)
+    if isinstance(dilation, int):
+        dilation = (dilation, dilation)
+
+    pad_h = ((kernel_size[0] - 1) * dilation[0]) // 2
+    pad_w = ((kernel_size[1] - 1) * dilation[1]) // 2
+    return (pad_h, pad_w)
+
+
+def create_activation_function(activation: str) -> nn.Module:
     """
     Retrieves an activation function from the PyTorch nn module based on its name, case-insensitively.
     """
@@ -59,7 +59,7 @@ def round_up(x: Union[int, Tensor], div: int = 1) -> Union[int, Tensor]:
     return x + (-x % div)
 
 
-def make_chunk(input_list, chunk_num):
+def divide_into_chunks(input_list, chunk_num):
     """
     Args: input_list: [0, 1, 2, 3, 4, 5], chunk: 2
     Return: [[0, 1, 2], [3, 4, 5]]