🔀 [Merge] branch 'DATASET' into TEST

.gitignore

@@ -111,8 +111,8 @@ dmypy.json
 
 # Machine learning specific folders and symlinks
 runs
-data
-datasets
+/data
+/datasets
 
 # Datasets and model checkpoints
 *.pth

README.md

@@ -3,7 +3,7 @@ An MIT license rewrite of YOLOv9
 
 ## To-Do Lists
 - [ ] Project Setup
-    - [ ] requirements
+    - [X] requirements
     - [ ] LICENSE
     - [ ] README
     - [ ] pytests

config/config.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from typing import List, Dict, Union
+
+
+@dataclass
+class Model:
+    anchor: List[List[int]]
+    model: Dict[str, List[Dict[str, Union[Dict, List, int]]]]
+
+
+@dataclass
+class Download:
+    auto: bool
+    path: str
+
+
+@dataclass
+class Config:
+    model: Model
+    download: Download

config/config.yaml

@@ -0,0 +1,10 @@
+hydra:
+  run:
+    dir: ./runs
+
+defaults:
+ - data: coco
+ - download: ../data/download
+ - augmentation: ../data/augmentation
+ - model: v7-base
+ - _self_

config/data/augmentation.yaml

@@ -0,0 +1,2 @@
+RandomHorizontalFlip: 0.5
+Mosaic: 0.5

config/data/coco.yaml

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

config/data/download.yaml

@@ -0,0 +1,17 @@
+auto: True
+path: data/coco
+images:
+  base_url: http://images.cocodataset.org/zips/
+  datasets:
+    train:
+      file_name: train2017.zip
+      file_num: 118287
+    val: 
+      file_name: val2017.zip
+      num_files: 5000
+    test:
+      file_name: test2017.zip
+      num_files: 40670
+hydra:
+  run:
+    dir: ./runs

model/yolo.py

@@ -1,25 +1,10 @@
-import inspect
 from typing import Any, Dict, List, Union
 
 import torch
 import torch.nn as nn
 from loguru import logger
-
-from model import module
-from utils.tools import load_model_cfg
-
-
-def get_layer_map():
-    """
-    Dynamically generates a dictionary mapping class names to classes,
-    filtering to include only those that are subclasses of nn.Module,
-    ensuring they are relevant neural network layers.
-    """
-    layer_map = {}
-    for name, obj in inspect.getmembers(module, inspect.isclass):
-        if issubclass(obj, nn.Module) and obj is not nn.Module:
-            layer_map[name] = obj
-    return layer_map
+from omegaconf import OmegaConf
+from tools.layer_helper import get_layer_map
 
 
 class YOLO(nn.Module):
@@ -35,16 +20,15 @@ class YOLO(nn.Module):
         super(YOLO, self).__init__()
         self.nc = model_cfg["nc"]
         self.layer_map = get_layer_map()  # Get the map Dict[str: Module]
-        self.build_model(model_cfg["model"])
+        self.build_model(model_cfg.model)
 
     def build_model(self, model_arch: Dict[str, List[Dict[str, Dict[str, Dict]]]]):
         model_list = nn.ModuleList()
         output_dim = [3]
         layer_indices_by_tag = {}
-
-        for arch_name, arch in model_arch.items():
+        for arch_name in model_arch:
             logger.info(f"🏗️  Building model-{arch_name}")
-            for layer_idx, layer_spec in enumerate(arch, start=1):
+            for layer_idx, layer_spec in enumerate(model_arch[arch_name], start=1):
                 layer_type, layer_info = next(iter(layer_spec.items()))
                 layer_args = layer_info.get("args", {})
                 source = layer_info.get("source", -1)
@@ -74,7 +58,7 @@ class YOLO(nn.Module):
         y = [x]
         output = []
         for layer in self.model:
-            if isinstance(layer.source, list):
+            if OmegaConf.is_list(layer.source):
                 model_input = [y[idx] for idx in layer.source]
             else:
                 model_input = y[layer.source]
@@ -113,6 +97,7 @@ def get_model(model_cfg: dict) -> YOLO:
     Returns:
         YOLO: An instance of the model defined by the given configuration.
     """
+    OmegaConf.set_struct(model_cfg, False)
     model = YOLO(model_cfg)
     logger.info("✅ Success load model")
     return model

requirements.txt

@@ -0,0 +1,9 @@
+hydra-core
+loguru
+numpy
+pytest
+pyyaml
+requests
+rich
+torch
+tqdm

tools/layer_helper.py

@@ -0,0 +1,20 @@
+import inspect
+import torch.nn as nn
+from model import module
+
+
+def auto_pad():
+    raise NotImplementedError
+
+
+def get_layer_map():
+    """
+    Dynamically generates a dictionary mapping class names to classes,
+    filtering to include only those that are subclasses of nn.Module,
+    ensuring they are relevant neural network layers.
+    """
+    layer_map = {}
+    for name, obj in inspect.getmembers(module, inspect.isclass):
+        if issubclass(obj, nn.Module) and obj is not nn.Module:
+            layer_map[name] = obj
+    return layer_map

tools/log_helper.py

@@ -0,0 +1,23 @@
+"""
+Module for initializing logging tools used in machine learning and data processing. 
+Supports integration with Weights & Biases (wandb), Loguru, TensorBoard, and other 
+logging frameworks as needed.
+
+This setup ensures consistent logging across various platforms, facilitating 
+effective monitoring and debugging.
+
+Example:
+    from tools.logger import custom_logger
+    custom_logger()
+"""
+
+import sys
+from loguru import logger
+
+
+def custom_logger():
+    logger.remove()
+    logger.add(
+        sys.stderr,
+        format="<green>{time:MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <level>{message}</level>",
+    )

train.py

@@ -1,26 +1,19 @@
-import argparse
 from loguru import logger
 from model.yolo import get_model
-from utils.tools import load_model_cfg, custom_logger
+from tools.log_helper import custom_logger
+from utils.get_dataset import prepare_dataset
+import hydra
+from config.config import Config
 
 
-def parse_arguments() -> argparse.Namespace:
-    """
-    Parse command-line arguments to get the model configuration file.
+@hydra.main(config_path="config", config_name="config", version_base=None)
+def main(cfg: Config):
+    if cfg.download.auto:
+        prepare_dataset(cfg.download)
 
-    Returns:
-        argparse.Namespace: The command-line arguments object with 'config' attribute.
-    """
-    parser = argparse.ArgumentParser(description="Load a YOLO model configuration and display the model.")
-    parser.add_argument(
-        "--model-config", type=str, default="v7-base", help="Name or path to the model configuration file."
-    )
-    return parser.parse_args()
+    model = get_model(cfg.model)
 
 
 if __name__ == "__main__":
     custom_logger()
-    args = parse_arguments()
-    model_cfg = load_model_cfg(args.model_config)
-    model = get_model(model_cfg)
-    logger.info("Success load model")
+    main()

utils/dataargument.py

@@ -0,0 +1,79 @@
+from PIL import Image
+import numpy as np
+import torch
+from torchvision.transforms import functional as TF
+
+
+class Compose:
+    """Composes several transforms together."""
+
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+        for transform in self.transforms:
+            if hasattr(transform, "set_parent"):
+                transform.set_parent(self)
+
+    def __call__(self, image, boxes):
+        for transform in self.transforms:
+            image, boxes = transform(image, boxes)
+        return image, boxes
+
+    def get_more_data(self):
+        raise NotImplementedError("This method should be overridden by subclass instances!")
+
+
+class RandomHorizontalFlip:
+    """Randomly horizontally flips the image along with the bounding boxes."""
+
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def __call__(self, image, boxes):
+        if torch.rand(1) < self.prob:
+            image = TF.hflip(image)
+            boxes[:, [1, 3]] = 1 - boxes[:, [3, 1]]
+        return image, boxes
+
+
+class Mosaic:
+    """Applies the Mosaic augmentation to a batch of images and their corresponding boxes."""
+
+    def __init__(self, prob=0.5):
+        self.prob = prob
+        self.parent = None
+
+    def set_parent(self, parent):
+        self.parent = parent
+
+    def __call__(self, image, boxes):
+        if torch.rand(1) >= self.prob:
+            return image, boxes
+
+        assert self.parent is not None, "Parent is not set. Mosaic cannot retrieve image size."
+
+        img_sz = self.parent.image_size  # Assuming `image_size` is defined in parent
+        more_data = self.parent.get_more_data(3)  # get 3 more images randomly
+
+        data = [(image, boxes)] + more_data
+        mosaic_image = Image.new("RGB", (2 * img_sz, 2 * img_sz))
+        vectors = np.array([(-1, -1), (0, -1), (-1, 0), (0, 0)])
+        center = np.array([img_sz, img_sz])
+        all_labels = []
+
+        for (image, boxes), vector in zip(data, vectors):
+            this_w, this_h = image.size
+            coord = tuple(center + vector * np.array([this_w, this_h]))
+
+            mosaic_image.paste(image, coord)
+            xmin, ymin, xmax, ymax = boxes[:, 1], boxes[:, 2], boxes[:, 3], boxes[:, 4]
+            xmin = (xmin * this_w + coord[0]) / (2 * img_sz)
+            xmax = (xmax * this_w + coord[0]) / (2 * img_sz)
+            ymin = (ymin * this_h + coord[1]) / (2 * img_sz)
+            ymax = (ymax * this_h + coord[1]) / (2 * img_sz)
+
+            adjusted_boxes = torch.stack([boxes[:, 0], xmin, ymin, xmax, ymax], dim=1)
+            all_labels.append(adjusted_boxes)
+
+        all_labels = torch.cat(all_labels, dim=0)
+        return mosaic_image, all_labels

utils/dataloader.py

@@ -0,0 +1,144 @@
+from PIL import Image
+from os import path, listdir
+
+import hydra
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from loguru import logger
+from tqdm.rich import tqdm
+import diskcache as dc
+from typing import Union
+from drawer import draw_bboxes
+from dataargument import Compose, RandomHorizontalFlip, Mosaic
+
+
+class YoloDataset(Dataset):
+    def __init__(self, dataset_cfg: dict, phase: str = "train", image_size: int = 640, transform=None):
+        phase_name = dataset_cfg.get(phase, phase)
+        self.image_size = image_size
+
+        self.transform = transform
+        self.transform.get_more_data = self.get_more_data
+        self.transform.image_size = self.image_size
+        self.data = self.load_data(dataset_cfg.path, phase_name)
+
+    def load_data(self, dataset_path, phase_name):
+        """
+        Loads data from a cache or generates a new cache for a specific dataset phase.
+
+        Parameters:
+            dataset_path (str): The root path to the dataset directory.
+            phase_name (str): The specific phase of the dataset (e.g., 'train', 'test') to load or generate data for.
+
+        Returns:
+            dict: The loaded data from the cache for the specified phase.
+        """
+        cache_path = path.join(dataset_path, ".cache")
+        cache = dc.Cache(cache_path)
+        data = cache.get(phase_name)
+
+        if data is None:
+            logger.info("Generating {} cache", phase_name)
+            images_path = path.join(dataset_path, phase_name, "images")
+            labels_path = path.join(dataset_path, phase_name, "labels")
+            data = self.filter_data(images_path, labels_path)
+            cache[phase_name] = data
+
+        cache.close()
+        logger.info("Loaded {} cache", phase_name)
+        data = cache[phase_name]
+        return data
+
+    def filter_data(self, images_path: str, labels_path: str) -> list:
+        """
+        Filters and collects dataset information by pairing images with their corresponding labels.
+
+        Parameters:
+            images_path (str): Path to the directory containing image files.
+            labels_path (str): Path to the directory containing label files.
+
+        Returns:
+            list: A list of tuples, each containing the path to an image file and its associated labels as a tensor.
+        """
+        data = []
+        valid_inputs = 0
+        images_list = sorted(listdir(images_path))
+        for image_name in tqdm(images_list, desc="Filtering data"):
+            if not image_name.lower().endswith((".jpg", ".jpeg", ".png")):
+                continue
+
+            img_path = path.join(images_path, image_name)
+            base_name, _ = path.splitext(image_name)
+            label_path = path.join(labels_path, f"{base_name}.txt")
+
+            if path.isfile(label_path):
+                labels = self.load_valid_labels(label_path)
+                if labels is not None:
+                    data.append((img_path, labels))
+                    valid_inputs += 1
+
+        logger.info("Recorded {}/{} valid inputs", valid_inputs, len(images_list))
+        return data
+
+    def load_valid_labels(self, label_path: str) -> Union[torch.Tensor, None]:
+        """
+        Loads and validates bounding box data is [0, 1] from a label file.
+
+        Parameters:
+            label_path (str): The filepath to the label file containing bounding box data.
+
+        Returns:
+            torch.Tensor or None: A tensor of all valid bounding boxes if any are found; otherwise, None.
+        """
+        bboxes = []
+        with open(label_path, "r") as file:
+            for line in file:
+                parts = list(map(float, line.strip().split()))
+                cls = parts[0]
+                points = np.array(parts[1:]).reshape(-1, 2)
+                valid_points = points[(points >= 0) & (points <= 1)].reshape(-1, 2)
+                if valid_points.size > 1:
+                    bbox = torch.tensor([cls, *valid_points.min(axis=0), *valid_points.max(axis=0)])
+                    bboxes.append(bbox)
+
+        if bboxes:
+            return torch.stack(bboxes)
+        else:
+            logger.warning("No valid BBox in {}", label_path)
+            return None
+
+    def get_data(self, idx):
+        img_path, bboxes = self.data[idx]
+        img = Image.open(img_path).convert("RGB")
+        return img, bboxes
+
+    def get_more_data(self, num: int = 1):
+        indices = torch.randint(0, len(self), (num,))
+        return [self.get_data(idx) for idx in indices]
+
+    def __getitem__(self, idx) -> Union[Image.Image, torch.Tensor]:
+        img, bboxes = self.get_data(idx)
+        if self.transform:
+            img, bboxes = self.transform(img, bboxes)
+        return img, bboxes
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+
+@hydra.main(config_path="../config", config_name="config", version_base=None)
+def main(cfg):
+    transform = Compose([eval(aug)(prob) for aug, prob in cfg.augmentation.items()])
+    dataset = YoloDataset(cfg.data, transform=transform)
+    draw_bboxes(*dataset[0])
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.path.append("./")
+    from tools.log_helper import custom_logger
+
+    custom_logger()
+    main()

utils/drawer.py

@@ -0,0 +1,32 @@
+from PIL import Image, ImageDraw, ImageFont
+
+
+def draw_bboxes(img, bboxes):
+    """
+    Draw bounding boxes on an image.
+
+    Args:
+    - image_path (str): Path to the image file.
+    - bboxes (list of lists/tuples): Bounding boxes with [x_min, y_min, x_max, y_max, class_id].
+    """
+    # Load an image
+    draw = ImageDraw.Draw(img)
+
+    # Font for class_id (optional)
+    try:
+        font = ImageFont.truetype("arial.ttf", 30)
+    except IOError:
+        font = ImageFont.load_default(30)
+    width, height = img.size
+
+    for bbox in bboxes:
+        class_id, x_min, y_min, x_max, y_max = bbox
+        x_min = x_min * width
+        x_max = x_max * width
+        y_min = y_min * height
+        y_max = y_max * height
+        shape = [(x_min, y_min), (x_max, y_max)]
+        draw.rectangle(shape, outline="red", width=2)
+        draw.text((x_min, y_min), str(int(class_id)), font=font, fill="blue")
+
+    img.save("output.jpg")

utils/get_dataset.py

@@ -0,0 +1,84 @@
+import os
+import zipfile
+
+import hydra
+from loguru import logger
+import requests
+from tqdm.rich import tqdm
+
+
+def download_file(url, dest_path):
+    """
+    Downloads a file from a specified URL to a destination path with progress logging.
+    """
+    logger.info(f"Downloading {os.path.basename(dest_path)}...")
+    with requests.get(url, stream=True) as r:
+        r.raise_for_status()
+        total_length = int(r.headers.get("content-length", 0))
+        with open(dest_path, "wb") as f, tqdm(
+            total=total_length, unit="iB", unit_scale=True, desc=os.path.basename(dest_path), leave=True
+        ) as bar:
+            for chunk in r.iter_content(chunk_size=1024 * 1024):
+                f.write(chunk)
+                bar.update(len(chunk))
+    logger.info("Download complete!")
+
+
+def unzip_file(zip_path, extract_to):
+    """
+    Unzips a ZIP file to a specified directory.
+    """
+    logger.info(f"Unzipping {os.path.basename(zip_path)}...")
+    with zipfile.ZipFile(zip_path, "r") as zip_ref:
+        zip_ref.extractall(extract_to)
+    os.remove(zip_path)
+    logger.info(f"Removed {zip_path}")
+
+
+def check_files(directory, expected_count):
+    """
+    Checks if the specified directory has the expected number of files.
+    """
+    num_files = len([name for name in os.listdir(directory) if os.path.isfile(os.path.join(directory, name))])
+    return num_files == expected_count
+
+
+@hydra.main(config_path="../config/data", config_name="download", version_base=None)
+def prepare_dataset(download_cfg):
+    data_dir = download_cfg.path
+    base_url = download_cfg.images.base_url
+    datasets = download_cfg.images.datasets
+
+    for dataset_type in datasets:
+        file_name, expected_files = datasets[dataset_type].values()
+        url = f"{base_url}{file_name}"
+        local_zip_path = os.path.join(data_dir, file_name)
+        extract_to = os.path.join(data_dir, dataset_type, "images")
+
+        # Ensure the extraction directory exists
+        os.makedirs(extract_to, exist_ok=True)
+
+        # Check if the correct number of files exists
+        if check_files(extract_to, expected_files):
+            logger.info(f"Dataset {dataset_type} already verified.")
+            continue
+
+        if os.path.exists(local_zip_path):
+            logger.info(f"Dataset {dataset_type} already downloaded.")
+        else:
+            download_file(url, local_zip_path)
+
+        unzip_file(local_zip_path, extract_to)
+
+        print(os.path.exists(local_zip_path), check_files(extract_to, expected_files))
+
+        # Additional verification post extraction
+        if not check_files(extract_to, expected_files):
+            logger.error(f"Error in verifying the {dataset_type} dataset after extraction.")
+
+
+if __name__ == "__main__":
+    from tools.log_helper import custom_logger
+
+    custom_logger()
+    prepare_dataset()

utils/tools.py

@@ -1,72 +0,0 @@
-import os
-import sys
-import yaml
-from loguru import logger
-from typing import Dict, Any
-
-
-def complete_path(file_name: str = "v7-base.yaml") -> str:
-    """
-    Ensures the path to a model configuration is a existing file
-
-    Parameters:
-        file_name (str): The filename or path, with default 'v7-base.yaml'.
-
-    Returns:
-        str: A complete path with necessary prefix and extension.
-    """
-    # Ensure the file has the '.yaml' extension if missing
-    if not file_name.endswith(".yaml"):
-        file_name += ".yaml"
-
-    # Add folder prefix if only the filename is provided
-    if os.path.dirname(file_name) == "":
-        file_name = os.path.join("./config/model", file_name)
-
-    return file_name
-
-
-def load_model_cfg(file_path: str) -> Dict[str, Any]:
-    """
-    Read a YAML configuration file, ensure necessary keys are present, and return its content as a dictionary.
-
-    Args:
-        file_path (str): The path to the YAML configuration file.
-
-    Returns:
-        Dict[str, Any]: The contents of the YAML file as a dictionary.
-
-    Raises:
-        FileNotFoundError: If the YAML file cannot be found.
-        yaml.YAMLError: If there is an error parsing the YAML file.
-    """
-    file_path = complete_path(file_path)
-    try:
-        with open(file_path, "r") as file:
-            model_cfg = yaml.safe_load(file) or {}
-
-        # Check for required keys and set defaults if not present
-        if "nc" not in model_cfg:
-            model_cfg["nc"] = 80
-            logger.warning("'nc' not found in the YAML file. Setting default 'nc' to 80.")
-
-        if "model" not in model_cfg:
-            logger.error("'model' is missing in the configuration file.")
-            raise ValueError("Missing required key: 'model'")
-
-        return model_cfg
-
-    except FileNotFoundError:
-        logger.error(f"YAML file not found: {file_path}")
-        raise
-    except yaml.YAMLError as e:
-        logger.error(f"Error parsing YAML file: {e}")
-        raise
-
-
-def custom_logger():
-    logger.remove()
-    logger.add(
-        sys.stderr,
-        format="<green>{time:MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <level>{message}</level>",
-    )