Feat/add yolov8 detector

.github/workflows/tests.yml

@@ -0,0 +1,42 @@
+# .github/workflows/test.yml
+name: Test
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  test:
+    name: Test
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version:
+          - "3.10"
+          - "3.11"
+          - "3.12"
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv and set the python version
+        uses: astral-sh/setup-uv@v5
+        with:
+          version: "0.5.23"
+          enable-cache: true
+          cache-dependency-glob: "uv.lock"
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install the project
+        run: uv sync --all-groups --no-group dev-model
+
+      - name: Checking linter and formatting
+        run: uvx ruff check
+
+      - name: Run tests
+        run: uv run pytest tests -v
+
+      - name: Test with Coverage
+        run: uv run pytest --cov=src tests/

.gitignore

@@ -8,3 +8,8 @@ wheels/
 
 # Virtual environments
 .venv
+*.onnx
+*.pt
+*.DS_Store
+.bak/
+*.coverage

.pre-commit-config.yaml

@@ -12,7 +12,7 @@ repos:
       - id: trailing-whitespace
   # python code formatting
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: v0.7.3
+    rev: v0.9.2
     hooks:
       - id: ruff
         types_or: [python, pyi, jupyter]
@@ -25,4 +25,4 @@ repos:
     rev: v1.5.5
     hooks:
       - id: forbid-crlf
-      - id: remove-crlf
+      - id: remove-crlf

Makefile

@@ -0,0 +1,11 @@
+yolo-export-onnx:
+	yolo export \
+	model=color_correction_asdfghjkl/asset/.model/yv8-det.pt \
+	format=onnx \
+	device=mps \
+	simplify=True \
+	dynamic=False \
+	half=True
+
+test:
+	pytest tests -v

README.md

@@ -40,7 +40,5 @@ This package is designed to perform color correction on images using the Color C
 - [Colour Science Python](https://www.colour-science.org/colour-checker-detection/)
 - [Fast and Robust Multiple ColorChecker Detection ()](https://github.com/pedrodiamel/colorchecker-detection)
 - [Automatic color correction with OpenCV and Python (PyImageSearch)](https://pyimagesearch.com/2021/02/15/automatic-color-correction-with-opencv-and-python/)
-
----
-
-Happy Color Correcting! 🌟
+- [ONNX-YOLOv8-Object-Detection](https://github.com/ibaiGorordo/ONNX-YOLOv8-Object-Detection)
+- [yolov8-triton](https://github.com/omarabid59/yolov8-triton/tree/main)

pyproject.toml

@@ -8,6 +8,7 @@ authors = [
 ]
 requires-python = ">=3.10"
 dependencies = [
+    "httpx>=0.28.1",
     "pydantic>=2.10.5",
     "pydantic-settings>=2.7.1",
 ]
@@ -24,15 +25,20 @@ analyze = [
     "matplotlib>=3.10.0",
 ]
 yolodet = [
+    "onnx>=1.17.0",
     "onnxruntime>=1.20.1",
 ]
 mccdet = [
     "opencv-contrib-python>=4.11.0.86",
 ]
 dev = [
+    "pytest-cov==6.0.0",
     "pytest>=8.3.4",
     "ruff>=0.9.2",
 ]
+dev-model = [
+    "ultralytics>=8.3.65",
+]
 
 # ---- ruff ----
 [tool.ruff]
@@ -129,4 +135,4 @@ unfixable = ["F401"]
 [tool.ruff.format]
 indent-style = "space"
 quote-style = "double"
-line-ending = "lf"
+line-ending = "lf"

src/color_correction_asdfghjkl/constant/yolov8_det.py

@@ -0,0 +1,4 @@
+class_names = [
+    "patch",
+    "card",
+]

src/color_correction_asdfghjkl/core/card_detection/base.py

@@ -0,0 +1,14 @@
+from abc import ABC, abstractmethod
+
+import numpy as np
+
+from color_correction_asdfghjkl.schemas.yolov8_det import DetectionResult
+
+
+class BaseCardDetector(ABC):
+    @abstractmethod
+    def detect(
+        self,
+        image: np.ndarray,
+        conf: float = 0.15,
+    ) -> DetectionResult: ...

src/color_correction_asdfghjkl/core/card_detection/yolov8_det_onnx.py

@@ -0,0 +1,194 @@
+import time
+
+import cv2
+import numpy as np
+import onnxruntime
+
+from color_correction_asdfghjkl.core.card_detection.base import BaseCardDetector
+from color_correction_asdfghjkl.schemas.yolov8_det import DetectionResult
+from color_correction_asdfghjkl.utils.downloader import downloader_model_yolov8
+from color_correction_asdfghjkl.utils.yolo_utils import (
+    multiclass_nms,
+    xywh2xyxy,
+)
+
+
+class YOLOv8CardDetector(BaseCardDetector):
+    """YOLOv8CardDetector is a class that implements card detection
+    using the YOLOv8 model.
+
+    Reference
+    ---------
+    https://github.com/ibaiGorordo/ONNX-YOLOv8-Object-Detection/blob/main/yolov8/YOLOv8.py
+
+    """
+
+    def __init__(
+        self,
+        conf_th: float = 0.15,
+        iou_th: float = 0.7,
+        path: str | None = None,
+        use_gpu: bool = False,
+    ) -> None:
+        self.conf_threshold = conf_th
+        self.iou_threshold = iou_th
+        self.use_gpu = use_gpu
+        if path is None:
+            print("Auto downloading YOLOv8 model...")
+            path = downloader_model_yolov8(use_gpu)
+        self.__initialize_model(path)
+
+    def detect(self, image: np.ndarray) -> DetectionResult:
+        """
+        Detect objects in the given image using YOLOv8 model.
+
+        Parameters
+        ----------
+        image : np.ndarray
+            The input image BGR in which to detect objects.
+
+        Returns
+        -------
+        DetectionResult
+            A dataclass containing detected bounding boxes, confidence scores,
+            and class IDs.
+        """
+        input_tensor = self.__prepare_input(image)
+        outputs = self.__inference(input_tensor)
+        boxes, scores, class_ids = self.__process_output(outputs)
+
+        det_res = DetectionResult(
+            boxes=boxes,
+            scores=scores,
+            class_ids=class_ids,
+        )
+
+        return det_res
+
+    # Service functions
+    def __initialize_model(self, path: str) -> None:
+        self.session = onnxruntime.InferenceSession(
+            path,
+            providers=onnxruntime.get_available_providers(),
+        )
+        # Get model info
+        self.__get_input_details()
+        self.__get_output_details()
+
+    def __prepare_input(self, original_image: np.ndarray) -> np.ndarray:
+        [height, width, _] = original_image.shape
+
+        # expected shape based on model input
+        expected_width = self.input_width
+        expected_height = self.input_height
+        expected_length = min((expected_height, expected_width))
+
+        length = max((height, width))
+        # self.scale_to_expected = expected_length / length
+        self.scale_to_ori = length / expected_length
+
+        image = np.zeros((length, length, 3), np.uint8)
+        image[0:height, 0:width] = original_image
+
+        input_image = cv2.resize(image, (expected_width, expected_height))
+
+        if self.use_gpu:
+            input_image = (input_image / 255.0).astype(np.float16)
+        else:
+            input_image = (input_image / 255.0).astype(np.float32)
+        # Channel first
+        input_image = input_image.transpose(2, 0, 1)
+
+        # Expand dimensions
+        input_image = np.expand_dims(input_image, axis=0)
+        return input_image
+
+    def __inference(self, input_tensor: np.ndarray) -> list[np.ndarray]:
+        start = time.perf_counter()  # noqa: F841
+        outputs = self.session.run(
+            self.output_names,
+            {self.input_names[0]: input_tensor},
+        )
+
+        # print(f"Inference time: {(time.perf_counter() - start) * 1000:.2f} ms")
+        return outputs
+
+    def __process_output(
+        self,
+        output: list[np.ndarray],
+    ) -> tuple[list[list[int]], list[float], list[int]]:
+        predictions = np.squeeze(output[0]).T
+
+        # Filter out object confidence scores below threshold
+        scores = np.max(predictions[:, 4:], axis=1)
+        predictions = predictions[scores > self.conf_threshold, :]
+        scores = scores[scores > self.conf_threshold]
+
+        if len(scores) == 0:
+            return [], [], []
+
+        # Get the class with the highest confidence
+        class_ids = np.argmax(predictions[:, 4:], axis=1)
+
+        # Get bounding boxes for each object
+        boxes = self.__extract_boxes(predictions)
+
+        # Apply non-maxima suppression to suppress weak, overlapping bounding boxes
+        # indices = nms(boxes, scores, self.iou_threshold)
+        indices = multiclass_nms(
+            boxes,
+            scores,
+            class_ids,
+            self.iou_threshold,
+        )
+
+        return (
+            boxes[indices].astype(int).tolist(),
+            scores[indices].tolist(),
+            class_ids[indices].tolist(),
+        )
+
+    # Helper functions
+    def __extract_boxes(self, predictions: np.ndarray) -> np.ndarray:
+        # Extract boxes from predictions
+        boxes = predictions[:, :4]
+
+        # Scale boxes to original image dimensions
+        boxes = self.__rescale_boxes(boxes)
+
+        # Convert boxes to xyxy format
+        boxes = xywh2xyxy(boxes)
+
+        return boxes
+
+    def __rescale_boxes(self, boxes: np.ndarray) -> np.ndarray:
+        # Rescale boxes to original image dimensions
+        boxes *= self.scale_to_ori
+        return boxes
+
+    def __get_input_details(self) -> None:
+        model_inputs = self.session.get_inputs()
+        self.input_names = [model_inputs[i].name for i in range(len(model_inputs))]
+
+        self.input_shape = model_inputs[0].shape
+        self.input_height = self.input_shape[2]
+        self.input_width = self.input_shape[3]
+
+    def __get_output_details(self) -> None:
+        model_outputs = self.session.get_outputs()
+        self.output_names = [model_outputs[i].name for i in range(len(model_outputs))]
+
+
+if __name__ == "__main__":
+    print("YOLOv8CardDetector")
+    model_path = "color_correction_asdfghjkl/asset/.model/yv8-det.onnx"
+    image_path = "color_correction_asdfghjkl/asset/images/cc-1.jpg"
+    image_path = "color_correction_asdfghjkl/asset/images/Test 19.png"
+    detector = YOLOv8CardDetector(conf_th=0.15, iou_th=0.7, use_gpu=True)
+
+    input_image = cv2.imread(image_path)
+    # input_image = cv2.resize(input_image, (640, 640))
+    result = detector.detect(input_image)
+    result.print_summary()
+    image_drawed = result.draw_detections(input_image)
+    cv2.imwrite("result.png", image_drawed)

src/color_correction_asdfghjkl/schemas/device.py

@@ -0,0 +1,30 @@
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+
+class GPUType(str, Enum):
+    NVIDIA = "NVIDIA"
+    AMD = "AMD"
+    APPLE = "Apple Integrated"
+    UNKNOWN = "Unknown GPU"
+
+
+class CPUArchitecture(str, Enum):
+    INTEL = "Intel"
+    AMD = "AMD"
+    ARM = "ARM"
+    APPLE = "Apple Silicon"
+    UNKNOWN = "Unknown"
+
+
+class DeviceSpecs(BaseModel):
+    """Device specifications schema."""
+
+    os_name: str = Field(..., description="Operating system name")
+    cpu_arch: CPUArchitecture = Field(
+        CPUArchitecture.UNKNOWN,
+        description="CPU architecture",
+    )
+    gpu_type: GPUType = Field(GPUType.UNKNOWN, description="GPU type")
+    is_apple_silicon: bool = Field(False, description="Whether device is Apple Silicon")