Upload folder using huggingface_hub

.github/workflows/deploy.yaml

@@ -0,0 +1,69 @@
+name: Deploy Mode Validation & Inference
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  deploy:
+    runs-on: ${{ matrix.operating-system }}
+
+    strategy:
+      matrix:
+        operating-system: [ubuntu-latest, macos-latest]
+        python-version: [3.8, '3.10', '3.12']
+      fail-fast: false
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v2
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Cache pip dependencies
+      uses: actions/cache@v2
+      with:
+        path: ~/.cache/pip
+        key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements.txt') }}-${{ matrix.python-version }}
+        restore-keys: |
+          ${{ runner.os }}-pip-${{ matrix.python-version }}
+          ${{ runner.os }}-pip-
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -r requirements.txt
+
+    - name: Install YOLO package
+      run: pip install -e .
+
+    - name: Cache model weights
+      id: cache-weights
+      uses: actions/cache@v2
+      with:
+        path: weights
+        key: ${{ runner.os }}-weights
+        restore-keys: |
+          ${{ runner.os }}-weights
+
+    - name: Run Validation
+      run: |
+        python yolo/lazy.py task=validation dataset=mock
+        python yolo/lazy.py task=validation dataset=mock model=v9-s
+        python yolo/lazy.py task=validation dataset=mock name=AnyNameYouWant
+
+    - name: Run Inference
+      run: |
+        python yolo/lazy.py task=inference
+        python yolo/lazy.py task=inference model=v7
+        python yolo/lazy.py task=inference +quite=True
+        python yolo/lazy.py task=inference name=AnyNameYouWant
+        python yolo/lazy.py task=inference image_size=\[480,640]
+        python yolo/lazy.py task=inference task.nms.min_confidence=0.1
+        python yolo/lazy.py task=inference task.fast_inference=deploy
+        python yolo/lazy.py task=inference task.data.source=tests/data/images/val

.github/workflows/develop.yaml

@@ -0,0 +1,52 @@
+name: Developer Mode Build & Test
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  build_and_test:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v3
+
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.10'
+
+    - name: Cache pip dependencies
+      uses: actions/cache@v2
+      with:
+        path: ~/.cache/pip
+        key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements.txt') }}
+        restore-keys: |
+          ${{ runner.os }}-pip-
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -r requirements-dev.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 hooks
+      run: pre-commit run --all-files
+
+    - name: Test with pytest
+      run: |
+        pip install pytest pytest-cov
+        pytest --cov=yolo

.github/workflows/docker.yaml

@@ -0,0 +1,42 @@
+name: Docker Image Deploy to DockerHub
+
+on:
+    push:
+      branches: [main]
+    pull_request:
+      branches: [main]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v3
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Log in to Docker Hub
+        uses: docker/login-action@v2
+        with:
+          username: henrytsui000
+          password: ${{ secrets.DOCKER_PASSWORD }}
+
+      - name: Extract metadata (tags, labels) for Docker
+        id: meta
+        uses: docker/metadata-action@v4
+        with:
+          images: henrytsui000/yolo
+
+      - name: Build and push Docker image
+        uses: docker/build-push-action@v4
+        with:
+          context: .
+          file: docker/Dockerfile
+          push: true
+          tags: ${{ steps.meta.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels }}
+
+      - name: Image digest
+        run: echo ${{ steps.meta.outputs.digest }}

.gitignore

@@ -114,6 +114,7 @@ dmypy.json
 runs
 /data
 /datasets
+*/data
 
 # Datasets and model checkpoints
 *.pth

.readthedocs.yaml

@@ -0,0 +1,8 @@
+version: 2
+
+sphinx:
+  configuration: docs/conf.py
+
+python:
+  install:
+    - requirements: docs/requirements.txt

README.md

@@ -6,22 +6,29 @@ sdk_version: 4.36.1
 ---
 # YOLO: Official Implementation of YOLOv9, YOLOv7
 
+[![Documentation Status](https://readthedocs.org/projects/yolo-docs/badge/?version=latest)](https://yolo-docs.readthedocs.io/en/latest/?badge=latest)
 ![GitHub License](https://img.shields.io/github/license/WongKinYiu/YOLO)
 ![WIP](https://img.shields.io/badge/status-WIP-orange)
-[![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/henry000/YOLO)
-> [!IMPORTANT]
+![](https://img.shields.io/github/actions/workflow/status/WongKinYiu/YOLO/deploy.yaml)
+
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/yolov9-learning-what-you-want-to-learn-using/real-time-object-detection-on-coco)](https://paperswithcode.com/sota/real-time-object-detection-on-coco)
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)]()
+[![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-green)](https://huggingface.co/spaces/henry000/YOLO)
+
+<!-- > [!IMPORTANT]
 > This project is currently a Work In Progress and may undergo significant changes. It is not recommended for use in production environments until further notice. Please check back regularly for updates.
 >
-> Use of this code is at your own risk and discretion. It is advisable to consult with the project owner before deploying or integrating into any critical systems.
+> Use of this code is at your own risk and discretion. It is advisable to consult with the project owner before deploying or integrating into any critical systems. -->
 
 Welcome to the official implementation of YOLOv7 and YOLOv9. This repository will contains the complete codebase, pre-trained models, and detailed instructions for training and deploying YOLOv9.
 
 ## TL;DR
 - This is the official YOLO model implementation with an MIT License.
-- For quick deployment: you can enter directly in the terminal:
+- For quick deployment: you can directly install by pip+git:
 ```shell
-pip install git+git@github.com:WongKinYiu/YOLO.git
-yolo task=inference task.source=0 # source could be a single file, video, image folder, webcam ID
+pip install git+https://github.com/WongKinYiu/YOLO.git
+yolo task.data.source=0 # source could be a single file, video, image folder, webcam ID
 ```
 
 ## Introduction
@@ -29,7 +36,7 @@ yolo task=inference task.source=0 # source could be a single file, video, image
 - [**YOLOv7**: Trainable Bag-of-Freebies Sets New State-of-the-Art for Real-Time Object Detectors](https://arxiv.org/abs/2207.02696)
 
 ## Installation
-To get started with YOLOv9, clone this repository and install the required dependencies:
+To get started using YOLOv9's developer mode, we recommand you clone this repository and install the required dependencies:
 ```shell
 git clone git@github.com:WongKinYiu/YOLO.git
 cd YOLO
@@ -66,12 +73,13 @@ pip install -r requirements.txt
 These are simple examples. For more customization details, please refer to [Notebooks](examples) and lower-level modifications **[HOWTO](docs/HOWTO.md)**.
 
 ## Training
-To train YOLO on your dataset:
+To train YOLO on your machine/dataset:
 
-1. Modify the configuration file `data/config.yaml` to point to your dataset.
+1. Modify the configuration file `yolo/config/dataset/**.yaml` to point to your dataset.
 2. Run the training script:
 ```shell
-python yolo/lazy.py task=train task.data.batch_size=8 model=v9-c
+python yolo/lazy.py task=train dataset=** use_wandb=True
+python yolo/lazy.py task=train task.data.batch_size=8 model=v9-c weight=False # or more args
 ```
 
 ### Transfer Learning
@@ -81,27 +89,70 @@ python yolo/lazy.py task=train task.data.batch_size=8 model=v9-c dataset={datase
 ```
 
 ### Inference
-To evaluate the model performance, use:
+To use a model for object detection, use:
 ```shell
-python yolo/lazy.py task=inference weight=weights/v9-c.pt model=v9-c task.fast_inference=deploy # use deploy weight
-python yolo/lazy.py task=inference # if cloned from GitHub
-yolo task=inference task.data.source={Any} # if pip installed
+python yolo/lazy.py # if cloned from GitHub
+python yolo/lazy.py task=inference \ # default is inference
+                    name=AnyNameYouWant \ # AnyNameYouWant
+                    device=cpu \ # hardware cuda, cpu, mps
+                    model=v9-s \ # model version: v9-c, m, s
+                    task.nms.min_confidence=0.1 \ # nms config
+                    task.fast_inference=onnx \ # onnx, trt, deploy
+                    task.data.source=data/toy/images/train \ # file, dir, webcam
+                    +quite=True \ # Quite Output
+yolo task.data.source={Any Source} # if pip installed
+yolo task=inference task.data.source={Any}
 ```
 
-### Validation [WIP]
-To validate the model performance, use:
+### Validation
+To validate model performance, or generate a json file in COCO format:
 ```shell
-# Work In Progress...
+python yolo/lazy.py task=validation
+python yolo/lazy.py task=validation dataset=toy
 ```
 
 ## Contributing
-Contributions to the YOLOv9 project are welcome! See [CONTRIBUTING](docs/CONTRIBUTING.md) for guidelines on how to contribute.
+Contributions to the YOLO project are welcome! See [CONTRIBUTING](docs/CONTRIBUTING.md) for guidelines on how to contribute.
+
+### TODO Diagrams
+```mermaid
+flowchart TB
+    subgraph Features
+      Taskv7-->Segmentation["#35 Segmentation"]
+      Taskv7-->Classification["#34 Classification"]
+      Taskv9-->Segmentation
+      Taskv9-->Classification
+      Trainv7
+    end
+    subgraph Model
+      MODELv7-->v7-X
+      MODELv7-->v7-E6
+      MODELv7-->v7-E6E
+      MODELv9-->v9-T
+      MODELv9-->v9-S
+      MODELv9-->v9-E
+    end
+    subgraph Bugs
+      Fix-->Fix1["#12 mAP > 1"]
+      Fix-->Fix2["v9 Gradient Bump"]
+      Reply-->Reply1["#39"]
+      Reply-->Reply2["#36"]
+    end
+```
 
 ## Star History
 [![Star History Chart](https://api.star-history.com/svg?repos=WongKinYiu/YOLO&type=Date)](https://star-history.com/#WongKinYiu/YOLO&Date)
 
 ## Citations
 ```
+@misc{wang2022yolov7,
+      title={YOLOv7: Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors},
+      author={Chien-Yao Wang and Alexey Bochkovskiy and Hong-Yuan Mark Liao},
+      year={2022},
+      eprint={2207.02696},
+      archivePrefix={arXiv},
+      primaryClass={id='cs.CV' full_name='Computer Vision and Pattern Recognition' is_active=True alt_name=None in_archive='cs' is_general=False description='Covers image processing, computer vision, pattern recognition, and scene understanding. Roughly includes material in ACM Subject Classes I.2.10, I.4, and I.5.'}
+}
 @misc{wang2024yolov9,
       title={YOLOv9: Learning What You Want to Learn Using Programmable Gradient Information},
       author={Chien-Yao Wang and I-Hau Yeh and Hong-Yuan Mark Liao},

demo/hf_demo.py

@@ -11,7 +11,7 @@ from yolo import (
     AugmentationComposer,
     NMSConfig,
     PostProccess,
-    Vec2Box,
+    create_converter,
     create_model,
     draw_bboxes,
 )
@@ -25,22 +25,22 @@ def load_model(model_name, device):
     model_cfg.model.auxiliary = {}
     model = create_model(model_cfg, True)
     model.to(device).eval()
-    return model
+    return model, model_cfg
 
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = load_model(DEFAULT_MODEL, device)
-v2b = Vec2Box(model, IMAGE_SIZE, device)
-class_list = OmegaConf.load("yolo/config/general.yaml").class_list
+model, model_cfg = load_model(DEFAULT_MODEL, device)
+converter = create_converter(model_cfg.name, model, model_cfg.anchor, IMAGE_SIZE, device)
+class_list = OmegaConf.load("yolo/config/dataset/coco.yaml").class_list
 
 transform = AugmentationComposer([])
 
 
 def predict(model_name, image, nms_confidence, nms_iou):
-    global DEFAULT_MODEL, model, device, v2b, class_list, post_proccess
+    global DEFAULT_MODEL, model, device, converter, class_list, post_proccess
     if model_name != DEFAULT_MODEL:
-        model = load_model(model_name, device)
-        v2b = Vec2Box(model, IMAGE_SIZE, device)
+        model, model_cfg = load_model(model_name, device)
+        converter = create_converter(model_cfg.name, model, model_cfg.anchor, IMAGE_SIZE, device)
         DEFAULT_MODEL = model_name
 
     image_tensor, _, rev_tensor = transform(image)
@@ -49,7 +49,7 @@ def predict(model_name, image, nms_confidence, nms_iou):
     rev_tensor = rev_tensor.to(device)[None]
 
     nms_config = NMSConfig(nms_confidence, nms_iou)
-    post_proccess = PostProccess(v2b, nms_config)
+    post_proccess = PostProccess(converter, nms_config)
 
     with torch.no_grad():
         predict = model(image_tensor)

docker/.dockerignore

@@ -0,0 +1,32 @@
+.devcontainer
+.vscode
+.git
+.github
+.gitignore
+.pre-commit-config.yaml
+Dockerfile
+LICENSE
+CITATION.cff
+README.md
+CODE_OF_CONDUCT.md
+CONTRIBUTING.md
+SECURITY.md
+
+data
+demo
+docs
+examples
+runs
+tests
+yolo
+build
+dist
+
+*.pt
+*.pth
+*.onnx
+*.hef
+*.engine
+*.profile
+
+*.png

docker/Dockerfile

@@ -0,0 +1,64 @@
+ARG BASE_IMG=nvidia/cuda:12.4.1-devel-ubuntu22.04
+FROM ${BASE_IMG}
+
+ARG USERNAME=user
+ARG WORKDIR=/home/${USERNAME}/YOLO
+
+SHELL ["/bin/bash", "-c"]
+
+
+ENV CUDA_HOME=/usr/local/cuda
+ENV PATH=${PATH}:${CUDA_HOME}/bin
+ENV LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${CUDA_HOME}/lib64
+
+RUN apt-get update \
+    && apt-get install -y \
+        sudo \
+        curl \
+        gcc \
+        git \
+        make \
+        wget \
+        zlib1g \
+        protobuf-compiler \
+        libgl1-mesa-dev \
+        graphviz \
+        python-is-python3 \
+        python3-pip \
+    && apt clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# Make user
+RUN echo "root:root" | chpasswd \
+    && useradd \
+        --create-home \
+        --home-dir /home/${USERNAME} \
+        --shell /bin/bash \
+        --user-group \
+        --groups adm,sudo \
+        ${USERNAME} \
+    && echo "${USERNAME}:${USERNAME}" | chpasswd \
+    && cat /dev/null > /etc/sudoers.d/${USERNAME} \
+    && echo "%${USERNAME}    ALL=(ALL)   NOPASSWD:    ALL" >> \
+        /etc/sudoers.d/${USERNAME} \
+    && mkdir -p ${WORKDIR} \
+    && chown ${USERNAME}:${USERNAME} ${WORKDIR}
+
+USER ${USERNAME}
+WORKDIR ${WORKDIR}
+
+COPY docker/requirements-locked.txt /app/requirements-locked.txt
+COPY requirements.txt /app/requirements.txt
+
+# Install any needed packages specified in requirements.txt
+RUN pip install --no-cache-dir -r /app/requirements-locked.txt
+RUN pip install --no-cache-dir -r /app/requirements.txt
+
+
+RUN git clone https://github.com/WongKinYiu/YOLO.git .
+
+# Ensure pip-installed packages are available in the PATH
+RUN echo 'export PATH=${PATH}:${HOME}/.local/bin' >> ~/.bashrc
+
+# Optional: Source .bashrc to apply changes in the current session
+RUN source ~/.bashrc

docker/requirements-locked.txt

@@ -0,0 +1,5 @@
+pycocotools==2.0.7
+torch==2.2.1
+torchvision==0.17.1
+setuptools>=60.0
+numpy==1.23.5

docs/0_get_start/0_quick_start.md

@@ -0,0 +1 @@
+# Quick Start

docs/0_get_start/1_installations.md

@@ -0,0 +1,4 @@
+# Install YOLO
+
+## GPU (Cuda / MPS)
+## CPU only

docs/0_get_start/2_git.md

@@ -0,0 +1,3 @@
+# Install via Git& GitHub
+
+## [WIP]

docs/0_get_start/3_pypi.md

@@ -0,0 +1,3 @@
+# Install via PyPi(pip install)
+
+## [WIP]

docs/0_get_start/4_docker.md

@@ -0,0 +1,3 @@
+# Install in Docker
+
+## [WIP]

docs/0_get_start/5_conda.md

@@ -0,0 +1,3 @@
+# Install in Conda env
+
+## [WIP]

docs/1_tutorials/0_train.md

@@ -0,0 +1,5 @@
+# Train
+
+## Train on COCO2017
+
+## Train on Cusom Dataset

docs/1_tutorials/1_validation.md

@@ -0,0 +1,5 @@
+# Validation
+
+## Validation on COCO2017
+
+## Validation on Custom Dataset

docs/2_model_zoo/0_object_detection.md

@@ -0,0 +1,5 @@
+# Object Detection
+
+## YOLOv7
+
+## YOLOv9

docs/2_model_zoo/1_segmentation.md

@@ -0,0 +1,5 @@
+# Segmentations
+
+## YOLOv7
+
+## YOLOv9

docs/2_model_zoo/2_classification.md

@@ -0,0 +1,5 @@
+# Classification
+
+## YOLOv7
+
+## YOLOv9

docs/3_custom/0_model.md

@@ -0,0 +1 @@
+# Model

docs/3_custom/1_data_augment.md

@@ -0,0 +1 @@
+# Data Augment

docs/3_custom/2_loss.md

@@ -0,0 +1 @@
+# Loss

docs/3_custom/3_task.md

@@ -0,0 +1 @@
+# Task

docs/4_deploy/1_deploy.md

@@ -0,0 +1,3 @@
+# Deploy YOLOv9
+
+# Deploy YOLOv7

docs/4_deploy/2_onnx.md

@@ -0,0 +1 @@
+# Compile to ONNX

docs/4_deploy/3_tensorrt.md

@@ -0,0 +1 @@
+# Compile to TensorRT

docs/HOWTO.md

@@ -2,6 +2,30 @@
 
 To facilitate easy customization of the YOLO model, we've structured the codebase to allow for changes through configuration files and minimal code adjustments. This guide will walk you through the steps to customize various components of the model including the architecture, blocks, data loaders, and loss functions.
 
+## Examples
+
+```shell
+# Train
+python yolo/lazy.py task=train dataset=dev use_wandb=True
+
+# Validate
+python yolo/lazy.py task=validation
+python yolo/lazy.py task=validation model=v9-s
+python yolo/lazy.py task=validation dataset=toy
+python yolo/lazy.py task=validation dataset=toy name=validation
+
+# Inference
+python yolo/lazy.py task=inference
+python yolo/lazy.py task=inference device=cpu
+python yolo/lazy.py task=inference +quite=True
+python yolo/lazy.py task=inference name=AnyNameYouWant
+python yolo/lazy.py task=inference image_size=\[480,640]
+python yolo/lazy.py task=inference task.nms.min_confidence=0.1
+python yolo/lazy.py task=inference task.fast_inference=deploy
+python yolo/lazy.py task=inference task.fast_inference=onnx device=cpu
+python yolo/lazy.py task=inference task.data.source=data/toy/images/train
+```
+
 ## Custom Model Architecture
 
 You can change the model architecture simply by modifying the YAML configuration file. Here's how:

docs/Makefile

@@ -0,0 +1,20 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = .
+BUILDDIR      = _build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

docs/conf.py

@@ -0,0 +1,37 @@
+# Configuration file for the Sphinx documentation builder.
+#
+# For the full list of built-in configuration values, see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+# -- Project information -----------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
+
+project = "YOLO-docs"
+copyright = "2024, Kin-Yiu, Wong and Hao-Tang, Tsui"
+author = "Kin-Yiu, Wong and Hao-Tang, Tsui"
+
+# -- General configuration ---------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
+
+extensions = [
+    "sphinx.ext.autodoc",
+    "sphinx_rtd_theme",
+    "myst_parser",
+]
+
+myst_enable_extensions = [
+    "dollarmath",
+    "amsmath",
+    "deflist",
+]
+html_theme = "sphinx_rtd_theme"
+
+
+templates_path = ["_templates"]
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
+
+
+# -- Options for HTML output -------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
+
+html_static_path = ["_static"]

docs/index.rst

@@ -0,0 +1,81 @@
+YOLO documentation
+=======================
+
+Introduction
+------------
+
+YOLO (You Only Look Once) is a state-of-the-art, real-time object detection system that is designed for both efficiency and accuracy. This documentation provides comprehensive guidance on how to set up, configure, and effectively use YOLO for object detection tasks.
+
+**Note: This project and some sections of this documentation are currently a work in progress.**
+
+Project Features
+----------------
+
+- **Real-time Processing**: YOLO can process images in real-time with high accuracy, making it suitable for applications that require instant detection.
+- **Multitasking Capabilities**: Our enhanced version of YOLO supports multitasking, allowing it to handle multiple object detection tasks simultaneously.
+- **Open Source**: YOLO is open source, released under the MIT License, encouraging a broad community of developers to contribute and build upon the existing framework.
+
+Documentation Contents
+----------------------
+
+Explore our documentation:
+
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Get Started
+
+   0_get_start/0_quick_start.md
+   0_get_start/1_installations.md
+   0_get_start/2_git.md
+   0_get_start/3_pypi.md
+   0_get_start/4_docker.md
+   0_get_start/5_conda.md
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Tutorials
+
+   1_tutorials/0_train.md
+   1_tutorials/1_validation.md
+
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Model Zoo
+
+   2_model_zoo/0_object_detection.md
+   2_model_zoo/1_segmentation.md
+   2_model_zoo/2_classification.md
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Custom YOLO
+
+   3_custom/0_model.md
+   3_custom/1_data_augment.md
+   3_custom/2_loss.md
+   3_custom/3_task.md
+
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Deploy
+
+   4_deploy/1_deploy.md
+   4_deploy/2_onnx.md
+   4_deploy/3_tensorrt.md
+
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Deploy
+
+   4_deploy/1_deploy.md
+   4_deploy/2_onnx.md
+   4_deploy/3_tensorrt.md
+
+License
+-------
+
+YOLO is provided under the MIT License, which allows extensive freedom for reuse and distribution. See the LICENSE file for full license text.

docs/make.bat

@@ -0,0 +1,35 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=.
+set BUILDDIR=_build
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+	echo.installed, then set the SPHINXBUILD environment variable to point
+	echo.to the full path of the 'sphinx-build' executable. Alternatively you
+	echo.may add the Sphinx directory to PATH.
+	echo.
+	echo.If you don't have Sphinx installed, grab it from
+	echo.https://www.sphinx-doc.org/
+	exit /b 1
+)
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd

docs/requirements.txt

@@ -0,0 +1,3 @@
+sphinx
+sphinx_rtd_theme
+myst-parser

examples/notebook_smallobject.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -22,7 +22,6 @@
     "import torch\n",
     "from hydra import compose, initialize\n",
     "from PIL import Image \n",
-    "from einops import rearrange\n",
     "\n",
     "# Ensure that the necessary repository is cloned and installed. You may need to run: \n",
     "# git clone git@github.com:WongKinYiu/YOLO.git\n",
@@ -30,13 +29,23 @@
     "# pip install .\n",
     "project_root = Path().resolve().parent\n",
     "sys.path.append(str(project_root))\n",
-    "from yolo.config.config import NMSConfig\n",
-    "from yolo import AugmentationComposer, bbox_nms, Config, create_model, custom_logger, draw_bboxes, Vec2Box"
+    "\n",
+    "from yolo import (\n",
+    "    AugmentationComposer, \n",
+    "    Config, \n",
+    "    NMSConfig, \n",
+    "    PostProccess,\n",
+    "    bbox_nms, \n",
+    "    create_model, \n",
+    "    create_converter, \n",
+    "    custom_logger, \n",
+    "    draw_bboxes, \n",
+    ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -55,32 +64,49 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[38;2;0;51;133m[07/31 12:22:22]\u001b[0m \u001b[1m  INFO  \u001b[0m| \u001b[1m🚜 Building YOLO\u001b[0m\n",
+      "\u001b[38;2;0;51;133m[07/31 12:22:22]\u001b[0m \u001b[1m  INFO  \u001b[0m| \u001b[1m  🏗️  Building backbone\u001b[0m\n",
+      "\u001b[38;2;0;51;133m[07/31 12:22:22]\u001b[0m \u001b[1m  INFO  \u001b[0m| \u001b[1m  🏗️  Building neck\u001b[0m\n",
+      "\u001b[38;2;0;51;133m[07/31 12:22:22]\u001b[0m \u001b[1m  INFO  \u001b[0m| \u001b[1m  🏗️  Building head\u001b[0m\n",
+      "\u001b[38;2;0;51;133m[07/31 12:22:22]\u001b[0m \u001b[1m  INFO  \u001b[0m| \u001b[1m  🏗️  Building detection\u001b[0m\n",
+      "\u001b[38;2;0;51;133m[07/31 12:22:22]\u001b[0m \u001b[1m  INFO  \u001b[0m| \u001b[1m  🏗️  Building auxiliary\u001b[0m\n",
+      "\u001b[38;2;0;51;133m[07/31 12:22:22]\u001b[0m \u001b[1m  INFO  \u001b[0m| \u001b[1m✅ Success load model & weight\u001b[0m\n",
+      "\u001b[38;2;0;51;133m[07/31 12:22:22]\u001b[0m \u001b[1m  INFO  \u001b[0m| \u001b[1m🈶 Found stride of model [8, 16, 32]\u001b[0m\n"
+     ]
+    }
+   ],
    "source": [
     "with initialize(config_path=CONFIG_PATH, version_base=None, job_name=\"notebook_job\"):\n",
     "    cfg: Config = compose(config_name=CONFIG_NAME, overrides=[\"task=inference\", f\"task.data.source={IMAGE_PATH}\", f\"model={MODEL}\"])\n",
     "    model = create_model(cfg.model, class_num=CLASS_NUM).to(device)\n",
     "    transform = AugmentationComposer([], cfg.image_size)\n",
-    "    vec2box = Vec2Box(model, cfg.image_size, device)"
+    "    converter = create_converter(cfg.model.name, model, cfg.model.anchor, cfg.image_size, device)\n",
+    "    post_proccess = PostProccess(converter, NMSConfig(0.5, 0.9))\n",
+    "    "
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
     "pil_image = Image.open(IMAGE_PATH)\n",
     "image, bbox, rev_tensor = transform(pil_image)\n",
     "image = image.to(device)[None]\n",
-    "rev_tensor = rev_tensor.to(device)"
+    "rev_tensor = rev_tensor.to(device)[None]"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -104,9 +130,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "ename": "NameError",
+     "evalue": "name 'vec2box' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[7], line 4\u001b[0m\n\u001b[1;32m      2\u001b[0m     total_image, total_shift \u001b[38;5;241m=\u001b[39m slide_image(image)\n\u001b[1;32m      3\u001b[0m     predict \u001b[38;5;241m=\u001b[39m model(total_image)\n\u001b[0;32m----> 4\u001b[0m     pred_class, _, pred_bbox \u001b[38;5;241m=\u001b[39m \u001b[43mvec2box\u001b[49m(predict[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMain\u001b[39m\u001b[38;5;124m\"\u001b[39m])\n\u001b[1;32m      5\u001b[0m pred_bbox[\u001b[38;5;241m1\u001b[39m:] \u001b[38;5;241m=\u001b[39m (pred_bbox[\u001b[38;5;241m1\u001b[39m: ] \u001b[38;5;241m+\u001b[39m total_shift[:, \u001b[38;5;28;01mNone\u001b[39;00m]) \u001b[38;5;241m/\u001b[39m SLIDE\n\u001b[1;32m      6\u001b[0m pred_bbox \u001b[38;5;241m=\u001b[39m pred_bbox\u001b[38;5;241m.\u001b[39mview(\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m4\u001b[39m)\n",
+      "\u001b[0;31mNameError\u001b[0m: name 'vec2box' is not defined"
+     ]
+    }
+   ],
    "source": [
     "with torch.no_grad():\n",
     "    total_image, total_shift = slide_image(image)\n",
@@ -114,7 +152,9 @@
     "    pred_class, _, pred_bbox = vec2box(predict[\"Main\"])\n",
     "pred_bbox[1:] = (pred_bbox[1: ] + total_shift[:, None]) / SLIDE\n",
     "pred_bbox = pred_bbox.view(1, -1, 4)\n",
-    "pred_class = pred_class.view(1, -1, 80)"
+    "pred_class = pred_class.view(1, -1, 80)\n",
+    "pred_bbox = (pred_bbox - rev_tensor[:, None, 1:]) / rev_tensor[:, 0:1, None]\n",
+    "predict_box = bbox_nms(pred_class, pred_bbox, NMSConfig(0.3, 0.5))\n"
    ]
   },
   {
@@ -123,7 +163,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "pred_bbox = (pred_bbox / rev_tensor[0] - rev_tensor[None, None, 1:]) "
+    "draw_bboxes(pil_image, predict_box, idx2label=cfg.class_list)"
    ]
   },
   {
@@ -131,10 +171,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "predict_box = bbox_nms(pred_class, pred_bbox, NMSConfig(0.5, 0.5))\n",
-    "draw_bboxes(pil_image, predict_box, idx2label=cfg.class_list)"
-   ]
+   "source": []
   }
  ],
  "metadata": {

requirements-dev.txt

@@ -0,0 +1,5 @@
+-r requirements.txt
+pytest
+pytest-cov
+pre-commit
+pycocotools

requirements.txt

@@ -5,11 +5,9 @@ loguru
 numpy
 opencv-python
 Pillow
-pytest
-pyyaml
+pycocotools
 requests
 rich
 torch
 torchvision
-tqdm
 wandb

tests/conftest.py

@@ -0,0 +1,118 @@
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+from hydra import compose, initialize
+
+project_root = Path(__file__).resolve().parent.parent
+sys.path.append(str(project_root))
+
+from yolo import Anc2Box, Config, Vec2Box, create_converter, create_model
+from yolo.model.yolo import YOLO
+from yolo.tools.data_loader import StreamDataLoader, YoloDataLoader
+from yolo.tools.dataset_preparation import prepare_dataset
+from yolo.utils.logging_utils import ProgressLogger, set_seed
+
+
+def pytest_configure(config):
+    config.addinivalue_line("markers", "requires_cuda: mark test to run only if CUDA is available")
+
+
+def get_cfg(overrides=[]) -> Config:
+    config_path = "../yolo/config"
+    with initialize(config_path=config_path, version_base=None):
+        cfg: Config = compose(config_name="config", overrides=overrides)
+        set_seed(cfg.lucky_number)
+        return cfg
+
+
+@pytest.fixture(scope="session")
+def train_cfg() -> Config:
+    return get_cfg(overrides=["task=train", "dataset=mock"])
+
+
+@pytest.fixture(scope="session")
+def validation_cfg():
+    return get_cfg(overrides=["task=validation", "dataset=mock"])
+
+
+@pytest.fixture(scope="session")
+def inference_cfg():
+    return get_cfg(overrides=["task=inference"])
+
+
+@pytest.fixture(scope="session")
+def inference_v7_cfg():
+    return get_cfg(overrides=["task=inference", "model=v7"])
+
+
+@pytest.fixture(scope="session")
+def device():
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+@pytest.fixture(scope="session")
+def train_progress_logger(train_cfg: Config):
+    progress_logger = ProgressLogger(train_cfg, exp_name=train_cfg.name)
+    return progress_logger
+
+
+@pytest.fixture(scope="session")
+def validation_progress_logger(validation_cfg: Config):
+    progress_logger = ProgressLogger(validation_cfg, exp_name=validation_cfg.name)
+    return progress_logger
+
+
+@pytest.fixture(scope="session")
+def model(train_cfg: Config, device) -> YOLO:
+    model = create_model(train_cfg.model)
+    return model.to(device)
+
+
+@pytest.fixture(scope="session")
+def model_v7(inference_v7_cfg: Config, device) -> YOLO:
+    model = create_model(inference_v7_cfg.model)
+    return model.to(device)
+
+
+@pytest.fixture(scope="session")
+def vec2box(train_cfg: Config, model: YOLO, device) -> Vec2Box:
+    vec2box = create_converter(train_cfg.model.name, model, train_cfg.model.anchor, train_cfg.image_size, device)
+    return vec2box
+
+
+@pytest.fixture(scope="session")
+def anc2box(inference_v7_cfg: Config, model: YOLO, device) -> Anc2Box:
+    anc2box = create_converter(
+        inference_v7_cfg.model.name, model, inference_v7_cfg.model.anchor, inference_v7_cfg.image_size, device
+    )
+    return anc2box
+
+
+@pytest.fixture(scope="session")
+def train_dataloader(train_cfg: Config):
+    prepare_dataset(train_cfg.dataset, task="train")
+    return YoloDataLoader(train_cfg.task.data, train_cfg.dataset, train_cfg.task.task)
+
+
+@pytest.fixture(scope="session")
+def validation_dataloader(validation_cfg: Config):
+    prepare_dataset(validation_cfg.dataset, task="val")
+    return YoloDataLoader(validation_cfg.task.data, validation_cfg.dataset, validation_cfg.task.task)
+
+
+@pytest.fixture(scope="session")
+def file_stream_data_loader(inference_cfg: Config):
+    return StreamDataLoader(inference_cfg.task.data)
+
+
+@pytest.fixture(scope="session")
+def file_stream_data_loader_v7(inference_v7_cfg: Config):
+    return StreamDataLoader(inference_v7_cfg.task.data)
+
+
+@pytest.fixture(scope="session")
+def directory_stream_data_loader(inference_cfg: Config):
+    inference_cfg.task.data.source = "tests/data/images/train"
+    return StreamDataLoader(inference_cfg.task.data)

tests/test_model/test_module.py

@@ -43,13 +43,6 @@ def test_adown():
     assert out.shape == (1, OUT_CHANNELS, 32, 32)
 
 
-def test_adown():
-    adown = ADown(IN_CHANNELS, OUT_CHANNELS)
-    x = torch.randn(1, IN_CHANNELS, 64, 64)
-    out = adown(x)
-    assert out.shape == (1, OUT_CHANNELS, 32, 32)
-
-
 def test_cblinear():
     cblinear = CBLinear(IN_CHANNELS, [5, 5])
     x = torch.randn(1, IN_CHANNELS, 64, 64)

tests/test_model/test_yolo.py

@@ -16,7 +16,7 @@ config_path = "../../yolo/config"
 config_name = "config"
 
 
-def test_build_model():
+def test_build_model_v9c():
     with initialize(config_path=config_path, version_base=None):
         cfg: Config = compose(config_name=config_name)
 
@@ -26,6 +26,26 @@ def test_build_model():
         assert len(model.model) == 39
 
 
+def test_build_model_v9m():
+    with initialize(config_path=config_path, version_base=None):
+        cfg: Config = compose(config_name=config_name, overrides=[f"model=v9-m"])
+
+        OmegaConf.set_struct(cfg.model, False)
+        cfg.weight = None
+        model = YOLO(cfg.model)
+        assert len(model.model) == 39
+
+
+def test_build_model_v7():
+    with initialize(config_path=config_path, version_base=None):
+        cfg: Config = compose(config_name=config_name, overrides=[f"model=v7"])
+
+        OmegaConf.set_struct(cfg.model, False)
+        cfg.weight = None
+        model = YOLO(cfg.model)
+        assert len(model.model) == 106
+
+
 @pytest.fixture
 def cfg() -> Config:
     with initialize(config_path="../../yolo/config", version_base=None):

tests/test_tools/test_data_augmentation.py

@@ -0,0 +1,71 @@
+import sys
+from pathlib import Path
+
+import torch
+from PIL import Image
+from torchvision.transforms import functional as TF
+
+project_root = Path(__file__).resolve().parent.parent.parent
+sys.path.append(str(project_root))
+
+from yolo.tools.data_augmentation import (
+    AugmentationComposer,
+    HorizontalFlip,
+    Mosaic,
+    VerticalFlip,
+)
+
+
+def test_horizontal_flip():
+    # Create a mock image and bounding boxes
+    img = Image.new("RGB", (100, 100), color="red")
+    boxes = torch.tensor([[1, 0.05, 0.1, 0.7, 0.9]])  # class, xmin, ymin, xmax, ymax
+
+    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.3, 0.1, 0.95, 0.9]])
+    assert torch.allclose(flipped_boxes, expected_boxes), "Bounding boxes were not flipped correctly"
+
+
+def test_compose():
+    # Define two mock transforms that simply return the inputs
+    def mock_transform(image, boxes):
+        return image, boxes
+
+    compose = AugmentationComposer([mock_transform, mock_transform])
+    img = Image.new("RGB", (640, 640), color="blue")
+    boxes = torch.tensor([[0, 0.2, 0.2, 0.8, 0.8]])
+
+    transformed_img, transformed_boxes, rev_tensor = compose(img, boxes)
+    tensor_img = TF.pil_to_tensor(img).to(torch.float32) / 255
+
+    assert (transformed_img == tensor_img).all(), "Image should not be altered"
+    assert torch.equal(transformed_boxes, boxes), "Boxes should not be altered"
+
+
+def test_mosaic():
+    img = Image.new("RGB", (100, 100), color="green")
+    boxes = torch.tensor([[0, 0.25, 0.25, 0.75, 0.75]])
+
+    # Mock parent with image_size and get_more_data method
+    class MockParent:
+        image_size = (100, 100)
+
+        def get_more_data(self, num_images):
+            return [(img, boxes) for _ in range(num_images)]
+
+    mosaic = Mosaic(prob=1)  # Ensure mosaic is applied
+    mosaic.set_parent(MockParent())
+
+    mosaic_img, mosaic_boxes = mosaic(img, boxes)
+
+    # 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 == (100, 100), "Mosaic image size should be same"
+    assert len(mosaic_boxes) > 0, "Should have some bounding boxes"

tests/test_tools/test_data_loader.py

@@ -0,0 +1,69 @@
+import sys
+from pathlib import Path
+
+project_root = Path(__file__).resolve().parent.parent.parent
+sys.path.append(str(project_root))
+
+from yolo.config.config import Config
+from yolo.tools.data_loader import StreamDataLoader, YoloDataLoader, create_dataloader
+
+
+def test_create_dataloader_cache(train_cfg: Config):
+    train_cfg.task.data.shuffle = False
+    train_cfg.task.data.batch_size = 2
+
+    cache_file = Path("tests/data/train.cache")
+    cache_file.unlink(missing_ok=True)
+
+    make_cache_loader = create_dataloader(train_cfg.task.data, train_cfg.dataset)
+    load_cache_loader = create_dataloader(train_cfg.task.data, train_cfg.dataset)
+    m_batch_size, m_images, _, m_reverse_tensors, m_image_paths = next(iter(make_cache_loader))
+    l_batch_size, l_images, _, l_reverse_tensors, l_image_paths = next(iter(load_cache_loader))
+    assert m_batch_size == l_batch_size
+    assert m_images.shape == l_images.shape
+    assert m_reverse_tensors.shape == l_reverse_tensors.shape
+    assert m_image_paths == l_image_paths
+
+
+def test_training_data_loader_correctness(train_dataloader: YoloDataLoader):
+    """Test that the training data loader produces correctly shaped data and metadata."""
+    batch_size, images, _, reverse_tensors, image_paths = next(iter(train_dataloader))
+    assert batch_size == 2
+    assert images.shape == (2, 3, 640, 640)
+    assert reverse_tensors.shape == (2, 5)
+    expected_paths = [
+        Path("tests/data/images/train/000000050725.jpg"),
+        Path("tests/data/images/train/000000167848.jpg"),
+    ]
+    assert list(image_paths) == list(expected_paths)
+
+
+def test_validation_data_loader_correctness(validation_dataloader: YoloDataLoader):
+    batch_size, images, targets, reverse_tensors, image_paths = next(iter(validation_dataloader))
+    assert batch_size == 4
+    assert images.shape == (4, 3, 640, 640)
+    assert targets.shape == (4, 18, 5)
+    assert reverse_tensors.shape == (4, 5)
+    expected_paths = [
+        Path("tests/data/images/val/000000151480.jpg"),
+        Path("tests/data/images/val/000000284106.jpg"),
+        Path("tests/data/images/val/000000323571.jpg"),
+        Path("tests/data/images/val/000000570456.jpg"),
+    ]
+    assert list(image_paths) == list(expected_paths)
+
+
+def test_file_stream_data_loader_frame(file_stream_data_loader: StreamDataLoader):
+    """Test the frame output from the file stream data loader."""
+    frame, rev_tensor, origin_frame = next(iter(file_stream_data_loader))
+    assert frame.shape == (1, 3, 640, 640)
+    assert rev_tensor.shape == (1, 5)
+    assert origin_frame.size == (1024, 768)
+
+
+def test_directory_stream_data_loader_frame(directory_stream_data_loader: StreamDataLoader):
+    """Test the frame output from the directory stream data loader."""
+    frame, rev_tensor, origin_frame = next(iter(directory_stream_data_loader))
+    assert frame.shape == (1, 3, 640, 640)
+    assert rev_tensor.shape == (1, 5)
+    assert origin_frame.size == (480, 640) or origin_frame.size == (512, 640)

tests/test_tools/test_dataset_preparation.py

@@ -0,0 +1,30 @@
+import os
+import shutil
+import sys
+from pathlib import Path
+
+project_root = Path(__file__).resolve().parent.parent.parent
+sys.path.append(str(project_root))
+
+from yolo.config.config import Config
+from yolo.tools.dataset_preparation import prepare_dataset, prepare_weight
+
+
+def test_prepare_dataset(train_cfg: Config):
+    dataset_path = Path("tests/data")
+    if dataset_path.exists():
+        shutil.rmtree(dataset_path)
+    prepare_dataset(train_cfg.dataset, task="train")
+    prepare_dataset(train_cfg.dataset, task="val")
+
+    images_path = Path("tests/data/images")
+    for data_type in images_path.iterdir():
+        assert len(os.listdir(data_type)) == 5
+
+    annotations_path = Path("tests/data/annotations")
+    assert "instances_val.json" in os.listdir(annotations_path)
+    assert "instances_train.json" in os.listdir(annotations_path)
+
+
+def test_prepare_weight():
+    prepare_weight()

tests/test_tools/test_drawer.py

@@ -0,0 +1,29 @@
+import sys
+from pathlib import Path
+
+from PIL import Image
+from torch import tensor
+
+project_root = Path(__file__).resolve().parent.parent.parent
+sys.path.append(str(project_root))
+
+from yolo.config.config import Config
+from yolo.model.yolo import YOLO
+from yolo.tools.drawer import draw_bboxes, draw_model
+
+
+def test_draw_model_by_config(train_cfg: Config):
+    """Test the drawing of a model based on a configuration."""
+    draw_model(model_cfg=train_cfg.model)
+
+
+def test_draw_model_by_model(model: YOLO):
+    """Test the drawing of a YOLO model."""
+    draw_model(model=model)
+
+
+def test_draw_bboxes():
+    """Test drawing bounding boxes on an image."""
+    predictions = tensor([[0, 60, 60, 160, 160, 0.5], [0, 40, 40, 120, 120, 0.5]])
+    pil_image = Image.open("tests/data/images/train/000000050725.jpg")
+    draw_bboxes(pil_image, [predictions])

tests/test_tools/test_loss_functions.py

@@ -0,0 +1,56 @@
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+from hydra import compose, initialize
+
+project_root = Path(__file__).resolve().parent.parent.parent
+sys.path.append(str(project_root))
+
+from yolo.config.config import Config
+from yolo.model.yolo import create_model
+from yolo.tools.loss_functions import DualLoss, create_loss_function
+from yolo.utils.bounding_box_utils import Vec2Box
+
+
+@pytest.fixture
+def cfg() -> Config:
+    with initialize(config_path="../../yolo/config", version_base=None):
+        cfg = compose(config_name="config", overrides=["task=train"])
+    return cfg
+
+
+@pytest.fixture
+def model(cfg: Config):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = create_model(cfg.model, weight_path=None)
+    return model.to(device)
+
+
+@pytest.fixture
+def vec2box(cfg: Config, model):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return Vec2Box(model, cfg.model.anchor, cfg.image_size, device)
+
+
+@pytest.fixture
+def loss_function(cfg, vec2box) -> DualLoss:
+    return create_loss_function(cfg, vec2box)
+
+
+@pytest.fixture
+def data():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    targets = torch.zeros(1, 20, 5, device=device)
+    predicts = [torch.zeros(1, 8400, *cn, device=device) for cn in [(80,), (4, 16), (4,)]]
+    return predicts, targets
+
+
+def test_yolo_loss(loss_function, data):
+    predicts, targets = data
+    loss, loss_dict = loss_function(predicts, predicts, targets)
+    assert torch.isnan(loss)
+    assert torch.isnan(loss_dict["BoxLoss"])
+    assert torch.isnan(loss_dict["DFLoss"])
+    assert torch.isinf(loss_dict["BCELoss"])

tests/test_tools/test_solver.py

@@ -0,0 +1,80 @@
+import sys
+from pathlib import Path
+
+import pytest
+from torch import allclose, tensor
+
+project_root = Path(__file__).resolve().parent.parent.parent
+sys.path.append(str(project_root))
+
+from yolo.config.config import Config
+from yolo.model.yolo import YOLO
+from yolo.tools.data_loader import StreamDataLoader, YoloDataLoader
+from yolo.tools.solver import ModelTester, ModelTrainer, ModelValidator
+from yolo.utils.bounding_box_utils import Anc2Box, Vec2Box
+
+
+@pytest.fixture
+def model_validator(validation_cfg: Config, model: YOLO, vec2box: Vec2Box, validation_progress_logger, device):
+    validator = ModelValidator(
+        validation_cfg.task, validation_cfg.dataset, model, vec2box, validation_progress_logger, device
+    )
+    return validator
+
+
+def test_model_validator_initialization(model_validator: ModelValidator):
+    assert isinstance(model_validator.model, YOLO)
+    assert hasattr(model_validator, "solve")
+
+
+def test_model_validator_solve_mock_dataset(model_validator: ModelValidator, validation_dataloader: YoloDataLoader):
+    mAPs = model_validator.solve(validation_dataloader)
+    except_mAPs = {"mAP.5": tensor(0.6969), "mAP.5:.95": tensor(0.4195)}
+    assert allclose(mAPs["mAP.5"], except_mAPs["mAP.5"], rtol=0.1)
+    print(mAPs)
+    assert allclose(mAPs["mAP.5:.95"], except_mAPs["mAP.5:.95"], rtol=0.1)
+
+
+@pytest.fixture
+def model_tester(inference_cfg: Config, model: YOLO, vec2box: Vec2Box, validation_progress_logger, device):
+    tester = ModelTester(inference_cfg, model, vec2box, validation_progress_logger, device)
+    return tester
+
+
+@pytest.fixture
+def modelv7_tester(inference_v7_cfg: Config, model_v7: YOLO, anc2box: Anc2Box, validation_progress_logger, device):
+    tester = ModelTester(inference_v7_cfg, model_v7, anc2box, validation_progress_logger, device)
+    return tester
+
+
+def test_model_tester_initialization(model_tester: ModelTester):
+    assert isinstance(model_tester.model, YOLO)
+    assert hasattr(model_tester, "solve")
+
+
+def test_model_tester_solve_single_image(model_tester: ModelTester, file_stream_data_loader: StreamDataLoader):
+    model_tester.solve(file_stream_data_loader)
+
+
+def test_modelv7_tester_solve_single_image(modelv7_tester: ModelTester, file_stream_data_loader_v7: StreamDataLoader):
+    modelv7_tester.solve(file_stream_data_loader_v7)
+
+
+@pytest.fixture
+def model_trainer(train_cfg: Config, model: YOLO, vec2box: Vec2Box, train_progress_logger, device):
+    train_cfg.task.epoch = 2
+    trainer = ModelTrainer(train_cfg, model, vec2box, train_progress_logger, device, use_ddp=False)
+    return trainer
+
+
+def test_model_trainer_initialization(model_trainer: ModelTrainer):
+
+    assert isinstance(model_trainer.model, YOLO)
+    assert hasattr(model_trainer, "solve")
+    assert model_trainer.optimizer is not None
+    assert model_trainer.scheduler is not None
+    assert model_trainer.loss_fn is not None
+
+
+# def test_model_trainer_solve_mock_dataset(model_trainer: ModelTrainer, train_dataloader: YoloDataLoader):
+#     model_trainer.solve(train_dataloader)

tests/test_utils/test_bounding_box_utils.py

@@ -0,0 +1,183 @@
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+from hydra import compose, initialize
+from torch import allclose, float32, isclose, tensor
+
+project_root = Path(__file__).resolve().parent.parent.parent
+sys.path.append(str(project_root))
+from yolo import Config, NMSConfig, create_model
+from yolo.config.config import AnchorConfig
+from yolo.utils.bounding_box_utils import (
+    Anc2Box,
+    Vec2Box,
+    bbox_nms,
+    calculate_iou,
+    calculate_map,
+    generate_anchors,
+    transform_bbox,
+)
+
+EPS = 1e-4
+
+
+@pytest.fixture
+def dummy_bboxes():
+    bbox1 = tensor([[50, 80, 150, 140], [30, 20, 100, 80]], dtype=float32)
+    bbox2 = tensor([[90, 70, 160, 160], [40, 40, 90, 120]], dtype=float32)
+    return bbox1, bbox2
+
+
+def test_calculate_iou_2d(dummy_bboxes):
+    bbox1, bbox2 = dummy_bboxes
+    iou = calculate_iou(bbox1, bbox2)
+    expected_iou = tensor([[0.4138, 0.1905], [0.0096, 0.3226]])
+    assert iou.shape == (2, 2)
+    assert allclose(iou, expected_iou, atol=EPS)
+
+
+def test_calculate_iou_3d(dummy_bboxes):
+    bbox1, bbox2 = dummy_bboxes
+    iou = calculate_iou(bbox1[None], bbox2[None])
+    expected_iou = tensor([[0.4138, 0.1905], [0.0096, 0.3226]])
+    assert iou.shape == (1, 2, 2)
+    assert allclose(iou, expected_iou, atol=EPS)
+
+
+def test_calculate_diou(dummy_bboxes):
+    bbox1, bbox2 = dummy_bboxes
+    iou = calculate_iou(bbox1, bbox2, "diou")
+    expected_diou = tensor([[0.3816, 0.0943], [-0.2048, 0.2622]])
+
+    assert iou.shape == (2, 2)
+    assert allclose(iou, expected_diou, atol=EPS)
+
+
+def test_calculate_ciou(dummy_bboxes):
+    bbox1, bbox2 = dummy_bboxes
+    iou = calculate_iou(bbox1, bbox2, metrics="ciou")
+    # TODO: check result!
+    expected_ciou = tensor([[0.3769, 0.0853], [-0.2050, 0.2602]])
+    assert iou.shape == (2, 2)
+    assert allclose(iou, expected_ciou, atol=EPS)
+
+    bbox1 = tensor([[50, 80, 150, 140], [30, 20, 100, 80]], dtype=float32)
+    bbox2 = tensor([[90, 70, 160, 160], [40, 40, 90, 120]], dtype=float32)
+
+
+def test_transform_bbox_xywh_to_Any(dummy_bboxes):
+    bbox1, _ = dummy_bboxes
+    transformed_bbox = transform_bbox(bbox1, "xywh -> xyxy")
+    expected_bbox = tensor([[50.0, 80.0, 200.0, 220.0], [30.0, 20.0, 130.0, 100.0]])
+    assert allclose(transformed_bbox, expected_bbox)
+
+
+def test_transform_bbox_xycwh_to_Any(dummy_bboxes):
+    bbox1, bbox2 = dummy_bboxes
+    transformed_bbox = transform_bbox(bbox1, "xycwh -> xycwh")
+    assert allclose(transformed_bbox, bbox1)
+
+    transformed_bbox = transform_bbox(bbox2, "xyxy -> xywh")
+    expected_bbox = tensor([[90.0, 70.0, 70.0, 90.0], [40.0, 40.0, 50.0, 80.0]])
+    assert allclose(transformed_bbox, expected_bbox)
+
+
+def test_transform_bbox_xyxy_to_Any(dummy_bboxes):
+    bbox1, bbox2 = dummy_bboxes
+    transformed_bbox = transform_bbox(bbox1, "xyxy -> xyxy")
+    assert allclose(transformed_bbox, bbox1)
+
+    transformed_bbox = transform_bbox(bbox2, "xyxy -> xycwh")
+    expected_bbox = tensor([[125.0, 115.0, 70.0, 90.0], [65.0, 80.0, 50.0, 80.0]])
+    assert allclose(transformed_bbox, expected_bbox)
+
+
+def test_transform_bbox_invalid_format(dummy_bboxes):
+    bbox, _ = dummy_bboxes
+
+    # Test invalid input format
+    with pytest.raises(ValueError, match="Invalid input or output format"):
+        transform_bbox(bbox, "invalid->xyxy")
+
+    # Test invalid output format
+    with pytest.raises(ValueError, match="Invalid input or output format"):
+        transform_bbox(bbox, "xywh->invalid")
+
+
+def test_generate_anchors():
+    image_size = [256, 256]
+    strides = [8, 16, 32]
+    anchors, scalers = generate_anchors(image_size, strides)
+    assert anchors.shape[0] == scalers.shape[0]
+    assert anchors.shape[1] == 2
+
+
+def test_vec2box_autoanchor():
+    with initialize(config_path="../../yolo/config", version_base=None):
+        cfg: Config = compose(config_name="config", overrides=["model=v9-m"])
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = create_model(cfg.model, weight_path=None).to(device)
+    vec2box = Vec2Box(model, cfg.model.anchor, cfg.image_size, device)
+    assert vec2box.strides == [8, 16, 32]
+
+    vec2box.update((320, 640))
+    assert vec2box.anchor_grid.shape == (4200, 2)
+    assert vec2box.scaler.shape == tuple([4200])
+
+
+def test_anc2box_autoanchor(inference_v7_cfg: Config):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = create_model(inference_v7_cfg.model, weight_path=None).to(device)
+    anchor_cfg: AnchorConfig = inference_v7_cfg.model.anchor.copy()
+    del anchor_cfg.strides
+    anc2box = Anc2Box(model, anchor_cfg, inference_v7_cfg.image_size, device)
+    assert anc2box.strides == [8, 16, 32]
+
+    anc2box.update((320, 640))
+    anchor_grids_shape = [anchor_grid.shape for anchor_grid in anc2box.anchor_grids]
+    assert anchor_grids_shape == [
+        torch.Size([1, 1, 80, 80, 2]),
+        torch.Size([1, 1, 40, 40, 2]),
+        torch.Size([1, 1, 20, 20, 2]),
+    ]
+    assert anc2box.anchor_scale.shape == torch.Size([3, 1, 3, 1, 1, 2])
+
+
+def test_bbox_nms():
+    cls_dist = tensor(
+        [[[0.1, 0.7, 0.2], [0.6, 0.3, 0.1]], [[0.4, 0.4, 0.2], [0.5, 0.4, 0.1]]]  # Example class distribution
+    )
+    bbox = tensor(
+        [[[50, 50, 100, 100], [60, 60, 110, 110]], [[40, 40, 90, 90], [70, 70, 120, 120]]],  # Example bounding boxes
+        dtype=float32,
+    )
+    nms_cfg = NMSConfig(min_confidence=0.5, min_iou=0.5)
+
+    expected_output = [
+        tensor(
+            [
+                [1.0000, 50.0000, 50.0000, 100.0000, 100.0000, 0.6682],
+                [0.0000, 60.0000, 60.0000, 110.0000, 110.0000, 0.6457],
+            ]
+        )
+    ]
+
+    output = bbox_nms(cls_dist, bbox, nms_cfg)
+
+    for out, exp in zip(output, expected_output):
+        assert allclose(out, exp, atol=1e-4), f"Output: {out} Expected: {exp}"
+
+
+def test_calculate_map():
+    predictions = tensor([[0, 60, 60, 160, 160, 0.5], [0, 40, 40, 120, 120, 0.5]])  # [class, x1, y1, x2, y2]
+    ground_truths = tensor([[0, 50, 50, 150, 150], [0, 30, 30, 100, 100]])  # [class, x1, y1, x2, y2]
+
+    mAP = calculate_map(predictions, ground_truths)
+
+    expected_ap50 = tensor(0.5)
+    expected_ap50_95 = tensor(0.2)
+
+    assert isclose(mAP["mAP.5"], expected_ap50, atol=1e-5), f"AP50 mismatch"
+    assert isclose(mAP["mAP.5:.95"], expected_ap50_95, atol=1e-5), f"Mean AP mismatch"

tests/test_utils/test_module_utils.py

@@ -0,0 +1,73 @@
+import re
+import sys
+from pathlib import Path
+
+import pytest
+from torch import nn
+
+project_root = Path(__file__).resolve().parent.parent.parent
+sys.path.append(str(project_root))
+from yolo.utils.module_utils import (
+    auto_pad,
+    create_activation_function,
+    divide_into_chunks,
+)
+
+
+@pytest.mark.parametrize(
+    "kernel_size, dilation, expected",
+    [
+        (3, 1, (1, 1)),
+        ((3, 3), (1, 1), (1, 1)),
+        (3, (2, 2), (2, 2)),
+        ((5, 5), 1, (2, 2)),
+        ((3, 5), (2, 1), (2, 2)),
+    ],
+)
+def test_auto_pad(kernel_size, dilation, expected):
+    assert auto_pad(kernel_size, dilation) == expected, "auto_pad does not calculate padding correctly"
+
+
+@pytest.mark.parametrize(
+    "activation_name, expected_type",
+    [("ReLU", nn.ReLU), ("leakyrelu", nn.LeakyReLU), ("none", nn.Identity), (None, nn.Identity), (False, nn.Identity)],
+)
+def test_get_activation(activation_name, expected_type):
+    result = create_activation_function(activation_name)
+    assert isinstance(result, expected_type), f"get_activation does not return correct type for {activation_name}"
+
+
+def test_get_activation_invalid():
+    with pytest.raises(ValueError):
+        create_activation_function("unsupported_activation")
+
+
+def test_divide_into_chunks():
+    input_list = [0, 1, 2, 3, 4, 5]
+    chunk_num = 2
+    expected_output = [[0, 1, 2], [3, 4, 5]]
+    assert divide_into_chunks(input_list, chunk_num) == expected_output
+
+
+def test_divide_into_chunks_non_divisible_length():
+    input_list = [0, 1, 2, 3, 4, 5]
+    chunk_num = 4
+    with pytest.raises(
+        ValueError,
+        match=re.escape("The length of the input list (6) must be exactly divisible by the number of chunks (4)."),
+    ):
+        divide_into_chunks(input_list, chunk_num)
+
+
+def test_divide_into_chunks_single_chunk():
+    input_list = [0, 1, 2, 3, 4, 5]
+    chunk_num = 1
+    expected_output = [[0, 1, 2, 3, 4, 5]]
+    assert divide_into_chunks(input_list, chunk_num) == expected_output
+
+
+def test_divide_into_chunks_equal_chunks():
+    input_list = [0, 1, 2, 3, 4, 5, 6, 7]
+    chunk_num = 4
+    expected_output = [[0, 1], [2, 3], [4, 5], [6, 7]]
+    assert divide_into_chunks(input_list, chunk_num) == expected_output

yolo/__init__.py

@@ -3,7 +3,7 @@ from yolo.model.yolo import create_model
 from yolo.tools.data_loader import AugmentationComposer, create_dataloader
 from yolo.tools.drawer import draw_bboxes
 from yolo.tools.solver import ModelTester, ModelTrainer, ModelValidator
-from yolo.utils.bounding_box_utils import Vec2Box, bbox_nms
+from yolo.utils.bounding_box_utils import Anc2Box, Vec2Box, bbox_nms, create_converter
 from yolo.utils.deploy_utils import FastModelLoader
 from yolo.utils.logging_utils import custom_logger
 from yolo.utils.model_utils import PostProccess
@@ -16,7 +16,9 @@ all = [
     "validate_log_directory",
     "draw_bboxes",
     "Vec2Box",
+    "Anc2Box",
     "bbox_nms",
+    "create_converter",
     "AugmentationComposer",
     "create_dataloader",
     "FastModelLoader",