File size: 2,405 Bytes
ebeed67 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | """YOLO-NAS postprocessor — Frigate-style flat batched-NMS ONNX export.
Frigate+ ``yolonas`` models (and any super-gradients ``model.export()`` with
fused NMS in FLAT format) emit ONE output tensor of shape ``[N, 7]`` where each
row is ``(batch_index, x1, y1, x2, y2, score, class_id)``:
- coordinates are in INPUT pixel space (0..inputSize), post-NMS,
- GPU exports pad unused slots with a NEGATIVE class id (skip them),
- class ids index the model's own label map (Frigate+ label set or COCO-80).
Output matches the yolo postprocessor:
``{"kind": "detections", "detections": [{"class", "score", "bbox": [x1,y1,x2,y2]}]}``
with pixel bboxes re-projected to ORIGINAL frame coordinates (letterbox undone
via the pool-provided ``scale`` + ``pad``).
"""
from __future__ import annotations
import numpy as np
from .yolo import COCO_80
def postprocess_yolonas(
predictions: dict,
config: dict,
orig_w: int,
orig_h: int,
scale: float,
pad: tuple[int, int],
) -> dict:
conf_threshold = float(config.get("confidence", 0))
labels = config.get("labels", COCO_80)
output = np.asarray(list(predictions.values())[0])
# Accept [1, N, 7] batched exports by squeezing the batch axis.
if output.ndim == 3:
output = output.reshape(-1, output.shape[-1])
if output.ndim != 2 or output.shape[-1] != 7:
raise ValueError(
f"yolonas postprocessor expects a flat [N, 7] output, got shape {output.shape}"
)
detections = []
for row in output:
_batch_idx, x1, y1, x2, y2, score, class_id = row
score = float(score)
class_id = int(class_id)
if class_id < 0:
continue
if score < conf_threshold:
continue
ox1 = max(0.0, min(float(orig_w), (float(x1) - pad[0]) / scale))
oy1 = max(0.0, min(float(orig_h), (float(y1) - pad[1]) / scale))
ox2 = max(0.0, min(float(orig_w), (float(x2) - pad[0]) / scale))
oy2 = max(0.0, min(float(orig_h), (float(y2) - pad[1]) / scale))
label = labels[class_id] if class_id < len(labels) else str(class_id)
detections.append(
{
"class": label,
"score": round(score, 4),
"bbox": [round(ox1, 1), round(oy1, 1), round(ox2, 1), round(oy2, 1)],
}
)
return {"kind": "detections", "detections": detections}
|