Add YOLOv26n RDK X5 model and quantized artifacts

x5_quantization/postprocess_yolo26.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+"""CPU post-processing for the six raw outputs of best_bpu.onnx."""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def _nms(boxes: np.ndarray, scores: np.ndarray, iou_threshold: float) -> np.ndarray:
+    order = scores.argsort()[::-1]
+    kept: list[int] = []
+    while order.size:
+        current = int(order[0])
+        kept.append(current)
+        if order.size == 1:
+            break
+        rest = order[1:]
+        xx1 = np.maximum(boxes[current, 0], boxes[rest, 0])
+        yy1 = np.maximum(boxes[current, 1], boxes[rest, 1])
+        xx2 = np.minimum(boxes[current, 2], boxes[rest, 2])
+        yy2 = np.minimum(boxes[current, 3], boxes[rest, 3])
+        inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
+        area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
+        union = area[current] + area[rest] - inter
+        order = rest[(inter / np.maximum(union, 1e-9)) <= iou_threshold]
+    return np.asarray(kept, dtype=np.int64)
+
+
+def postprocess(
+    outputs: list[np.ndarray] | tuple[np.ndarray, ...],
+    classes: int = 4,
+    score_threshold: float = 0.25,
+    iou_threshold: float = 0.7,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Decode `(cls_s8, box_s8, cls_s16, box_s16, cls_s32, box_s32)` outputs.
+
+    Returns `(boxes_xyxy, scores, class_ids)` in the 640x640 model coordinate
+    system. Scale boxes back to the camera image after this function when using
+    letterbox preprocessing.
+    """
+    if len(outputs) != 6:
+        raise ValueError(f"expected six outputs, got {len(outputs)}")
+    conf_raw = -np.log(1.0 / score_threshold - 1.0)
+    detections: list[np.ndarray] = []
+    for index, stride in enumerate((8, 16, 32)):
+        cls = np.asarray(outputs[index * 2]).reshape(-1, classes).astype(np.float32)
+        box = np.asarray(outputs[index * 2 + 1]).reshape(-1, 4).astype(np.float32)
+        max_logits = cls.max(axis=1)
+        valid = np.flatnonzero(max_logits >= conf_raw)
+        if valid.size == 0:
+            continue
+        height = width = 640 // stride
+        grid_y, grid_x = np.indices((height, width))
+        grid = np.stack((grid_x, grid_y), axis=-1).reshape(-1, 2).astype(np.float32) + 0.5
+        anchor = grid[valid]
+        distances = box[valid]
+        xyxy = np.column_stack(
+            ((anchor[:, 0] - distances[:, 0]) * stride,
+             (anchor[:, 1] - distances[:, 1]) * stride,
+             (anchor[:, 0] + distances[:, 2]) * stride,
+             (anchor[:, 1] + distances[:, 3]) * stride)
+        )
+        scores = 1.0 / (1.0 + np.exp(-max_logits[valid]))
+        ids = cls[valid].argmax(axis=1).astype(np.float32)
+        detections.append(np.column_stack((xyxy, scores, ids)))
+
+    if not detections:
+        return np.empty((0, 4), np.float32), np.empty((0,), np.float32), np.empty((0,), np.int32)
+    dets = np.concatenate(detections, axis=0)
+    kept: list[np.ndarray] = []
+    for class_id in np.unique(dets[:, 5]).astype(np.int32):
+        class_dets = dets[dets[:, 5] == class_id]
+        indices = _nms(class_dets[:, :4], class_dets[:, 4], iou_threshold)
+        kept.append(class_dets[indices])
+    if not kept:
+        return np.empty((0, 4), np.float32), np.empty((0,), np.float32), np.empty((0,), np.int32)
+    result = np.concatenate(kept, axis=0)
+    return result[:, :4], result[:, 4], result[:, 5].astype(np.int32)