diff --git a/.github/workflows/detector-export.yml b/.github/workflows/detector-export.yml index d252b02d..350a706d 100644 --- a/.github/workflows/detector-export.yml +++ b/.github/workflows/detector-export.yml @@ -4,21 +4,29 @@ on: pull_request: paths: - ".github/workflows/detector-export.yml" + - "complexity/deploy/onnx_detector/**" - "complexity/generative/detection/**" - "scripts/check_onnx_parity.py" + - "scripts/onnx_detect.py" - "scripts/export_onnx.py" - "scripts/export_tensorrt.py" - "tests/test_detector_export.py" + - "tests/test_onnx_detect_cli.py" + - "tests/test_onnx_detector_*.py" - "pyproject.toml" push: branches: [main] paths: - ".github/workflows/detector-export.yml" + - "complexity/deploy/onnx_detector/**" - "complexity/generative/detection/**" - "scripts/check_onnx_parity.py" + - "scripts/onnx_detect.py" - "scripts/export_onnx.py" - "scripts/export_tensorrt.py" - "tests/test_detector_export.py" + - "tests/test_onnx_detect_cli.py" + - "tests/test_onnx_detector_*.py" - "pyproject.toml" permissions: @@ -42,10 +50,24 @@ jobs: - name: Lint detector export code run: >- ruff check + complexity/deploy/onnx_detector complexity/generative/detection/exporting.py scripts/check_onnx_parity.py + scripts/onnx_detect.py scripts/export_onnx.py scripts/export_tensorrt.py tests/test_detector_export.py + tests/test_onnx_detect_cli.py + tests/test_onnx_detector_core.py + tests/test_onnx_detector_metadata.py + tests/test_onnx_detector_pipeline.py + tests/test_onnx_detector_skeleton.py - name: Test ONNX branches and dynamic batch parity - run: pytest -q tests/test_detector_export.py + run: >- + pytest -q + tests/test_detector_export.py + tests/test_onnx_detect_cli.py + tests/test_onnx_detector_core.py + tests/test_onnx_detector_metadata.py + tests/test_onnx_detector_pipeline.py + tests/test_onnx_detector_skeleton.py diff --git a/complexity/deploy/__init__.py b/complexity/deploy/__init__.py new file mode 100644 index 00000000..a63acf62 --- /dev/null +++ b/complexity/deploy/__init__.py @@ -0,0 +1,4 @@ +"""Deployment helpers for exported Complexity models.""" + +__all__ = ["onnx_detector"] + diff --git a/complexity/deploy/onnx_detector/__init__.py b/complexity/deploy/onnx_detector/__init__.py new file mode 100644 index 00000000..7d918199 --- /dev/null +++ b/complexity/deploy/onnx_detector/__init__.py @@ -0,0 +1,31 @@ +"""ONNX Runtime deployment helpers for TR-Hash Vision detectors.""" + +from .metadata import ( + DEFAULT_CONFIDENCE_THRESHOLD, + DEFAULT_IOU_THRESHOLD, + DEFAULT_MAX_DETECTIONS, + BranchType, + OnnxDetectorMetadata, +) +from .pipeline import OnnxDetectorPipeline +from .preprocess import ImageGeometry, PreprocessResult, preprocess_image, restore_boxes +from .session import OnnxDetectorSession, OrtSessionConfig +from .types import Detection, DetectionResult, TimingBreakdown + +__all__ = [ + "BranchType", + "DEFAULT_CONFIDENCE_THRESHOLD", + "DEFAULT_IOU_THRESHOLD", + "DEFAULT_MAX_DETECTIONS", + "Detection", + "DetectionResult", + "ImageGeometry", + "OnnxDetectorMetadata", + "OnnxDetectorPipeline", + "OnnxDetectorSession", + "OrtSessionConfig", + "PreprocessResult", + "TimingBreakdown", + "preprocess_image", + "restore_boxes", +] diff --git a/complexity/deploy/onnx_detector/dfl.py b/complexity/deploy/onnx_detector/dfl.py new file mode 100644 index 00000000..9de8da16 --- /dev/null +++ b/complexity/deploy/onnx_detector/dfl.py @@ -0,0 +1,49 @@ +"""DFL box decode for TR-Hash Vision v8 ONNX detector exports.""" + +from __future__ import annotations + +import numpy as np + +from .grid import GridGeometry +from .metadata import OnnxDetectorMetadata + + +def decode_dfl_boxes( + regression_logits: np.ndarray, + metadata: OnnxDetectorMetadata, + geometry: GridGeometry, +) -> np.ndarray: + """Decode raw LTRB DFL logits into input-pixel xyxy boxes.""" + + logits = np.asarray(regression_logits, dtype=np.float32) + if logits.ndim not in {2, 3}: + raise ValueError("regression_logits must have shape [N, C] or [B, N, C]") + if logits.shape[-1] != metadata.regression_width: + raise ValueError( + "regression_logits last dimension does not match metadata.regression_width" + ) + if logits.shape[-2] != geometry.centers_xy.shape[0]: + raise ValueError("regression_logits cell count does not match grid geometry") + + bins = metadata.dfl_bins + distances = logits.reshape(*logits.shape[:-1], 4, bins) + if metadata.reg_max: + shifted = distances - distances.max(axis=-1, keepdims=True) + probabilities = np.exp(shifted) + probabilities /= probabilities.sum(axis=-1, keepdims=True) + bucket_indices = np.arange(bins, dtype=np.float32) + distances = (probabilities * bucket_indices).sum(axis=-1) + else: + # DFL disabled mirrors PyTorch's single-bin regression width contract. + distances = np.log1p(np.exp(distances[..., 0])) + + distances_px = distances * geometry.strides.reshape(1, -1, 1) + centers = geometry.centers_xy.reshape(1, -1, 2) + top_left = centers - distances_px[..., (0, 1)] + bottom_right = centers + distances_px[..., (2, 3)] + boxes = np.concatenate((top_left, bottom_right), axis=-1) + boxes = np.clip(boxes, 0.0, float(metadata.image_size)).astype(np.float32) + + if logits.ndim == 2: + return boxes[0] + return boxes diff --git a/complexity/deploy/onnx_detector/grid.py b/complexity/deploy/onnx_detector/grid.py new file mode 100644 index 00000000..779f0159 --- /dev/null +++ b/complexity/deploy/onnx_detector/grid.py @@ -0,0 +1,47 @@ +"""Anchor-center generation for TR-Hash Vision v8 ONNX detector exports.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Sequence + +import numpy as np + + +@dataclass(frozen=True) +class GridGeometry: + """Concatenated row-major feature-grid geometry.""" + + centers_xy: np.ndarray + strides: np.ndarray + grid_sizes: tuple[int, ...] + + +def generate_grid_geometry(image_size: int, grid_sizes: Sequence[int]) -> GridGeometry: + """Generate concatenated anchor centers and strides for each grid level.""" + + if image_size <= 0: + raise ValueError("image_size must be positive") + + resolved_grid_sizes = tuple(int(grid) for grid in grid_sizes) + if not resolved_grid_sizes: + raise ValueError("grid_sizes must not be empty") + if any(grid <= 0 for grid in resolved_grid_sizes): + raise ValueError("grid_sizes must contain only positive values") + + centers: list[np.ndarray] = [] + strides: list[np.ndarray] = [] + for grid in resolved_grid_sizes: + stride = float(image_size) / float(grid) + coords = (np.arange(grid, dtype=np.float32) + 0.5) * np.float32(stride) + x_centers, y_centers = np.meshgrid(coords, coords, indexing="xy") + centers.append( + np.stack((x_centers.reshape(-1), y_centers.reshape(-1)), axis=1) + ) + strides.append(np.full((grid * grid,), stride, dtype=np.float32)) + + return GridGeometry( + centers_xy=np.concatenate(centers, axis=0).astype(np.float32, copy=False), + strides=np.concatenate(strides, axis=0).astype(np.float32, copy=False), + grid_sizes=resolved_grid_sizes, + ) diff --git a/complexity/deploy/onnx_detector/metadata.py b/complexity/deploy/onnx_detector/metadata.py new file mode 100644 index 00000000..0851bed9 --- /dev/null +++ b/complexity/deploy/onnx_detector/metadata.py @@ -0,0 +1,185 @@ +"""Metadata sidecar schema for TR-Hash Vision v8 ONNX detector exports.""" + +from __future__ import annotations + +import json +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Literal, Mapping, cast + +BranchType = Literal["nms-free", "o2m"] + +DEFAULT_CONFIDENCE_THRESHOLD = 0.25 +DEFAULT_IOU_THRESHOLD = 0.45 +DEFAULT_MAX_DETECTIONS = 300 + + +@dataclass(frozen=True) +class OnnxDetectorMetadata: + """Validated deployment metadata loaded from an ONNX sidecar JSON file.""" + + architecture_version: int + image_size: int + num_classes: int + num_cells: int + regression_width: int + reg_max: int + scale_factors: tuple[int, ...] + grid_sizes: tuple[int, ...] + p2_head: bool + branch: BranchType + requires_nms: bool + output_semantics: str + confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD + iou_threshold: float = DEFAULT_IOU_THRESHOLD + max_detections: int = DEFAULT_MAX_DETECTIONS + class_names: tuple[str, ...] | None = None + + @property + def dfl_bins(self) -> int: + return self.reg_max + 1 if self.reg_max else 1 + + @property + def prediction_width(self) -> int: + return self.regression_width + self.num_classes + + @property + def strides(self) -> tuple[float, ...]: + return tuple(float(self.image_size) / float(grid) for grid in self.grid_sizes) + + def as_dict(self) -> dict[str, object]: + return { + "architecture_version": self.architecture_version, + "image_size": self.image_size, + "num_classes": self.num_classes, + "num_cells": self.num_cells, + "regression_width": self.regression_width, + "reg_max": self.reg_max, + "scale_factors": list(self.scale_factors), + "grid_sizes": list(self.grid_sizes), + "p2_head": self.p2_head, + "branch": self.branch, + "requires_nms": self.requires_nms, + "output_semantics": self.output_semantics, + "confidence_threshold": self.confidence_threshold, + "iou_threshold": self.iou_threshold, + "max_detections": self.max_detections, + "class_names": list(self.class_names) if self.class_names is not None else None, + } + + +def load_metadata(path: str | Path) -> OnnxDetectorMetadata: + """Load and validate an ONNX detector metadata sidecar.""" + + metadata_path = Path(path) + return metadata_from_mapping(json.loads(metadata_path.read_text())) + + +def metadata_from_mapping(values: Mapping[str, object]) -> OnnxDetectorMetadata: + """Create validated metadata from a mapping.""" + + branch = _branch(values.get("branch")) + requires_nms = bool(values.get("requires_nms", branch == "o2m")) + if requires_nms is not (branch == "o2m"): + raise ValueError("requires_nms must be true only for the o2m branch") + + class_names_value = values.get("class_names") + class_names = None + if class_names_value is not None: + if not isinstance(class_names_value, Sequence) or isinstance( + class_names_value, (str, bytes) + ): + raise ValueError("class_names must be a sequence of strings") + class_names = tuple(str(name) for name in class_names_value) + + metadata = OnnxDetectorMetadata( + architecture_version=_positive_int(values, "architecture_version"), + image_size=_positive_int(values, "image_size"), + num_classes=_positive_int(values, "num_classes"), + num_cells=_positive_int(values, "num_cells"), + regression_width=_positive_int(values, "regression_width"), + reg_max=_non_negative_int(values, "reg_max"), + scale_factors=_positive_int_tuple(values, "scale_factors"), + grid_sizes=_positive_int_tuple(values, "grid_sizes"), + p2_head=bool(values.get("p2_head", False)), + branch=branch, + requires_nms=requires_nms, + output_semantics=str(values.get("output_semantics", "")), + confidence_threshold=float( + values.get("confidence_threshold", DEFAULT_CONFIDENCE_THRESHOLD) + ), + iou_threshold=float(values.get("iou_threshold", DEFAULT_IOU_THRESHOLD)), + max_detections=int(values.get("max_detections", DEFAULT_MAX_DETECTIONS)), + class_names=class_names, + ) + _validate_metadata(metadata) + return metadata + + +def validate_output_shape( + metadata: OnnxDetectorMetadata, + output_shape: Sequence[int | str | None], +) -> None: + """Validate an ONNX output shape against the sidecar metadata.""" + + if len(output_shape) != 3: + raise ValueError(f"ONNX output must be rank 3, got shape {tuple(output_shape)}") + _validate_axis(output_shape[1], metadata.num_cells, "num_cells") + _validate_axis(output_shape[2], metadata.prediction_width, "prediction_width") + + +def _validate_metadata(metadata: OnnxDetectorMetadata) -> None: + if metadata.architecture_version != 8: + raise ValueError("only TR-Hash detector architecture v8 metadata is supported") + if metadata.num_cells != sum(grid * grid for grid in metadata.grid_sizes): + raise ValueError("num_cells must equal sum(grid ** 2 for grid_sizes)") + if metadata.regression_width != 4 * metadata.dfl_bins: + raise ValueError("regression_width must equal 4 * DFL bin count") + if metadata.max_detections <= 0: + raise ValueError("max_detections must be positive") + if not 0.0 <= metadata.confidence_threshold <= 1.0: + raise ValueError("confidence_threshold must be in [0, 1]") + if not 0.0 <= metadata.iou_threshold <= 1.0: + raise ValueError("iou_threshold must be in [0, 1]") + if ( + metadata.class_names is not None + and len(metadata.class_names) != metadata.num_classes + ): + raise ValueError("class_names length must match num_classes") + + +def _branch(value: object) -> BranchType: + normalized = "nms-free" if value == "nms_free" else value + if normalized not in {"nms-free", "o2m"}: + raise ValueError(f"unsupported ONNX detector branch: {value!r}") + return cast(BranchType, normalized) + + +def _positive_int(values: Mapping[str, object], key: str) -> int: + value = int(values[key]) + if value <= 0: + raise ValueError(f"{key} must be positive") + return value + + +def _non_negative_int(values: Mapping[str, object], key: str) -> int: + value = int(values[key]) + if value < 0: + raise ValueError(f"{key} must be non-negative") + return value + + +def _positive_int_tuple(values: Mapping[str, object], key: str) -> tuple[int, ...]: + raw = values[key] + if not isinstance(raw, Sequence) or isinstance(raw, (str, bytes)): + raise ValueError(f"{key} must be a sequence of positive integers") + resolved = tuple(int(item) for item in raw) + if not resolved or any(item <= 0 for item in resolved): + raise ValueError(f"{key} must contain positive integers") + return resolved + + +def _validate_axis(axis: int | str | None, expected: int, name: str) -> None: + if isinstance(axis, int) and axis != expected: + raise ValueError(f"ONNX output {name} mismatch: expected {expected}, got {axis}") diff --git a/complexity/deploy/onnx_detector/pipeline.py b/complexity/deploy/onnx_detector/pipeline.py new file mode 100644 index 00000000..d94e3887 --- /dev/null +++ b/complexity/deploy/onnx_detector/pipeline.py @@ -0,0 +1,148 @@ +"""High-level preprocessing, inference, decode, and postprocess pipeline.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from time import perf_counter +from typing import Sequence + +import numpy as np + +from . import postprocess +from .dfl import decode_dfl_boxes +from .grid import GridGeometry, generate_grid_geometry +from .metadata import OnnxDetectorMetadata, load_metadata, validate_output_shape +from .preprocess import preprocess_image, restore_boxes +from .session import OnnxDetectorSession, OrtSessionConfig +from .types import Detection, DetectionResult, TimingBreakdown + + +@dataclass +class OnnxDetectorPipeline: + """Deployment pipeline configured by ONNX model and metadata sidecar.""" + + metadata: OnnxDetectorMetadata + session: OnnxDetectorSession + geometry: GridGeometry | None = None + + def __post_init__(self) -> None: + if self.geometry is None: + self.geometry = generate_grid_geometry( + self.metadata.image_size, + self.metadata.grid_sizes, + ) + + @classmethod + def from_files( + cls, + model_path: str | Path, + metadata_path: str | Path, + providers: Sequence[str] = ("CPUExecutionProvider",), + ) -> "OnnxDetectorPipeline": + metadata = load_metadata(metadata_path) + session = cls.create_session(model_path, providers).open() + validate_output_shape(metadata, session._require_session().get_outputs()[0].shape) + if session.config.warmup_runs: + session.warmup((1, 3, metadata.image_size, metadata.image_size)) + return cls(metadata=metadata, session=session) + + def predict(self, image: object) -> DetectionResult: + preprocess_start = perf_counter() + preprocessed = preprocess_image(image, self.metadata.image_size) + preprocess_ms = (perf_counter() - preprocess_start) * 1000.0 + + inference_start = perf_counter() + predictions = self.session.run(preprocessed.pixel_values) + inference_ms = (perf_counter() - inference_start) * 1000.0 + + postprocess_start = perf_counter() + detections = self._decode_and_postprocess(predictions, preprocessed.geometry) + postprocess_ms = (perf_counter() - postprocess_start) * 1000.0 + + return DetectionResult( + detections=tuple(detections), + timing=TimingBreakdown(preprocess_ms, inference_ms, postprocess_ms), + provider_used=self.session.provider_used, + branch_type=self.metadata.branch, + metadata=self.metadata.as_dict(), + ) + + def postprocess_single_image( + self, + boxes: np.ndarray, + scores: np.ndarray, + classes: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + filtered_boxes, filtered_scores, filtered_classes, _ = ( + postprocess.filter_by_confidence( + boxes, + scores, + classes, + self.metadata.confidence_threshold, + ) + ) + if self.metadata.branch == "o2m": + keep = postprocess.class_aware_nms( + filtered_boxes, + filtered_scores, + filtered_classes, + self.metadata.iou_threshold, + self.metadata.max_detections, + ) + return filtered_boxes[keep], filtered_scores[keep], filtered_classes[keep] + + order = np.argsort(filtered_scores)[::-1] + keep = order[: self.metadata.max_detections] + return filtered_boxes[keep], filtered_scores[keep], filtered_classes[keep] + + @staticmethod + def create_session( + model_path: str | Path, + providers: Sequence[str] = ("CPUExecutionProvider",), + ) -> OnnxDetectorSession: + return OnnxDetectorSession( + OrtSessionConfig(Path(model_path), providers=tuple(providers)) + ) + + def _decode_and_postprocess( + self, + predictions: np.ndarray, + image_geometry: object, + ) -> list[Detection]: + if predictions.ndim != 3 or predictions.shape[0] != 1: + raise ValueError("pipeline.predict currently expects ONNX output shape [1, N, C]") + validate_output_shape(self.metadata, predictions.shape) + + regression = predictions[..., : self.metadata.regression_width] + class_logits = predictions[..., self.metadata.regression_width :] + boxes_input = decode_dfl_boxes(regression, self.metadata, self.geometry)[0] + class_scores = _sigmoid(class_logits[0]) + scores = class_scores.max(axis=-1) + classes = class_scores.argmax(axis=-1).astype(np.int64, copy=False) + + kept_boxes, kept_scores, kept_classes = self.postprocess_single_image( + boxes_input, + scores, + classes, + ) + boxes_pixel = restore_boxes(kept_boxes, image_geometry) + boxes_norm = np.clip(kept_boxes / float(self.metadata.image_size), 0.0, 1.0) + return [ + Detection( + box_norm=tuple(float(value) for value in norm_box), + box_pixel=tuple(float(value) for value in pixel_box), + class_id=int(class_id), + score=float(score), + ) + for norm_box, pixel_box, class_id, score in zip( + boxes_norm, + boxes_pixel, + kept_classes, + kept_scores, + ) + ] + + +def _sigmoid(values: np.ndarray) -> np.ndarray: + return 1.0 / (1.0 + np.exp(-values)) diff --git a/complexity/deploy/onnx_detector/postprocess.py b/complexity/deploy/onnx_detector/postprocess.py new file mode 100644 index 00000000..f32f03db --- /dev/null +++ b/complexity/deploy/onnx_detector/postprocess.py @@ -0,0 +1,108 @@ +"""Confidence filtering and class-aware NMS for ONNX detector outputs.""" + +from __future__ import annotations + +import numpy as np + + +def filter_by_confidence( + boxes: np.ndarray, + scores: np.ndarray, + classes: np.ndarray, + conf_threshold: float, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Return candidate detections whose score meets the inclusive threshold.""" + + resolved_boxes = np.asarray(boxes, dtype=np.float32) + resolved_scores = np.asarray(scores, dtype=np.float32) + resolved_classes = np.asarray(classes) + + if resolved_boxes.ndim != 2 or resolved_boxes.shape[1] != 4: + raise ValueError("boxes must have shape [N, 4]") + if resolved_scores.shape != (resolved_boxes.shape[0],): + raise ValueError("scores must have shape [N]") + if resolved_classes.shape != (resolved_boxes.shape[0],): + raise ValueError("classes must have shape [N]") + + indices = np.nonzero(resolved_scores >= conf_threshold)[0] + return ( + resolved_boxes[indices], + resolved_scores[indices], + resolved_classes[indices], + indices, + ) + + +def class_aware_nms( + boxes: np.ndarray, + scores: np.ndarray, + classes: np.ndarray, + iou_threshold: float, + max_detections: int, +) -> np.ndarray: + """Run class-aware greedy NMS for the O2M branch.""" + + if max_detections <= 0: + raise ValueError("max_detections must be positive") + + resolved_boxes = np.asarray(boxes, dtype=np.float32) + resolved_scores = np.asarray(scores, dtype=np.float32) + resolved_classes = np.asarray(classes) + if resolved_boxes.ndim != 2 or resolved_boxes.shape[1] != 4: + raise ValueError("boxes must have shape [N, 4]") + if resolved_scores.shape != (resolved_boxes.shape[0],): + raise ValueError("scores must have shape [N]") + if resolved_classes.shape != (resolved_boxes.shape[0],): + raise ValueError("classes must have shape [N]") + if resolved_boxes.shape[0] == 0: + return np.empty((0,), dtype=np.int64) + + kept: list[np.ndarray] = [] + for class_id in np.unique(resolved_classes): + class_indices = np.nonzero(resolved_classes == class_id)[0] + local_keep = _greedy_nms( + resolved_boxes[class_indices], + resolved_scores[class_indices], + iou_threshold, + ) + kept.append(class_indices[local_keep]) + + keep = np.concatenate(kept).astype(np.int64, copy=False) + order = np.argsort(resolved_scores[keep])[::-1] + return keep[order[:max_detections]] + + +def _greedy_nms( + boxes: np.ndarray, + scores: np.ndarray, + iou_threshold: float, +) -> np.ndarray: + order = np.argsort(scores)[::-1] + kept: list[int] = [] + + while order.size: + current = int(order[0]) + kept.append(current) + if order.size == 1: + break + + rest = order[1:] + ious = _box_iou(boxes[current], boxes[rest]) + order = rest[ious <= iou_threshold] + + return np.asarray(kept, dtype=np.int64) + + +def _box_iou(box: np.ndarray, boxes: np.ndarray) -> np.ndarray: + x1 = np.maximum(box[0], boxes[:, 0]) + y1 = np.maximum(box[1], boxes[:, 1]) + x2 = np.minimum(box[2], boxes[:, 2]) + y2 = np.minimum(box[3], boxes[:, 3]) + + intersection = np.maximum(x2 - x1, 0.0) * np.maximum(y2 - y1, 0.0) + box_area = np.maximum(box[2] - box[0], 0.0) * np.maximum(box[3] - box[1], 0.0) + boxes_area = np.maximum(boxes[:, 2] - boxes[:, 0], 0.0) * np.maximum( + boxes[:, 3] - boxes[:, 1], 0.0 + ) + union = box_area + boxes_area - intersection + return intersection / np.maximum(union, 1e-7) diff --git a/complexity/deploy/onnx_detector/preprocess.py b/complexity/deploy/onnx_detector/preprocess.py new file mode 100644 index 00000000..c13d088f --- /dev/null +++ b/complexity/deploy/onnx_detector/preprocess.py @@ -0,0 +1,79 @@ +"""Image preprocessing and box restoration for ONNX detector inference.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np + + +@dataclass(frozen=True) +class ImageGeometry: + """Geometry required to undo square letterbox preprocessing.""" + + original_width: int + original_height: int + image_size: int + scale: float + left: int + top: int + + +@dataclass(frozen=True) +class PreprocessResult: + """Preprocessed model input and the geometry needed for restoration.""" + + pixel_values: np.ndarray + geometry: ImageGeometry + + +def preprocess_image(image: Any, image_size: int) -> PreprocessResult: + """Mirror PyTorch RGB letterbox preprocessing for detector exports.""" + + from PIL import Image + + if isinstance(image, (str, Path)): + pil_image = Image.open(image) + elif isinstance(image, Image.Image): + pil_image = image + else: + pil_image = Image.fromarray(np.asarray(image)) + + pil_image = pil_image.convert("RGB") + original_width, original_height = pil_image.size + scale = min(image_size / original_width, image_size / original_height) + resized_width = max(1, round(original_width * scale)) + resized_height = max(1, round(original_height * scale)) + left = (image_size - resized_width) // 2 + top = (image_size - resized_height) // 2 + + resized = pil_image.resize((resized_width, resized_height), Image.Resampling.BILINEAR) + canvas = Image.new("RGB", (image_size, image_size), (114, 114, 114)) + canvas.paste(resized, (left, top)) + + pixels = np.asarray(canvas, dtype=np.float32).transpose(2, 0, 1) / 255.0 + pixels = (pixels - 0.5) / 0.5 + pixel_values = np.expand_dims(pixels, axis=0).astype(np.float32, copy=False) + geometry = ImageGeometry( + original_width=original_width, + original_height=original_height, + image_size=image_size, + scale=scale, + left=left, + top=top, + ) + return PreprocessResult(pixel_values=pixel_values, geometry=geometry) + + +def restore_boxes(boxes: np.ndarray, geometry: ImageGeometry) -> np.ndarray: + """Map input-square pixel xyxy boxes back to source-image pixel xyxy boxes.""" + + restored = np.asarray(boxes, dtype=np.float32).copy() + restored[:, (0, 2)] = (restored[:, (0, 2)] - geometry.left) / geometry.scale + restored[:, (1, 3)] = (restored[:, (1, 3)] - geometry.top) / geometry.scale + restored[:, (0, 2)] = np.clip(restored[:, (0, 2)], 0.0, geometry.original_width) + restored[:, (1, 3)] = np.clip(restored[:, (1, 3)], 0.0, geometry.original_height) + return restored + diff --git a/complexity/deploy/onnx_detector/session.py b/complexity/deploy/onnx_detector/session.py new file mode 100644 index 00000000..803acfb5 --- /dev/null +++ b/complexity/deploy/onnx_detector/session.py @@ -0,0 +1,106 @@ +"""ONNX Runtime session wrapper for detector exports.""" + +from __future__ import annotations + +import os +import site +from dataclasses import dataclass +from pathlib import Path +from typing import Sequence + +import numpy as np + + +@dataclass(frozen=True) +class OrtSessionConfig: + """Configuration for constructing an ONNX Runtime detector session.""" + + model_path: Path | str + providers: tuple[str, ...] = ("CPUExecutionProvider",) + warmup_runs: int = 1 + + +class OnnxDetectorSession: + """Thin wrapper around an ONNX Runtime inference session.""" + + def __init__(self, config: OrtSessionConfig, session: object | None = None) -> None: + self.config = config + self._session = session + self._dll_directory_handles: list[object] = [] + + def open(self) -> "OnnxDetectorSession": + if self._session is None: + import onnxruntime as ort + + if _needs_cuda_dlls(self.config.providers) and hasattr(ort, "preload_dlls"): + # Prefer NVIDIA site-package DLLs over an unrelated PyTorch CUDA build. + ort.preload_dlls(directory="") + self._add_tensorrt_dll_directories() + self._session = ort.InferenceSession( + str(self.config.model_path), + providers=list(self.config.providers), + ) + return self + + @property + def provider_used(self) -> str: + providers = self._require_session().get_providers() + return providers[0] if providers else "" + + @property + def input_name(self) -> str: + return self._require_session().get_inputs()[0].name + + @property + def output_name(self) -> str: + return self._require_session().get_outputs()[0].name + + def warmup(self, input_shape: Sequence[int]) -> None: + dummy = np.zeros(tuple(input_shape), dtype=np.float32) + for _ in range(self.config.warmup_runs): + self.run(dummy) + + def run(self, pixel_values: np.ndarray) -> np.ndarray: + output = self._require_session().run( + [self.output_name], + {self.input_name: pixel_values}, + )[0] + return np.asarray(output, dtype=np.float32) + + def _add_tensorrt_dll_directories(self) -> None: + if not _needs_tensorrt_dlls(self.config.providers): + return + for site_packages in _candidate_site_packages(): + tensorrt_libs = site_packages / "tensorrt_libs" + if not tensorrt_libs.is_dir(): + continue + if hasattr(os, "add_dll_directory"): + self._dll_directory_handles.append( + os.add_dll_directory(str(tensorrt_libs)) + ) + os.environ["PATH"] = str(tensorrt_libs) + os.pathsep + os.environ.get( + "PATH", "" + ) + + def _require_session(self): + if self._session is None: + raise RuntimeError("ONNX Runtime session is not open") + return self._session + + +def _needs_cuda_dlls(providers: Sequence[str]) -> bool: + cuda_providers = {"CUDAExecutionProvider", "TensorrtExecutionProvider"} + return any(provider in cuda_providers for provider in providers) + + +def _needs_tensorrt_dlls(providers: Sequence[str]) -> bool: + return any(provider == "TensorrtExecutionProvider" for provider in providers) + + +def _candidate_site_packages() -> tuple[Path, ...]: + candidates: list[Path] = [] + for raw_path in [*site.getsitepackages(), site.getusersitepackages()]: + path = Path(raw_path) + if path not in candidates: + candidates.append(path) + return tuple(candidates) diff --git a/complexity/deploy/onnx_detector/types.py b/complexity/deploy/onnx_detector/types.py new file mode 100644 index 00000000..1a767266 --- /dev/null +++ b/complexity/deploy/onnx_detector/types.py @@ -0,0 +1,45 @@ +"""Stable output schema for ONNX detector inference.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field + +Box = tuple[float, float, float, float] + + +@dataclass(frozen=True) +class Detection: + """One decoded object detection.""" + + box_norm: Box + box_pixel: Box + class_id: int + score: float + + +@dataclass(frozen=True) +class TimingBreakdown: + """Wall-clock timing captured around the ONNX inference pipeline.""" + + preprocess_ms: float = 0.0 + inference_ms: float = 0.0 + postprocess_ms: float = 0.0 + + def as_dict(self) -> dict[str, float]: + return { + "preprocess_ms": self.preprocess_ms, + "inference_ms": self.inference_ms, + "postprocess_ms": self.postprocess_ms, + } + + +@dataclass(frozen=True) +class DetectionResult: + """Stable result returned by the ONNX detector pipeline.""" + + detections: Sequence[Detection] = field(default_factory=tuple) + timing: TimingBreakdown = field(default_factory=TimingBreakdown) + provider_used: str = "" + branch_type: str = "" + metadata: Mapping[str, object] = field(default_factory=dict) diff --git a/docs/index.md b/docs/index.md index 144013a1..f4b19ac6 100644 --- a/docs/index.md +++ b/docs/index.md @@ -90,6 +90,7 @@ GQA + TR-MoE architecture. - [TR-Hash image-text-to-text](tr-hash-image-text-to-text.md) - [TR-Hash text-to-image](tr-hash-text-to-image.md) - [TR-Hash object detection and serving](tr-hash-object-detection.md) +- [TR-Hash Vision ONNX deployment](onnx_deploy.md) - [Detector specialization and ablations](TR_HASH_DETECTOR_SPECIALIZATION.md) - [TR-Hash sensor fusion](tr_hash_sensor_fusion.md) - [Vision dependency stack](vision-dependency-stack.md) diff --git a/docs/onnx_deploy.md b/docs/onnx_deploy.md new file mode 100644 index 00000000..1017c989 --- /dev/null +++ b/docs/onnx_deploy.md @@ -0,0 +1,105 @@ +# TR-Hash Vision v8 ONNX Deployment + +This guide runs exported TR-Hash Vision v8 detector models with ONNX Runtime. +The ONNX file returns raw predictions; the deployment pipeline performs the +same preprocessing, DFL decode, confidence filtering, and branch-specific +postprocessing documented in [the contract](onnx_deploy_contract.md). + +## Installation + +Install the framework with its ONNX export/runtime dependencies: + +```bash +python -m pip install -e ".[export]" +``` + +CUDA and TensorRT execution providers still require the matching NVIDIA runtime +libraries for the installed ONNX Runtime build. + +## Export Inputs + +Each deployed model needs two files: + +- the `.onnx` model exported by `scripts/export_onnx.py`; +- the JSON sidecar written next to it by the exporter. + +Do not rename one without passing both paths to the CLI. + +## CPU Example + +```bash +python scripts/onnx_detect.py \ + --model tr_hash_v8_o2m.onnx \ + --metadata tr_hash_v8_o2m.json \ + --image sample.jpg \ + --provider cpu \ + --pretty +``` + +The output is JSON: + +```json +{ + "provider_used": "CPUExecutionProvider", + "branch_type": "o2m", + "timing": { + "preprocess_ms": 0.0, + "inference_ms": 0.0, + "postprocess_ms": 0.0 + }, + "detections": [] +} +``` + +Actual timings and detections depend on the image and hardware. + +## CUDA Example + +```bash +python scripts/onnx_detect.py \ + --model tr_hash_v8_nms_free.onnx \ + --metadata tr_hash_v8_nms_free.json \ + --image sample.jpg \ + --provider cuda \ + --pretty +``` + +`--provider cuda` requests `CUDAExecutionProvider` first and falls back to +`CPUExecutionProvider`. The JSON `provider_used` field reports what ONNX Runtime +actually selected after session creation. + +TensorRT can be requested with: + +```bash +python scripts/onnx_detect.py --model model.onnx --metadata model.json --image sample.jpg --provider tensorrt +``` + +That expands to TensorRT, CUDA, then CPU fallback. + +## Branch Behavior + +- `o2m` exports run confidence filtering followed by class-aware NMS. +- `nms-free` exports run confidence filtering and top-k score selection only. +- `--iou-threshold` only affects `o2m`; the CLI warns if it is passed for an + NMS-free export. +- `--conf-threshold` overrides the default confidence threshold for either + branch. + +The default thresholds match the PyTorch detector path: + +- confidence threshold: `0.25`; +- O2M IoU threshold: `0.45`; +- max detections: `300`. + +## Output Schema + +Each detection contains: + +- `box_norm`: normalized `xyxy` relative to the square model input; +- `box_pixel`: restored `xyxy` in original source-image pixels; +- `class_id`: integer class index; +- `score`: sigmoid quality-class score. + +See [the deployment contract](onnx_deploy_contract.md) for the exact tensor +layout, grid mapping, DFL decode formula, preprocessing, and sidecar validation +rules. diff --git a/docs/onnx_deploy_contract.md b/docs/onnx_deploy_contract.md new file mode 100644 index 00000000..419c5548 --- /dev/null +++ b/docs/onnx_deploy_contract.md @@ -0,0 +1,164 @@ +# TR-Hash Vision v8 ONNX Deployment Contract + +This document records the PyTorch detector contract that deployment runtimes must +preserve when using exported TR-Hash Vision v8 ONNX models. The ONNX graph exports +one raw prediction branch only; preprocessing, decode, confidence filtering, and +optional NMS are runtime responsibilities. + +## Source Of Truth + +- Export wrapper: `complexity/generative/detection/exporting.py` +- Export CLI and sidecar metadata: `scripts/export_onnx.py` +- Detector geometry and decode: `complexity/generative/detection/model.py` +- Class-aware NMS: `complexity/generative/detection/ops.py` +- Image preprocessing and box restoration: `complexity/generative/detection/hub.py` +- Detector shape properties: `complexity/generative/detection/config.py` + +## ONNX Tensor Contract + +- Input name: `pixel_values` +- Input dtype and shape: `float32[B, 3, image_size, image_size]` +- Output name: `predictions` +- Output dtype and shape: `float32[B, num_cells, regression_width + num_classes]` +- Default opset: `17`, unless `scripts/export_onnx.py --opset` overrides it. +- Batch dimension is static `1` unless exported with `--dynamic-batch`. + +For the current v8 640 px COCO exports: + +- `image_size`: `640` +- `num_classes`: `80` +- `reg_max`: `16` +- `dfl_bins`: `17` +- `regression_width`: `68` +- `grid_sizes`: `[160, 80, 40, 20]` +- `num_cells`: `34000` +- `prediction_width`: `148` +- `predictions` shape: `[B, 34000, 148]` + +The first `regression_width` channels are LTRB distribution-focal-loss logits. +The remaining `num_classes` channels are unified quality-class logits. + +## Sidecar Metadata + +The exporter writes a JSON sidecar next to each `.onnx` file with: + +- `architecture_version` +- `image_size` +- `num_classes` +- `num_cells` +- `regression_width` +- `reg_max` +- `scale_factors` +- `grid_sizes` +- `p2_head` +- `branch` +- `requires_nms` +- `output_semantics` + +Deployment code should validate on first inference that: + +- `architecture_version == 8` +- `num_cells == sum(grid ** 2 for grid in grid_sizes)` +- `regression_width == 4 * (reg_max + 1 if reg_max else 1)` +- `prediction_width == regression_width + num_classes` +- ONNX output shape matches `[B, num_cells, prediction_width]` + +The sidecar currently does not encode confidence threshold, IoU threshold, max +detections, class names, preprocessing constants, original-image geometry, input +and output names, or dynamic axes. Those values must be supplied by the runtime +from this contract or future expanded metadata. + +## Preprocessing Contract + +Preprocessing must match `preprocess_detector_image`: + +1. Convert the source image to RGB. +2. Compute `scale = min(image_size / original_width, image_size / original_height)`. +3. Resize with bilinear filtering to `round(original_width * scale)` by + `round(original_height * scale)`, clamped to at least `1` pixel per dimension. +4. Center the resized image on an `image_size x image_size` RGB canvas filled with + `(114, 114, 114)`. +5. Convert to CHW float tensor in `[0, 1]`. +6. Normalize with `(pixel - 0.5) / 0.5`, yielding the training-time `[-1, 1]` + scale. +7. Add batch dimension before ONNX inference. + +The runtime must retain `original_width`, `original_height`, `image_size`, +`scale`, `left`, and `top` so decoded normalized boxes can be restored to source +pixels. + +## Grid And Decode Contract + +Prediction cells are ordered by feature level, then row-major within each level. +For every `grid` in `grid_sizes`: + +- `row = 0..grid-1` +- `col = 0..grid-1` +- normalized center x is `(col + 0.5) / grid` +- normalized center y is `(row + 0.5) / grid` + +For the current 640 px v8 model with `p2_head=true`, grids `[160, 80, 40, 20]` +correspond to input-pixel strides `[4, 8, 16, 32]`. Do not assume a fixed +`[8, 16, 32]` pyramid; derive it from sidecar `grid_sizes`. + +Decode raw predictions as: + +1. `regression = predictions[..., :regression_width]` +2. `class_logits = predictions[..., regression_width:]` +3. If `reg_max > 0`, reshape regression to `[B, N, 4, reg_max + 1]`, softmax over + the final bin axis, and compute the expected value against bins + `0..reg_max`. +4. Divide each expected LTRB distance by that cell's `grid` size to get + normalized distances. +5. Convert to normalized `xyxy` with + `[center_x - left, center_y - top, center_x + right, center_y + bottom]`. +6. Clamp normalized boxes to `[0, 1]`. +7. Compute `class_scores = sigmoid(class_logits)`. +8. For each cell, choose `score, label = max(class_scores, dim=-1)`. + +`reg_max` is the maximum DFL bin index, not the number of bins. With +`reg_max=16`, there are `17` bins and `4 * 17 = 68` regression channels. + +## Branch Postprocessing + +The export wrapper exposes exactly one raw branch: + +- `branch == "nms-free"` returns the one-to-one end-to-end head. +- `branch == "o2m"` returns the one-to-many dense head. +- `branch == "auto"` resolves to `nms-free` when the checkpoint has a one-to-one + head, otherwise `o2m`. + +Both branches use the same decode contract. Their postprocessing differs after +scores and labels are computed: + +- `nms-free`: keep predictions with `score >= confidence_threshold`, then keep + the top `max_detections` by score. PyTorch defaults are + `confidence_threshold=0.25` and `max_detections=300`. +- `o2m`: keep predictions with `score >= confidence_threshold`, then run + class-aware NMS with `iou_threshold=0.45` and `max_detections=300`. + +The score threshold is inclusive. Empty inputs must return empty boxes, scores, +and labels without error. + +## Box Restoration Contract + +Decoded boxes are normalized `xyxy` coordinates in the square letterboxed image. +Map them back to original source pixels as: + +- `x = (x * image_size - left) / scale` +- `y = (y * image_size - top) / scale` +- clamp x coordinates to `[0, original_width]` +- clamp y coordinates to `[0, original_height]` + +## Runtime Defaults To Preserve + +Until these are encoded in sidecar metadata, deployment runtimes should preserve +the PyTorch defaults: + +- confidence threshold: `0.25` +- O2M IoU threshold: `0.45` +- max detections: `300` +- letterbox fill: RGB `(114, 114, 114)` +- resize filter: bilinear +- normalization: `(pixel / 255.0 - 0.5) / 0.5` + diff --git a/pyproject.toml b/pyproject.toml index 9e05ce62..fc1b04e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -111,6 +111,7 @@ serve = [ "python-multipart>=0.0.9", ] export = [ + "Pillow>=10.0.0", "onnx>=1.12.0; platform_system != 'Darwin'", "onnx>=1.12.0,<1.18.0; platform_system == 'Darwin' and python_version < '3.13'", "onnx>=1.20.0; platform_system == 'Darwin' and python_version >= '3.13'", diff --git a/scripts/onnx_detect.py b/scripts/onnx_detect.py new file mode 100644 index 00000000..dd015870 --- /dev/null +++ b/scripts/onnx_detect.py @@ -0,0 +1,110 @@ +"""Run TR-Hash Vision v8 ONNX detector inference on one image. + +Example: + python scripts/onnx_detect.py --model tr_hash_v8_o2m.onnx --metadata tr_hash_v8_o2m.json --image sample.jpg --provider cpu +""" + +from __future__ import annotations + +import argparse +import json +import sys +from dataclasses import replace +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", type=Path, required=True, help="ONNX model path") + parser.add_argument("--metadata", type=Path, required=True, help="JSON sidecar path") + parser.add_argument("--image", type=Path, required=True, help="Input image path") + parser.add_argument( + "--provider", + action="append", + default=None, + help=( + "Provider alias or ORT provider name. May be repeated or comma-separated. " + "Aliases: cpu, cuda, tensorrt." + ), + ) + parser.add_argument("--conf-threshold", type=float, default=None) + parser.add_argument("--iou-threshold", type=float, default=None) + parser.add_argument("--pretty", action="store_true", help="Pretty-print JSON output") + return parser.parse_args() + + +def provider_names(values: list[str] | None) -> tuple[str, ...]: + if not values: + values = ["cpu"] + + resolved: list[str] = [] + for value in values: + for token in value.split(","): + provider = token.strip() + if not provider: + continue + lowered = provider.lower() + if lowered == "cpu": + candidates = ("CPUExecutionProvider",) + elif lowered == "cuda": + candidates = ("CUDAExecutionProvider", "CPUExecutionProvider") + elif lowered == "tensorrt": + candidates = ( + "TensorrtExecutionProvider", + "CUDAExecutionProvider", + "CPUExecutionProvider", + ) + else: + candidates = (provider,) + for candidate in candidates: + if candidate not in resolved: + resolved.append(candidate) + return tuple(resolved) + + +def main() -> None: + args = parse_args() + + from complexity.deploy.onnx_detector import OnnxDetectorPipeline + from complexity.deploy.onnx_detector.metadata import load_metadata + + metadata = load_metadata(args.metadata) + if args.conf_threshold is not None: + metadata = replace(metadata, confidence_threshold=args.conf_threshold) + if args.iou_threshold is not None: + if metadata.branch == "nms-free": + print( + "warning: --iou-threshold is ignored for nms-free exports", + file=sys.stderr, + ) + metadata = replace(metadata, iou_threshold=args.iou_threshold) + + session = OnnxDetectorPipeline.create_session( + args.model, + providers=provider_names(args.provider), + ).open() + pipeline = OnnxDetectorPipeline(metadata=metadata, session=session) + result = pipeline.predict(args.image) + payload = { + "provider_used": result.provider_used, + "branch_type": result.branch_type, + "timing": result.timing.as_dict(), + "detections": [ + { + "box_norm": detection.box_norm, + "box_pixel": detection.box_pixel, + "class_id": detection.class_id, + "score": detection.score, + } + for detection in result.detections + ], + } + print(json.dumps(payload, indent=2 if args.pretty else None)) + + +if __name__ == "__main__": + main() diff --git a/tests/test_onnx_detect_cli.py b/tests/test_onnx_detect_cli.py new file mode 100644 index 00000000..b33e81d0 --- /dev/null +++ b/tests/test_onnx_detect_cli.py @@ -0,0 +1,22 @@ +from scripts.onnx_detect import provider_names + + +def test_provider_aliases_expand_with_fallbacks() -> None: + assert provider_names(None) == ("CPUExecutionProvider",) + assert provider_names(["cuda"]) == ( + "CUDAExecutionProvider", + "CPUExecutionProvider", + ) + assert provider_names(["tensorrt"]) == ( + "TensorrtExecutionProvider", + "CUDAExecutionProvider", + "CPUExecutionProvider", + ) + + +def test_provider_names_accept_repeated_and_explicit_values() -> None: + assert provider_names(["cuda,CPUExecutionProvider", "CustomProvider"]) == ( + "CUDAExecutionProvider", + "CPUExecutionProvider", + "CustomProvider", + ) diff --git a/tests/test_onnx_detector_core.py b/tests/test_onnx_detector_core.py new file mode 100644 index 00000000..5dc1ee71 --- /dev/null +++ b/tests/test_onnx_detector_core.py @@ -0,0 +1,163 @@ +from pathlib import Path + +import numpy as np + +from complexity.deploy.onnx_detector.dfl import decode_dfl_boxes +from complexity.deploy.onnx_detector.grid import generate_grid_geometry +from complexity.deploy.onnx_detector.metadata import BranchType, OnnxDetectorMetadata +from complexity.deploy.onnx_detector.pipeline import OnnxDetectorPipeline +from complexity.deploy.onnx_detector.postprocess import ( + class_aware_nms, + filter_by_confidence, +) +from complexity.deploy.onnx_detector.session import ( + OnnxDetectorSession, + OrtSessionConfig, + _needs_cuda_dlls, +) + + +def _metadata(branch: BranchType = "o2m") -> OnnxDetectorMetadata: + return OnnxDetectorMetadata( + architecture_version=8, + image_size=16, + num_classes=3, + num_cells=16, + regression_width=12, + reg_max=2, + scale_factors=(1,), + grid_sizes=(4,), + p2_head=False, + branch=branch, + requires_nms=branch == "o2m", + output_semantics="raw_ltrb_dfl_and_quality_class_logits", + max_detections=2, + ) + + +def test_grid_mapping_is_row_major_and_counts_real_v8_cells() -> None: + geometry = generate_grid_geometry(16, (4, 2)) + + np.testing.assert_allclose( + geometry.centers_xy[:5], + np.array( + [ + [2.0, 2.0], + [6.0, 2.0], + [10.0, 2.0], + [14.0, 2.0], + [2.0, 6.0], + ], + dtype=np.float32, + ), + ) + np.testing.assert_allclose(geometry.strides[:4], np.full((4,), 4.0)) + np.testing.assert_allclose(geometry.centers_xy[16], np.array([4.0, 4.0])) + + real_geometry = generate_grid_geometry(640, (160, 80, 40, 20)) + assert real_geometry.centers_xy.shape == (34000, 2) + assert real_geometry.strides.shape == (34000,) + + +def test_cuda_dll_preload_is_limited_to_nvidia_providers() -> None: + assert _needs_cuda_dlls(("CUDAExecutionProvider",)) + assert _needs_cuda_dlls(("TensorrtExecutionProvider", "CPUExecutionProvider")) + assert not _needs_cuda_dlls(("CPUExecutionProvider",)) + assert not _needs_cuda_dlls(("CoreMLExecutionProvider",)) + assert not _needs_cuda_dlls(("DmlExecutionProvider",)) + + +def test_dfl_decode_matches_hand_computed_expectation() -> None: + metadata = _metadata() + geometry = generate_grid_geometry(metadata.image_size, metadata.grid_sizes) + logits = np.full((1, metadata.num_cells, metadata.regression_width), -40.0) + # First cell center is (2, 2) and stride is 4. Distances are LTRB + # [0, 0, 1, 1] bins, so the pixel box is [2, 2, 6, 6]. + logits[0, 0, 0] = 40.0 + logits[0, 0, 3] = 40.0 + logits[0, 0, 7] = 40.0 + logits[0, 0, 10] = 40.0 + + boxes = decode_dfl_boxes(logits, metadata, geometry) + + np.testing.assert_allclose(boxes[0, 0], np.array([2.0, 2.0, 6.0, 6.0]), atol=1e-4) + + +def test_confidence_filter_includes_exact_threshold() -> None: + boxes = np.zeros((3, 4), dtype=np.float32) + scores = np.array([0.249, 0.25, 0.251], dtype=np.float32) + classes = np.array([0, 1, 2]) + + _, filtered_scores, filtered_classes, indices = filter_by_confidence( + boxes, + scores, + classes, + conf_threshold=0.25, + ) + + np.testing.assert_array_equal(indices, np.array([1, 2])) + np.testing.assert_allclose(filtered_scores, np.array([0.25, 0.251], dtype=np.float32)) + np.testing.assert_array_equal(filtered_classes, np.array([1, 2])) + + +def test_o2m_nms_suppresses_same_class_only_and_caps_by_score() -> None: + boxes = np.array( + [ + [0.0, 0.0, 10.0, 10.0], + [1.0, 1.0, 9.0, 9.0], + [1.0, 1.0, 9.0, 9.0], + [20.0, 20.0, 30.0, 30.0], + ], + dtype=np.float32, + ) + scores = np.array([0.9, 0.8, 0.7, 0.95], dtype=np.float32) + classes = np.array([1, 1, 2, 1]) + + keep = class_aware_nms( + boxes, + scores, + classes, + iou_threshold=0.5, + max_detections=2, + ) + + np.testing.assert_array_equal(keep, np.array([3, 0])) + + +def test_nms_free_pipeline_path_never_calls_nms(monkeypatch) -> None: + metadata = _metadata(branch="nms-free") + pipeline = OnnxDetectorPipeline( + metadata=metadata, + session=OnnxDetectorSession(OrtSessionConfig(model_path=Path("model.onnx"))), + ) + boxes = np.array( + [ + [0.0, 0.0, 10.0, 10.0], + [1.0, 1.0, 9.0, 9.0], + [20.0, 20.0, 30.0, 30.0], + ], + dtype=np.float32, + ) + scores = np.array([0.8, 0.95, 0.7], dtype=np.float32) + classes = np.array([1, 1, 2]) + + def fail_nms(*_args: object, **_kwargs: object) -> np.ndarray: + raise AssertionError("NMS-free path must not call class_aware_nms") + + monkeypatch.setattr( + "complexity.deploy.onnx_detector.postprocess.class_aware_nms", + fail_nms, + ) + + filtered_boxes, filtered_scores, filtered_classes = pipeline.postprocess_single_image( + boxes, + scores, + classes, + ) + + np.testing.assert_allclose(filtered_scores, np.array([0.95, 0.8], dtype=np.float32)) + np.testing.assert_allclose( + filtered_boxes, + np.array([[1.0, 1.0, 9.0, 9.0], [0.0, 0.0, 10.0, 10.0]], dtype=np.float32), + ) + np.testing.assert_array_equal(filtered_classes, np.array([1, 1])) diff --git a/tests/test_onnx_detector_metadata.py b/tests/test_onnx_detector_metadata.py new file mode 100644 index 00000000..e17e76f7 --- /dev/null +++ b/tests/test_onnx_detector_metadata.py @@ -0,0 +1,70 @@ +import json + +import pytest + +from complexity.deploy.onnx_detector.metadata import ( + DEFAULT_CONFIDENCE_THRESHOLD, + DEFAULT_IOU_THRESHOLD, + load_metadata, + metadata_from_mapping, + validate_output_shape, +) + + +def _sidecar(branch: str = "o2m") -> dict[str, object]: + return { + "architecture_version": 8, + "image_size": 640, + "num_classes": 80, + "num_cells": 34000, + "regression_width": 68, + "reg_max": 16, + "scale_factors": [1, 2, 4], + "grid_sizes": [160, 80, 40, 20], + "p2_head": True, + "branch": branch, + "requires_nms": branch == "o2m", + "output_semantics": "raw_ltrb_dfl_and_quality_class_logits", + } + + +def test_load_metadata_applies_defaults_and_derives_strides(tmp_path) -> None: + metadata_path = tmp_path / "model.json" + metadata_path.write_text(json.dumps(_sidecar())) + + metadata = load_metadata(metadata_path) + + assert metadata.branch == "o2m" + assert metadata.requires_nms is True + assert metadata.dfl_bins == 17 + assert metadata.prediction_width == 148 + assert metadata.strides == (4.0, 8.0, 16.0, 32.0) + assert metadata.confidence_threshold == DEFAULT_CONFIDENCE_THRESHOLD + assert metadata.iou_threshold == DEFAULT_IOU_THRESHOLD + + +def test_metadata_normalizes_nms_free_branch_spelling() -> None: + values = _sidecar("nms_free") + values["requires_nms"] = False + + metadata = metadata_from_mapping(values) + + assert metadata.branch == "nms-free" + assert metadata.requires_nms is False + + +def test_metadata_rejects_inconsistent_cell_count() -> None: + values = _sidecar() + values["num_cells"] = 123 + + with pytest.raises(ValueError, match="num_cells"): + metadata_from_mapping(values) + + +def test_validate_output_shape_accepts_dynamic_batch_and_rejects_width_mismatch() -> None: + metadata = metadata_from_mapping(_sidecar()) + + validate_output_shape(metadata, ("batch_size", 34000, 148)) + + with pytest.raises(ValueError, match="prediction_width"): + validate_output_shape(metadata, (1, 34000, 149)) diff --git a/tests/test_onnx_detector_pipeline.py b/tests/test_onnx_detector_pipeline.py new file mode 100644 index 00000000..03444941 --- /dev/null +++ b/tests/test_onnx_detector_pipeline.py @@ -0,0 +1,215 @@ +import json + +import numpy as np +import pytest +import torch +from PIL import Image + +from complexity.deploy.onnx_detector.metadata import OnnxDetectorMetadata +from complexity.deploy.onnx_detector.pipeline import OnnxDetectorPipeline +from complexity.deploy.onnx_detector.preprocess import ImageGeometry, restore_boxes +from complexity.generative.detection import TRHashDetectorConfig, TRHashObjectDetector +from complexity.generative.detection.exporting import RawDetectorExport +from complexity.generative.detection.hub import ( + preprocess_detector_image, + restore_detector_boxes, +) + + +class FakeSession: + provider_used = "FakeExecutionProvider" + + def __init__(self, predictions: np.ndarray) -> None: + self.predictions = predictions + self.input_shape: tuple[int, ...] | None = None + + def run(self, pixel_values: np.ndarray) -> np.ndarray: + self.input_shape = pixel_values.shape + return self.predictions + + +def _metadata() -> OnnxDetectorMetadata: + return OnnxDetectorMetadata( + architecture_version=8, + image_size=16, + num_classes=3, + num_cells=16, + regression_width=12, + reg_max=2, + scale_factors=(1,), + grid_sizes=(4,), + p2_head=False, + branch="nms-free", + requires_nms=False, + output_semantics="raw_ltrb_dfl_and_quality_class_logits", + confidence_threshold=0.9, + max_detections=1, + ) + + +def _tiny_detector_config() -> TRHashDetectorConfig: + return TRHashDetectorConfig( + architecture_version=8, + image_size=32, + patch_size=8, + vision_hidden_size=16, + vision_layers=3, + vision_heads=4, + vision_num_experts=2, + vision_top_k=1, + vision_expert_width=8, + vision_stage_depths=(1, 1, 1), + vision_window_size=2, + vision_precision="fp32", + num_classes=3, + reg_max=4, + head_hidden_size=16, + end_to_end=True, + ) + + +def _export_metadata(config: TRHashDetectorConfig, branch: str) -> dict[str, object]: + return { + "architecture_version": config.architecture_version, + "image_size": config.image_size, + "num_classes": config.num_classes, + "num_cells": config.num_cells, + "regression_width": config.regression_width, + "reg_max": config.reg_max, + "scale_factors": list(config.scale_factors), + "grid_sizes": list(config.grid_sizes), + "p2_head": config.p2_head, + "branch": branch, + "requires_nms": branch == "o2m", + "output_semantics": "raw_ltrb_dfl_and_quality_class_logits", + "confidence_threshold": 0.0, + "iou_threshold": 0.45, + "max_detections": 5, + } + + +def test_restore_boxes_undoes_letterbox_for_non_square_image() -> None: + boxes = np.array([[2.0, 2.0, 6.0, 6.0]], dtype=np.float32) + geometry = ImageGeometry( + original_width=16, + original_height=8, + image_size=16, + scale=1.0, + left=0, + top=4, + ) + + restored = restore_boxes(boxes, geometry) + + np.testing.assert_allclose(restored, np.array([[2.0, 0.0, 6.0, 2.0]])) + + +def test_pipeline_predict_runs_fake_session_and_returns_stable_schema() -> None: + metadata = _metadata() + predictions = np.full( + (1, metadata.num_cells, metadata.prediction_width), + -20.0, + dtype=np.float32, + ) + # First cell center is (2, 2), stride is 4, DFL distances are [0, 0, 1, 1]. + predictions[0, 0, 0] = 40.0 + predictions[0, 0, 3] = 40.0 + predictions[0, 0, 7] = 40.0 + predictions[0, 0, 10] = 40.0 + predictions[0, 0, metadata.regression_width + 2] = 20.0 + session = FakeSession(predictions) + pipeline = OnnxDetectorPipeline(metadata=metadata, session=session) # type: ignore[arg-type] + image = np.zeros((8, 16, 3), dtype=np.uint8) + + result = pipeline.predict(image) + + assert session.input_shape == (1, 3, 16, 16) + assert result.provider_used == "FakeExecutionProvider" + assert result.branch_type == "nms-free" + assert len(result.detections) == 1 + detection = result.detections[0] + assert detection.class_id == 2 + assert detection.score > 0.99 + np.testing.assert_allclose(detection.box_norm, (0.125, 0.125, 0.375, 0.375)) + np.testing.assert_allclose(detection.box_pixel, (2.0, 0.0, 6.0, 2.0)) + assert result.timing.preprocess_ms >= 0.0 + assert result.timing.inference_ms >= 0.0 + assert result.timing.postprocess_ms >= 0.0 + + +@pytest.mark.parametrize("branch", ("o2m", "nms-free")) +def test_onnx_pipeline_matches_pytorch_postprocessing_on_fixed_input( + tmp_path, + branch: str, +) -> None: + pytest.importorskip("onnx") + pytest.importorskip("onnxruntime") + torch.manual_seed(13) + model = TRHashObjectDetector(_tiny_detector_config()).eval() + export_model = RawDetectorExport(model, branch).eval() + onnx_path = tmp_path / f"detector-{branch}.onnx" + metadata_path = tmp_path / f"detector-{branch}.json" + metadata_path.write_text(json.dumps(_export_metadata(model.config, branch))) + dummy_input = torch.randn(1, 3, model.config.image_size, model.config.image_size) + + torch.onnx.export( + export_model, + dummy_input, + str(onnx_path), + opset_version=17, + input_names=["pixel_values"], + output_names=["predictions"], + do_constant_folding=True, + dynamo=False, + ) + + pixels = np.arange(20 * 28 * 3, dtype=np.uint8).reshape(20, 28, 3) + image = Image.fromarray(pixels) + pipeline = OnnxDetectorPipeline.from_files( + onnx_path, + metadata_path, + providers=("CPUExecutionProvider",), + ) + onnx_result = pipeline.predict(image) + + pytorch_pixels, geometry = preprocess_detector_image(image, model.config.image_size) + with torch.no_grad(): + if branch == "nms-free": + reference = model.predict_end_to_end( + pytorch_pixels.unsqueeze(0), + confidence_threshold=0.0, + max_detections=5, + )[0] + else: + reference = model.predict( + pytorch_pixels.unsqueeze(0), + confidence_threshold=0.0, + iou_threshold=0.45, + max_detections=5, + nms_free=False, + )[0] + reference_boxes_pixel = restore_detector_boxes(reference["boxes"], geometry) + + assert len(onnx_result.detections) == len(reference["scores"]) + np.testing.assert_array_equal( + np.array([detection.class_id for detection in onnx_result.detections]), + reference["labels"].cpu().numpy(), + ) + np.testing.assert_allclose( + np.array([detection.score for detection in onnx_result.detections]), + reference["scores"].cpu().numpy(), + atol=1e-3, + rtol=1e-3, + ) + np.testing.assert_allclose( + np.array([detection.box_norm for detection in onnx_result.detections]), + reference["boxes"].cpu().numpy(), + atol=1e-3, + rtol=1e-3, + ) + np.testing.assert_allclose( + np.array([detection.box_pixel for detection in onnx_result.detections]), + reference_boxes_pixel.cpu().numpy(), + atol=1e-3, + rtol=1e-3, + ) diff --git a/tests/test_onnx_detector_skeleton.py b/tests/test_onnx_detector_skeleton.py new file mode 100644 index 00000000..09760181 --- /dev/null +++ b/tests/test_onnx_detector_skeleton.py @@ -0,0 +1,53 @@ +from pathlib import Path + +from complexity.deploy.onnx_detector import ( + DEFAULT_CONFIDENCE_THRESHOLD, + DEFAULT_IOU_THRESHOLD, + DEFAULT_MAX_DETECTIONS, + Detection, + DetectionResult, + OnnxDetectorMetadata, + OnnxDetectorPipeline, + TimingBreakdown, +) + + +def test_onnx_detector_package_imports_and_schema_constructs() -> None: + metadata = OnnxDetectorMetadata( + architecture_version=8, + image_size=640, + num_classes=80, + num_cells=34000, + regression_width=68, + reg_max=16, + scale_factors=(1, 2, 4), + grid_sizes=(160, 80, 40, 20), + p2_head=True, + branch="o2m", + requires_nms=True, + output_semantics="raw_ltrb_dfl_and_quality_class_logits", + ) + session = OnnxDetectorPipeline.create_session( + Path("model.onnx"), + providers=("CPUExecutionProvider",), + ) + detection = Detection( + box_norm=(0.1, 0.2, 0.3, 0.4), + box_pixel=(64.0, 128.0, 192.0, 256.0), + class_id=1, + score=0.9, + ) + result = DetectionResult( + detections=(detection,), + timing=TimingBreakdown(1.0, 2.0, 3.0), + provider_used="CPUExecutionProvider", + branch_type=metadata.branch, + ) + + assert metadata.dfl_bins == 17 + assert metadata.prediction_width == 148 + assert metadata.confidence_threshold == DEFAULT_CONFIDENCE_THRESHOLD + assert metadata.iou_threshold == DEFAULT_IOU_THRESHOLD + assert metadata.max_detections == DEFAULT_MAX_DETECTIONS + assert session.config.model_path == Path("model.onnx") + assert result.detections[0].box_norm == (0.1, 0.2, 0.3, 0.4)