"""Composable pre/post-processing helpers shared across model families."""
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Literal
import numpy as np
@dataclass(frozen=True, slots=True)
class LetterboxResult:
"""Letterboxed image plus the mapping back to source pixel coordinates."""
array: np.ndarray
scale: float
pad_x: int
pad_y: int
[docs]
def letterbox(
array: np.ndarray,
size: tuple[int, int],
*,
pad_value: int = 114,
center: bool = False,
) -> LetterboxResult:
"""Resize an HWC image to (height, width) preserving aspect ratio with padding."""
import cv2
target_h, target_w = size
source_h, source_w = array.shape[:2]
scale = min(target_h / source_h, target_w / source_w)
resized_w = max(1, int(round(source_w * scale)))
resized_h = max(1, int(round(source_h * scale)))
resized = cv2.resize(array, (resized_w, resized_h), interpolation=cv2.INTER_LINEAR)
channels = () if array.ndim == 2 else (array.shape[2],)
padded = np.full((target_h, target_w, *channels), pad_value, dtype=array.dtype)
pad_x = (target_w - resized_w) // 2 if center else 0
pad_y = (target_h - resized_h) // 2 if center else 0
padded[pad_y : pad_y + resized_h, pad_x : pad_x + resized_w] = resized
return LetterboxResult(array=padded, scale=scale, pad_x=pad_x, pad_y=pad_y)
[docs]
def image_to_tensor(
array: np.ndarray,
*,
layout: Literal["NCHW", "NHWC"] = "NCHW",
channel_order: Literal["BGR", "RGB"] = "BGR",
dtype: np.dtype | type = np.float32,
scale: float | None = None,
mean: Sequence[float] | None = None,
std: Sequence[float] | None = None,
) -> np.ndarray:
"""Convert an HWC BGR image array into a batched model input tensor.
Applies optional value scaling (e.g. 1/255) and per-channel mean/std
normalization, then transposes to the requested layout with batch dim 1.
"""
if array.ndim != 3:
raise ValueError(f"Expected an HWC image array, got shape {array.shape}")
tensor = array[:, :, ::-1] if channel_order == "RGB" else array
tensor = tensor.astype(dtype, copy=True)
if scale is not None:
tensor *= scale
if mean is not None:
tensor -= np.asarray(mean, dtype=tensor.dtype)
if std is not None:
tensor /= np.asarray(std, dtype=tensor.dtype)
if layout == "NCHW":
tensor = tensor.transpose(2, 0, 1)
return np.ascontiguousarray(tensor[np.newaxis, ...])
[docs]
def nms(boxes_xyxy: np.ndarray, scores: np.ndarray, iou_threshold: float) -> list[int]:
"""Greedy non-maximum suppression; returns kept indices by descending score."""
if len(boxes_xyxy) == 0:
return []
x1, y1, x2, y2 = boxes_xyxy[:, 0], boxes_xyxy[:, 1], boxes_xyxy[:, 2], boxes_xyxy[:, 3]
areas = np.maximum(0.0, x2 - x1) * np.maximum(0.0, y2 - y1)
order = scores.argsort()[::-1]
keep: list[int] = []
while order.size > 0:
best = int(order[0])
keep.append(best)
rest = order[1:]
inter_w = np.maximum(0.0, np.minimum(x2[best], x2[rest]) - np.maximum(x1[best], x1[rest]))
inter_h = np.maximum(0.0, np.minimum(y2[best], y2[rest]) - np.maximum(y1[best], y1[rest]))
intersection = inter_w * inter_h
union = areas[best] + areas[rest] - intersection
iou = np.where(union > 0, intersection / union, 0.0)
order = rest[iou <= iou_threshold]
return keep
[docs]
def softmax(values: np.ndarray, axis: int = -1) -> np.ndarray:
"""Numerically stable softmax."""
shifted = values - values.max(axis=axis, keepdims=True)
exp = np.exp(shifted)
return exp / exp.sum(axis=axis, keepdims=True)
__all__ = [
"LetterboxResult",
"image_to_tensor",
"letterbox",
"nms",
"softmax",
]