🔀 [Merge] branch 'DATASET' of https://github.com/WongKinYiu/yolov9mit into DATASET

.github/workflows/main.yaml

@@ -1,4 +1,4 @@
-name: YOLOv9 - Model test
+name: YOLOv9 - Model test and Code Style Check
 
 on:
   push:
@@ -8,7 +8,6 @@ on:
 
 jobs:
   build:
-
     runs-on: ubuntu-latest
 
     steps:
@@ -17,10 +16,25 @@ jobs:
       uses: actions/setup-python@v2
       with:
         python-version: 3.8
+
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
         pip install -r requirements.txt
+
+    - name: Install pre-commit
+      run: pip install pre-commit
+
+    - name: Cache pre-commit environment
+      uses: actions/cache@v2
+      with:
+        path: ~/.cache/pre-commit
+        key: ${{ runner.os }}-precommit-${{ hashFiles('**/.pre-commit-config.yaml') }}
+        restore-keys: |
+          ${{ runner.os }}-precommit-
+
+    - name: Run pre-commit (black and isort)
+      run: pre-commit run --all-files
+
     - name: Test with pytest
-      run: |
-        pytest
+      run: pytest

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Kin-Yiu, Wong
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -13,20 +13,22 @@ While the project's structure is still being finalized, we ask that potential co
 
 If you are interested in contributing, please keep an eye on project updates or contact us directly at [[email protected]](mailto:[email protected]) for more information.
 
+## Star History
 
+[![Star History Chart](https://api.star-history.com/svg?repos=WongKinYiu/yolov9mit&type=Date)](https://star-history.com/#WongKinYiu/yolov9mit&Date)
 
 ## To-Do Lists
 - [ ] Project Setup
     - [X] requirements
-    - [ ] LICENSE
+    - [x] LICENSE
     - [ ] README
-    - [ ] pytests
+    - [x] pytests
     - [ ] setup.py/pip install
-    - [ ] log format
+    - [x] log format
     - [ ] hugging face
 - [ ] Data proccess
     - [ ] Dataset
-        - [ ] Download script
+        - [x] Download script
         - [ ] Auto Download
         - [ ] xywh, xxyy, xcyc
     - [ ] Dataloder
@@ -35,14 +37,14 @@ If you are interested in contributing, please keep an eye on project updates or
     - [ ] load model
         - [ ] from yaml
         - [ ] from github
-    - [ ] trainer
-        - [ ] train_one_iter
-        - [ ] train_one_epoch
-    - [ ] DDP, EMA, OTA
+    - [x] trainer
+        - [x] train_one_iter
+        - [x] train_one_epoch
+    - [ ] DDP
+    - [x] EMA, OTA
+- [ ] Loss
 - [ ] Run
     - [ ] train
     - [ ] test
     - [ ] demo
-- [ ] Configuration
-    - [ ] hyperparams: dataclass
-    - [ ] model cfg: yaml
+- [x] Configuration

config/config.py

@@ -14,6 +14,57 @@ class Download:
     path: str
 
 
+@dataclass
+class DataLoaderConfig:
+    batch_size: int
+    shuffle: bool
+    num_workers: int
+    pin_memory: bool
+
+
+@dataclass
+class OptimizerArgs:
+    lr: float
+    weight_decay: float
+
+
+@dataclass
+class OptimizerConfig:
+    type: str
+    args: OptimizerArgs
+
+
+@dataclass
+class SchedulerArgs:
+    step_size: int
+    gamma: float
+
+
+@dataclass
+class SchedulerConfig:
+    type: str
+    args: SchedulerArgs
+
+
+@dataclass
+class EMAConfig:
+    enabled: bool
+    decay: float
+
+
+@dataclass
+class TrainConfig:
+    optimizer: OptimizerConfig
+    scheduler: SchedulerConfig
+    ema: EMAConfig
+
+
+@dataclass
+class HyperConfig:
+    data: DataLoaderConfig
+    train: TrainConfig
+
+
 @dataclass
 class Dataset:
     file_name: str
@@ -37,3 +88,4 @@ class Download:
 class Config:
     model: Model
     download: Download
+    hyper: HyperConfig

config/hyper/default.yaml

@@ -3,3 +3,17 @@ data:
   shuffle: True
   num_workers: 4
   pin_memory: True
+train:
+  optimizer:
+    type: Adam
+    args:
+      lr: 0.001
+      weight_decay: 0.0001
+  scheduler:
+    type: StepLR
+    args:
+      step_size: 10
+      gamma: 0.1
+  ema:
+      enabled: true
+      decay: 0.995

tests/test_utils/test_dataaugment.py

@@ -6,23 +6,22 @@ from PIL import Image
 from torchvision.transforms import functional as TF
 
 sys.path.append("./")
-from utils.data_augment import Compose, Mosaic, RandomHorizontalFlip
+from utils.data_augment import Compose, HorizontalFlip, Mosaic, VerticalFlip
 
 
-def test_random_horizontal_flip():
+def test_horizontal_flip():
     # Create a mock image and bounding boxes
     img = Image.new("RGB", (100, 100), color="red")
-    boxes = torch.tensor([[1, 0.1, 0.1, 0.9, 0.9]])  # class, xmin, ymin, xmax, ymax
+    boxes = torch.tensor([[1, 0.05, 0.1, 0.7, 0.9]])  # class, xmin, ymin, xmax, ymax
 
-    flip_transform = RandomHorizontalFlip(prob=1)  # Set probability to 1 to ensure flip
+    flip_transform = HorizontalFlip(prob=1)  # Set probability to 1 to ensure flip
     flipped_img, flipped_boxes = flip_transform(img, boxes)
 
     # Assert image is flipped by comparing it to a manually flipped image
     assert TF.hflip(img) == flipped_img
 
     # Assert bounding boxes are flipped correctly
-    expected_boxes = torch.tensor([[1, 0.1, 0.1, 0.9, 0.9]])
-    expected_boxes[:, [1, 3]] = 1 - expected_boxes[:, [3, 1]]
+    expected_boxes = torch.tensor([[1, 0.3, 0.1, 0.95, 0.9]])
     assert torch.allclose(flipped_boxes, expected_boxes), "Bounding boxes were not flipped correctly"
 
 
@@ -60,5 +59,5 @@ def test_mosaic():
     # Checks here would depend on the exact expected behavior of the mosaic function,
     # such as dimensions and content of the output image and boxes.
 
-    assert mosaic_img.size == (200, 200), "Mosaic image size should be doubled"
+    assert mosaic_img.size == (100, 100), "Mosaic image size should be same"
     assert len(mosaic_boxes) > 0, "Should have some bounding boxes"

tools/model_helper.py

@@ -0,0 +1,51 @@
+from typing import Any, Dict, Type
+
+import torch
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler
+
+from config.config import OptimizerConfig, SchedulerConfig
+
+
+class EMA:
+    def __init__(self, model: torch.nn.Module, decay: float):
+        self.model = model
+        self.decay = decay
+        self.shadow = {name: param.clone().detach() for name, param in model.named_parameters()}
+
+    def update(self):
+        """Update the shadow parameters using the current model parameters."""
+        for name, param in self.model.named_parameters():
+            assert name in self.shadow, "All model parameters should have a corresponding shadow parameter."
+            new_average = (1.0 - self.decay) * param.data + self.decay * self.shadow[name]
+            self.shadow[name] = new_average.clone()
+
+    def apply_shadow(self):
+        """Apply the shadow parameters to the model."""
+        for name, param in self.model.named_parameters():
+            param.data.copy_(self.shadow[name])
+
+    def restore(self):
+        """Restore the original parameters from the shadow."""
+        for name, param in self.model.named_parameters():
+            self.shadow[name].copy_(param.data)
+
+
+def get_optimizer(model_parameters, optim_cfg: OptimizerConfig) -> Optimizer:
+    """Create an optimizer for the given model parameters based on the configuration.
+
+    Returns:
+        An instance of the optimizer configured according to the provided settings.
+    """
+    optimizer_class: Type[Optimizer] = getattr(torch.optim, optim_cfg.type)
+    return optimizer_class(model_parameters, **optim_cfg.args)
+
+
+def get_scheduler(optimizer: Optimizer, schedul_cfg: SchedulerConfig) -> _LRScheduler:
+    """Create a learning rate scheduler for the given optimizer based on the configuration.
+
+    Returns:
+        An instance of the scheduler configured according to the provided settings.
+    """
+    scheduler_class: Type[_LRScheduler] = getattr(torch.optim.lr_scheduler, schedul_cfg.type)
+    return scheduler_class(optimizer, **schedul_cfg.args)

tools/trainer.py

@@ -0,0 +1,63 @@
+import torch
+from loguru import logger
+from tqdm import tqdm
+
+from config.config import TrainConfig
+from model.yolo import YOLO
+from tools.model_helper import EMA, get_optimizer, get_scheduler
+from utils.loss import get_loss_function
+
+
+class Trainer:
+    def __init__(self, model: YOLO, train_cfg: TrainConfig, device):
+        self.model = model.to(device)
+        self.device = device
+        self.optimizer = get_optimizer(model.parameters(), train_cfg.optimizer)
+        self.scheduler = get_scheduler(self.optimizer, train_cfg.scheduler)
+        self.loss_fn = get_loss_function()
+
+        if train_cfg.ema.get("enabled", False):
+            self.ema = EMA(model, decay=train_cfg.ema.decay)
+        else:
+            self.ema = None
+
+    def train_one_batch(self, data, targets):
+        data, targets = data.to(self.device), targets.to(self.device)
+        self.optimizer.zero_grad()
+        outputs = self.model(data)
+        loss = self.loss_fn(outputs, targets)
+        loss.backward()
+        self.optimizer.step()
+        if self.ema:
+            self.ema.update()
+        return loss.item()
+
+    def train_one_epoch(self, dataloader):
+        self.model.train()
+        total_loss = 0
+        for data, targets in tqdm(dataloader, desc="Training"):
+            loss = self.train_one_batch(data, targets)
+            total_loss += loss
+        if self.scheduler:
+            self.scheduler.step()
+        return total_loss / len(dataloader)
+
+    def save_checkpoint(self, epoch, filename="checkpoint.pt"):
+        checkpoint = {
+            "epoch": epoch,
+            "model_state_dict": self.model.state_dict(),
+            "optimizer_state_dict": self.optimizer.state_dict(),
+        }
+        if self.ema:
+            self.ema.apply_shadow()
+            checkpoint["model_state_dict_ema"] = self.model.state_dict()
+            self.ema.restore()
+        torch.save(checkpoint, filename)
+
+    def train(self, dataloader, num_epochs):
+        logger.info("start train")
+        for epoch in range(num_epochs):
+            epoch_loss = self.train_one_epoch(dataloader)
+            logger.info(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}")
+            if (epoch + 1) % 5 == 0:
+                self.save_checkpoint(epoch, f"checkpoint_epoch_{epoch+1}.pth")

train.py

@@ -1,20 +1,27 @@
 import hydra
+import torch
 from loguru import logger
 
 from config.config import Config
 from model.yolo import get_model
 from tools.log_helper import custom_logger
-from utils.dataloader import YoloDataset
+from tools.trainer import Trainer
+from utils.dataloader import get_dataloader
 from utils.get_dataset import prepare_dataset
 
 
 @hydra.main(config_path="config", config_name="config", version_base=None)
 def main(cfg: Config):
-    dataset = YoloDataset(cfg)
     if cfg.download.auto:
         prepare_dataset(cfg.download)
 
+    dataloader = get_dataloader(cfg)
     model = get_model(cfg.model)
+    # TODO: get_device or rank, for DDP mode
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    trainer = Trainer(model, cfg.hyper.train, device)
+    trainer.train(dataloader, 10)
 
 
 if __name__ == "__main__":

utils/converter_json2txt.py

@@ -0,0 +1,86 @@
+import json
+import os
+from typing import Dict, List, Optional
+
+from tqdm import tqdm
+
+
+def discretize_categories(categories: List[Dict[str, int]]) -> Dict[int, int]:
+    """
+    Maps each unique 'id' in the list of category dictionaries to a sequential integer index.
+    Indices are assigned based on the sorted 'id' values.
+    """
+    sorted_categories = sorted(categories, key=lambda category: category["id"])
+    return {category["id"]: index for index, category in enumerate(sorted_categories)}
+
+
+def process_annotations(
+    image_annotations: Dict[int, List[Dict]],
+    image_info_dict: Dict[int, tuple],
+    output_dir: str,
+    id_to_idx: Optional[Dict[int, int]] = None,
+) -> None:
+    """
+    Process and save annotations to files, with option to remap category IDs.
+    """
+    for image_id, annotations in tqdm(image_annotations.items(), desc="Processing annotations"):
+        file_path = os.path.join(output_dir, f"{image_id:0>12}.txt")
+        if not annotations:
+            continue
+        with open(file_path, "w") as file:
+            for annotation in annotations:
+                process_annotation(annotation, image_info_dict[image_id], id_to_idx, file)
+
+
+def process_annotation(annotation: Dict, image_dims: tuple, id_to_idx: Optional[Dict[int, int]], file) -> None:
+    """
+    Convert a single annotation's segmentation and write it to the open file handle.
+    """
+    category_id = annotation["category_id"]
+    segmentation = (
+        annotation["segmentation"][0]
+        if annotation["segmentation"] and isinstance(annotation["segmentation"][0], list)
+        else None
+    )
+
+    if segmentation is None:
+        return
+
+    img_width, img_height = image_dims
+    normalized_segmentation = normalize_segmentation(segmentation, img_width, img_height)
+
+    if id_to_idx:
+        category_id = id_to_idx.get(category_id, category_id)
+
+    file.write(f"{category_id} {' '.join(normalized_segmentation)}\n")
+
+
+def normalize_segmentation(segmentation: List[float], img_width: int, img_height: int) -> List[str]:
+    """
+    Normalize and format segmentation coordinates.
+    """
+    return [f"{x/img_width:.6f}" if i % 2 == 0 else f"{x/img_height:.6f}" for i, x in enumerate(segmentation)]
+
+
+def convert_annotations(json_file: str, output_dir: str) -> None:
+    """
+    Load annotation data from a JSON file and process all annotations.
+    """
+    with open(json_file) as file:
+        data = json.load(file)
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    image_info_dict = {img["id"]: (img["width"], img["height"]) for img in data.get("images", [])}
+    id_to_idx = discretize_categories(data.get("categories", [])) if "categories" in data else None
+    image_annotations = {img_id: [] for img_id in image_info_dict}
+
+    for annotation in data.get("annotations", []):
+        if not annotation.get("iscrowd", False):
+            image_annotations[annotation["image_id"]].append(annotation)
+
+    process_annotations(image_annotations, image_info_dict, output_dir, id_to_idx)
+
+
+convert_annotations("./data/coco/annotations/instances_train2017.json", "./data/coco/labels/train2017/")
+convert_annotations("./data/coco/annotations/instances_val2017.json", "./data/coco/labels/val2017/")

utils/dataloader.py

@@ -96,7 +96,7 @@ class YoloDataset(Dataset):
             cache[phase_name] = data
 
         cache.close()
-        logger.info("Loaded {} cache", phase_name)
+        logger.info("📦 Loaded {} cache", phase_name)
         data = cache[phase_name]
         return data
 

utils/loss.py

@@ -0,0 +1,2 @@
+def get_loss_function(*args, **kwargs):
+    raise NotImplementedError