File size: 2,194 Bytes
1f1eb3d | 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 | #!/usr/bin/env python3
"""
Minimal demo: batched inference with the fixed det_500m model.
Stacks N input images into a single batch and runs one forward pass.
Prints output shapes and a per-frame summary (count of cls anchors above a
score threshold per scale). Intended as a sanity check, not a full detector —
plug your own NMS / anchor decoding for actual face boxes.
"""
import argparse
import cv2
import numpy as np
import onnxruntime as ort
DET_SIZE = (640, 640)
SCORE_THRESHOLD = 0.5
parser = argparse.ArgumentParser()
parser.add_argument('--model', required=True, help='Path to det_500m_fixed.onnx')
parser.add_argument('--images', nargs='+', required=True,
help='One or more image paths (any number — they form the batch)')
args = parser.parse_args()
def preprocess(path):
img = cv2.imread(path)
if img is None:
raise FileNotFoundError(path)
blob = cv2.dnn.blobFromImage(
cv2.resize(img, DET_SIZE),
1.0 / 128.0, DET_SIZE,
(127.5, 127.5, 127.5), swapRB=True,
)
return blob[0]
# Build batch
batch = np.stack([preprocess(p) for p in args.images], axis=0)
print(f'Input batch: {batch.shape} ({len(args.images)} image(s))')
# One forward pass
sess = ort.InferenceSession(args.model, providers=['CPUExecutionProvider'])
inp_name = sess.get_inputs()[0].name
out_names = [o.name for o in sess.get_outputs()]
outputs = sess.run(None, {inp_name: batch})
# Output shapes
print('\nOutputs:')
for n, o in zip(out_names, outputs):
print(f' {n}: {list(o.shape)}')
# Per-frame summary using the 3 cls heads (post-Sigmoid, indices 0/1/2)
strides = [8, 16, 32]
cls_outputs = outputs[:3]
print(f'\nPer-frame anchor counts above score {SCORE_THRESHOLD}:')
print(f' {"image":<40} {"stride 8":>10} {"stride 16":>10} {"stride 32":>10} {"total":>8}')
for n in range(batch.shape[0]):
counts = [int((cls[n] > SCORE_THRESHOLD).sum()) for cls in cls_outputs]
name = args.images[n]
if len(name) > 38:
name = '...' + name[-35:]
print(f' {name:<40} {counts[0]:>10} {counts[1]:>10} {counts[2]:>10} {sum(counts):>8}')
print('\nIf all images are the same, the per-frame counts must match exactly.')
|