"""YOLOX-family detection adapter over the generic AI inference API."""
from collections.abc import Sequence
from typing import TYPE_CHECKING, Union
import numpy as np
from olo.ai.processing import image_to_tensor, letterbox, nms
from olo.ai.types import Detection
from olo.core.images import Image
from olo.vision import BoundingBox
if TYPE_CHECKING:
from olo.ai import Model
#: COCO class labels used by the pretrained YOLOX models.
COCO_LABELS: tuple[str, ...] = (
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck",
"boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench",
"bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra",
"giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
"skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
"fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
"broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
"remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
"refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
"hair drier", "toothbrush",
)
_DEFAULT_STRIDES = (8, 16, 32)
[docs]
def decode_yolox(
raw: np.ndarray,
input_size: tuple[int, int],
*,
strides: Sequence[int] = _DEFAULT_STRIDES,
) -> np.ndarray:
"""Decode raw YOLOX output rows into (N, 5+C) pixel-space predictions.
YOLOX exports emit grid-relative rows (cx, cy, log w, log h, objectness,
class scores...). This applies the per-stride grid offsets and exp scaling,
yielding center-based boxes in input-image pixels.
"""
if raw.ndim == 3:
raw = raw[0]
height, width = input_size
grids = []
expanded_strides = []
for stride in strides:
grid_h, grid_w = height // stride, width // stride
xv, yv = np.meshgrid(np.arange(grid_w), np.arange(grid_h))
grid = np.stack((xv, yv), axis=2).reshape(-1, 2)
grids.append(grid)
expanded_strides.append(np.full((grid.shape[0], 1), stride))
grid = np.concatenate(grids, axis=0)
expanded_stride = np.concatenate(expanded_strides, axis=0)
if raw.shape[0] != grid.shape[0]:
raise ValueError(
f"YOLOX output has {raw.shape[0]} rows but input size {input_size} "
f"with strides {tuple(strides)} implies {grid.shape[0]}"
)
decoded = raw.copy()
decoded[:, :2] = (raw[:, :2] + grid) * expanded_stride
decoded[:, 2:4] = np.exp(raw[:, 2:4]) * expanded_stride
return decoded
[docs]
class YoloDetector:
"""Runs a YOLOX ONNX model through a loaded Model and decodes detections."""
[docs]
def __init__(
self,
model: "Model",
*,
labels: Sequence[str] = COCO_LABELS,
input_size: tuple[int, int] | None = None,
strides: Sequence[int] = _DEFAULT_STRIDES,
) -> None:
"""Bind to a loaded model; input size is read from the model when static."""
self._model = model
self._labels = tuple(labels)
self._strides = tuple(strides)
info = model.info
if len(info.inputs) != 1 or len(info.outputs) != 1:
raise ValueError(
f"Expected a single-input single-output YOLOX model, "
f"got {len(info.inputs)} inputs / {len(info.outputs)} outputs"
)
self._input_name = info.inputs[0].name
self._output_name = info.outputs[0].name
if input_size is None:
shape = info.inputs[0].shape
if len(shape) != 4 or shape[2] < 0 or shape[3] < 0:
raise ValueError(
f"Model input shape {shape} is dynamic; pass input_size=(height, width)"
)
input_size = (int(shape[2]), int(shape[3]))
self._input_size = input_size
@property
def input_size(self) -> tuple[int, int]:
"""Model input size as (height, width)."""
return self._input_size
[docs]
def detect(
self,
image: Union[Image, np.ndarray],
*,
conf: float = 0.25,
iou: float = 0.45,
) -> list[Detection]:
"""Detect objects in a BGR image and return score-sorted detections."""
array = image.array if isinstance(image, Image) else np.asarray(image)
if array.ndim != 3:
raise ValueError(f"Expected an HWC BGR image, got shape {array.shape}")
source_h, source_w = array.shape[:2]
boxed = letterbox(array, self._input_size)
# YOLOX takes raw 0-255 BGR values; no mean/std normalization.
tensor = image_to_tensor(boxed.array, layout="NCHW", channel_order="BGR")
outputs = self._model.infer({self._input_name: tensor})
decoded = decode_yolox(outputs[self._output_name], self._input_size, strides=self._strides)
class_scores = decoded[:, 5:]
class_ids = class_scores.argmax(axis=1)
scores = decoded[:, 4] * class_scores[np.arange(len(class_ids)), class_ids]
mask = scores >= conf
if not mask.any():
return []
decoded, scores, class_ids = decoded[mask], scores[mask], class_ids[mask]
# cxcywh in letterboxed pixels -> xyxy in source-image pixels.
cx = (decoded[:, 0] - boxed.pad_x) / boxed.scale
cy = (decoded[:, 1] - boxed.pad_y) / boxed.scale
half_w = decoded[:, 2] / (2.0 * boxed.scale)
half_h = decoded[:, 3] / (2.0 * boxed.scale)
boxes = np.stack(
(
np.clip(cx - half_w, 0, source_w),
np.clip(cy - half_h, 0, source_h),
np.clip(cx + half_w, 0, source_w),
np.clip(cy + half_h, 0, source_h),
),
axis=1,
)
detections: list[Detection] = []
for index in nms(boxes, scores, iou):
x_min, y_min, x_max, y_max = boxes[index]
if x_max <= x_min or y_max <= y_min:
continue
class_id = int(class_ids[index])
label = self._labels[class_id] if class_id < len(self._labels) else str(class_id)
detections.append(
Detection(
label=label,
class_id=class_id,
score=float(scores[index]),
box=BoundingBox.from_xyxy(float(x_min), float(y_min), float(x_max), float(y_max)),
)
)
return detections
__all__ = ["COCO_LABELS", "YoloDetector", "decode_yolox"]