79 lines
3.3 KiB
Python
79 lines
3.3 KiB
Python
#!/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)
|