yolonnx 0.1.1__tar.gz

yolonnx-0.1.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Kasper Fromm Pedersen
+
+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.

yolonnx-0.1.1/PKG-INFO

@@ -0,0 +1,66 @@
+Metadata-Version: 2.1
+Name: yolonnx
+Version: 0.1.1
+Summary: You Only Look ONNX
+Keywords: ONNX,YOLOv8,onnxruntime,vision
+Home-page: https://www.cs.aau.dk/
+Author-Email: Kasper Fromm Pedersen <kasperf@cs.aau.dk>
+License: MIT
+Project-URL: Homepage, https://www.cs.aau.dk/
+Project-URL: Repository, https://github.com/fromm1990/yolonnx
+Requires-Python: >=3.10
+Requires-Dist: numpy==1.26.*
+Requires-Dist: aau-label==0.2.*
+Description-Content-Type: text/markdown
+
+# You Only Look ONNX
+This repository is a light weight library to ease the use of ONNX models exported by the Ultralytics YOLOv8 framework.
+
+## Example Detector Usage
+```python
+from pathlib import Path
+
+from onnxruntime import InferenceSession
+from PIL import Image
+
+from yolonnx.services import Detector
+from yolonnx.to_tensor_strategies import PillowToTensorContainStrategy
+
+model = Path("path/to/file.onnx")
+session = InferenceSession(
+    model.as_posix(),
+    providers=[
+        "CUDAExecutionProvider",
+        "CPUExecutionProvider",
+    ],
+)
+
+predictor = Detector(session, PillowToTensorContainStrategy())
+img = Image.open("path/to/image.jpg")
+print(predictor.run(img))
+```
+
+## Example Classifier Usage
+```python
+from pathlib import Path
+
+from onnxruntime import InferenceSession
+from PIL import Image
+
+from yolonnx.services import Classifier
+from yolonnx.to_tensor_strategies import PillowToTensorContainStrategy
+
+model = Path("path/to/file.onnx")
+session = InferenceSession(
+    model.as_posix(),
+    providers=[
+        "CUDAExecutionProvider",
+        "CPUExecutionProvider",
+    ],
+)
+
+predictor = Classifier(session, PillowToTensorContainStrategy())
+img = Image.open("path/to/image.jpg")
+print(predictor.run(img))
+
+```

yolonnx-0.1.1/README.md

@@ -0,0 +1,51 @@
+# You Only Look ONNX
+This repository is a light weight library to ease the use of ONNX models exported by the Ultralytics YOLOv8 framework.
+
+## Example Detector Usage
+```python
+from pathlib import Path
+
+from onnxruntime import InferenceSession
+from PIL import Image
+
+from yolonnx.services import Detector
+from yolonnx.to_tensor_strategies import PillowToTensorContainStrategy
+
+model = Path("path/to/file.onnx")
+session = InferenceSession(
+    model.as_posix(),
+    providers=[
+        "CUDAExecutionProvider",
+        "CPUExecutionProvider",
+    ],
+)
+
+predictor = Detector(session, PillowToTensorContainStrategy())
+img = Image.open("path/to/image.jpg")
+print(predictor.run(img))
+```
+
+## Example Classifier Usage
+```python
+from pathlib import Path
+
+from onnxruntime import InferenceSession
+from PIL import Image
+
+from yolonnx.services import Classifier
+from yolonnx.to_tensor_strategies import PillowToTensorContainStrategy
+
+model = Path("path/to/file.onnx")
+session = InferenceSession(
+    model.as_posix(),
+    providers=[
+        "CUDAExecutionProvider",
+        "CPUExecutionProvider",
+    ],
+)
+
+predictor = Classifier(session, PillowToTensorContainStrategy())
+img = Image.open("path/to/image.jpg")
+print(predictor.run(img))
+
+```

yolonnx-0.1.1/pyproject.toml

@@ -0,0 +1,40 @@
+[project]
+name = "yolonnx"
+version = "0.1.1"
+description = "You Only Look ONNX"
+authors = [
+    { name = "Kasper Fromm Pedersen", email = "kasperf@cs.aau.dk" },
+]
+dependencies = [
+    "numpy==1.26.*",
+    "aau-label==0.2.*",
+]
+requires-python = ">=3.10"
+readme = "README.md"
+keywords = [
+    "ONNX",
+    "YOLOv8",
+    "onnxruntime",
+    "vision",
+]
+
+[project.license]
+text = "MIT"
+
+[project.urls]
+homepage = "https://www.cs.aau.dk/"
+repository = "https://github.com/fromm1990/yolonnx"
+
+[build-system]
+requires = [
+    "pdm-backend",
+]
+build-backend = "pdm.backend"
+
+[tool.pdm]
+distribution = true
+
+[tool.pdm.dev-dependencies]
+test = [
+    "onnxruntime==1.17.*",
+]

yolonnx-0.1.1/src/yolonnx/model.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from aau_label.model import AAULabel
+from numpy.typing import NDArray
+
+
+@dataclass
+class Size:
+    width: float
+    height: float
+
+
+@dataclass
+class ImgTensor:
+    scale: Size
+    data: NDArray
+
+
+@dataclass
+class ClassifierResult:
+    name: str
+    score: float
+
+
+@dataclass
+class DetectorResult(ClassifierResult, AAULabel): ...

yolonnx-0.1.1/src/yolonnx/protocols.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from typing import Any, Callable, Protocol, TypeVar, runtime_checkable
+
+from numpy import int64
+from numpy.typing import NDArray
+
+from .model import ImgTensor
+
+T = TypeVar("T", contravariant=True)
+
+
+class PNodeArg(Protocol):
+    @property
+    def name(self) -> str: ...
+
+    @property
+    def shape(self) -> Any: ...
+
+    @property
+    def type(self) -> str: ...
+
+
+@runtime_checkable
+class SparseTensorProtocol(Protocol):
+    def dense_shape(self) -> NDArray[int64]: ...
+    def values(self) -> NDArray: ...
+    def device_name(self) -> str: ...
+
+
+class ModelMetadataProtocol(Protocol):
+    @property
+    def custom_metadata_map(self) -> dict: ...
+
+
+class InferenceSessionProtocol(Protocol):
+    def run(
+        self, output_names, input_feed: dict[str, Any], run_options=None
+    ) -> list[NDArray] | list[list] | list[dict] | list[SparseTensorProtocol]: ...
+
+    def run_async(
+        self,
+        output_names,
+        input_feed: dict[str, Any],
+        callback: Callable,
+        user_data: Any,
+        run_options=None,
+    ): ...
+
+    def get_inputs(self) -> list[PNodeArg]: ...
+
+    def get_modelmeta(self) -> ModelMetadataProtocol: ...
+
+
+class ToTensorStrategyProtocol(Protocol[T]):
+    def __call__(self, img: T, tensor_width: int, tensor_height: int) -> ImgTensor: ...

yolonnx-0.1.1/src/yolonnx/services/__init__.py

@@ -0,0 +1,4 @@
+from .classifier import Classifier
+from .detector import Detector
+
+__all__ = ["Classifier", "Detector"]

yolonnx-0.1.1/src/yolonnx/services/classifier.py

@@ -0,0 +1,71 @@
+import ast
+import logging
+from typing import Generic, Sequence, TypeVar
+
+import numpy
+from numpy import float32
+from numpy.typing import NDArray
+
+from ..model import ClassifierResult
+from ..protocols import (
+    InferenceSessionProtocol,
+    SparseTensorProtocol,
+    ToTensorStrategyProtocol,
+)
+
+T = TypeVar("T")
+logger = logging.getLogger(__name__)
+
+
+class Classifier(Generic[T]):
+    def __init__(
+        self,
+        session: InferenceSessionProtocol,
+        to_tensor_strategy: ToTensorStrategyProtocol[T],
+        threshold: float = 0.1,
+    ) -> None:
+        self.__session = session
+        self.__to_tensor_strategy = to_tensor_strategy
+        self.__threshold = threshold
+        meta = self.__session.get_modelmeta()
+        self.__names: dict[int, str] = ast.literal_eval(
+            meta.custom_metadata_map["names"]
+        )
+
+    @property
+    def threshold(self) -> float:
+        return self.__threshold
+
+    @property
+    def shape(self) -> tuple[int, int]:
+        return self.__session.get_inputs()[0].shape[2:]
+
+    @property
+    def names(self) -> dict[int, str]:
+        return self.__names
+
+    def warmup(self) -> None:
+        tensor = numpy.zeros((1, 3, self.shape[0], self.shape[1]), float32)
+        self.__session.run(None, {"images": tensor})
+
+    def __result_handler(self, results: NDArray) -> Sequence[ClassifierResult]:
+        labels_idx = numpy.argwhere(results > self.threshold)
+        scores: NDArray[float32] = results[labels_idx]
+
+        rv: list[ClassifierResult] = []
+        for i in range(len(labels_idx)):
+            id = labels_idx[i][0]
+            label_name = self.names[id]
+            rv.append(ClassifierResult(name=label_name, score=scores[i]))
+
+        rv.sort(key=lambda x: x.score, reverse=True)
+        return rv
+
+    def run(self, img: T) -> Sequence[ClassifierResult]:
+        tensor = self.__to_tensor_strategy(img, *self.shape)
+        results = self.__session.run(None, {"images": tensor.data})[0]
+
+        if isinstance(results, SparseTensorProtocol):
+            results = results.values()
+
+        return self.__result_handler(results[0])

yolonnx-0.1.1/src/yolonnx/services/detector.py

@@ -0,0 +1,94 @@
+import ast
+import logging
+from typing import Generic, Sequence, TypeVar, cast
+
+import numpy
+from numpy.typing import NDArray
+
+from .. import utils
+from ..model import DetectorResult, ImgTensor
+from ..protocols import InferenceSessionProtocol, ToTensorStrategyProtocol
+
+T = TypeVar("T")
+logger = logging.getLogger(__name__)
+
+
+class Detector(Generic[T]):
+    def __init__(
+        self,
+        session: InferenceSessionProtocol,
+        to_tensor_strategy: ToTensorStrategyProtocol[T],
+        conf_threshold: float = 0.25,
+        iou_threshold: float = 0.7,
+    ) -> None:
+        self.__session = session
+        self.__to_tensor_strategy = to_tensor_strategy
+        self.__conf_threshold = conf_threshold
+        self.__iou_threshold = iou_threshold
+
+        meta = self.__session.get_modelmeta()
+        self.__names: dict[int, str] = ast.literal_eval(
+            meta.custom_metadata_map["names"]
+        )
+
+    @property
+    def conf_threshold(self) -> float:
+        return self.__conf_threshold
+
+    @property
+    def iou_threshold(self) -> float:
+        return self.__iou_threshold
+
+    @property
+    def shape(self) -> tuple[int, int]:
+        return self.__session.get_inputs()[0].shape[2:]
+
+    @property
+    def names(self) -> dict[int, str]:
+        return self.__names
+
+    def __result_handler(
+        self, results: NDArray, tensor: ImgTensor
+    ) -> Sequence[DetectorResult]:
+        predictions = numpy.squeeze(results[0]).T
+
+        scores = numpy.max(predictions[:, 4:], axis=1)
+        keep = scores > self.__conf_threshold
+        predictions = predictions[keep, :]
+        scores = scores[keep]
+        class_ids = numpy.argmax(predictions[:, 4:], axis=1)
+
+        boxes = predictions[:, :4]
+        # Make x0, y0 left upper corner instead of box center
+        boxes[:, 0:2] -= boxes[:, 2:4] / 2
+        boxes /= numpy.array(
+            [
+                tensor.scale.width,
+                tensor.scale.height,
+                tensor.scale.width,
+                tensor.scale.height,
+            ],
+            dtype=numpy.float32,
+        )
+        boxes = boxes.astype(numpy.int32)
+
+        keep = utils.nms(boxes, scores, self.__iou_threshold)
+        rv = []
+        for bbox, label, score in zip(boxes[keep], class_ids[keep], scores[keep]):
+            rv.append(
+                DetectorResult(
+                    x=bbox[0].item(),
+                    y=bbox[1].item(),
+                    width=bbox[2].item(),
+                    height=bbox[3].item(),
+                    name=self.__names[label],
+                    score=score.item(),
+                )
+            )
+
+        return rv
+
+    def run(self, img: T) -> Sequence[DetectorResult]:
+        tensor = self.__to_tensor_strategy(img, *self.shape)
+        results = cast(list[NDArray], self.__session.run(None, {"images": tensor.data}))
+        return self.__result_handler(results[0], tensor)

yolonnx-0.1.1/src/yolonnx/to_tensor_strategies.py

@@ -0,0 +1,44 @@
+import numpy
+from PIL import ImageOps
+from PIL.Image import Image as ImageType
+from PIL.Image import Resampling
+
+from .model import ImgTensor, Size
+from .protocols import ToTensorStrategyProtocol
+
+
+class PillowToTensorContainStrategy(ToTensorStrategyProtocol[ImageType]):
+    """
+    Takes a Pillow image and scales it such it fits inside the tensor. The image aspect ratio is preserved.
+
+    The image is then padded with the color (114, 114, 114) to make the image fit inside the tensor.
+    """
+
+    def __call__(
+        self, img: ImageType, tensor_width: int, tensor_height: int
+    ) -> ImgTensor:
+        tensor_dims = (tensor_width, tensor_height)
+        original_width, original_height = img.size
+
+        img = ImageOps.contain(img, tensor_dims, Resampling.BILINEAR)
+        new_width, new_height = img.size
+
+        img = ImageOps.pad(
+            img,
+            tensor_dims,
+            Resampling.BILINEAR,
+            (114, 114, 114),
+            (0, 0),
+        )
+        data = numpy.array(img)
+
+        data = data / 255.0
+        data = data.transpose(2, 0, 1)
+        tensor = data[numpy.newaxis, :, :, :].astype(numpy.float32)
+
+        scale = Size(
+            width=new_width / original_width,
+            height=new_height / original_height,
+        )
+
+        return ImgTensor(scale, tensor)

yolonnx-0.1.1/src/yolonnx/utils.py

@@ -0,0 +1,48 @@
+import numpy
+from numpy import float32, float64, int32, int64
+from numpy.typing import NDArray
+
+
+def compute_iou(box: NDArray[int32], boxes: NDArray[int32]) -> NDArray[float64]:
+    # Compute xmin, ymin, xmax, ymax for both boxes
+    xmin = numpy.minimum(box[0], boxes[:, 0])
+    ymin = numpy.minimum(box[1], boxes[:, 1])
+    xmax = numpy.maximum(box[0] + box[2], boxes[:, 0] + boxes[:, 2])
+    ymax = numpy.maximum(box[1] + box[3], boxes[:, 1] + boxes[:, 3])
+
+    # Compute intersection area
+    intersection_area = numpy.maximum(0, xmax - xmin) * numpy.maximum(0, ymax - ymin)
+
+    # Compute union area
+    box_area = box[2] * box[3]
+    boxes_area = boxes[:, 2] * boxes[:, 3]
+    union_area = box_area + boxes_area - intersection_area
+
+    # Compute IoU
+    iou = intersection_area / union_area
+
+    return iou
+
+
+def nms(
+    boxes: NDArray[int32], scores: NDArray[float32], iou_threshold: float
+) -> list[int64]:
+    # Sort by score
+    sorted_indices = numpy.argsort(scores)[::-1]
+
+    keep_boxes = []
+    while sorted_indices.size > 0:
+        # Pick the last box
+        box_id = sorted_indices[0]
+        keep_boxes.append(box_id)
+
+        # Compute IoU of the picked box with the rest
+        ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :])
+
+        # Remove boxes with IoU over the threshold
+        keep_indices = numpy.where(ious < iou_threshold)[0]
+
+        # print(keep_indices.shape, sorted_indices.shape)
+        sorted_indices = sorted_indices[keep_indices + 1]
+
+    return keep_boxes