File size: 1,724 Bytes
eb9afd8 | 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 | #!/usr/bin/env python3
"""
det01 — Inspect det_500m.onnx input/output shapes and test batch > 1.
Answers: does the model support batched frame input, or is it fixed at batch=1?
"""
import argparse
import numpy as np
import onnxruntime as ort
import onnx
parser = argparse.ArgumentParser()
parser.add_argument('--model', required=True, help='Path to det_500m.onnx')
args = parser.parse_args()
# ── Declared shapes from ONNX graph ──────────────────────────────────────────
model = onnx.load(args.model)
print('=== Declared input shapes ===')
for inp in model.graph.input:
shape = [d.dim_param if d.HasField('dim_param') else d.dim_value
for d in inp.type.tensor_type.shape.dim]
print(f' {inp.name}: {shape}')
print('\n=== Declared output shapes ===')
for out in model.graph.output:
shape = [d.dim_param if d.HasField('dim_param') else d.dim_value
for d in out.type.tensor_type.shape.dim]
print(f' {out.name}: {shape}')
# ── Runtime test with N=1 and N=2 ────────────────────────────────────────────
print('\n=== Runtime batch test ===')
sess = ort.InferenceSession(args.model, providers=['CPUExecutionProvider'])
inp0 = sess.get_inputs()[0]
out_names = [o.name for o in sess.get_outputs()]
for N in (1, 2):
dummy = np.random.randint(0, 255, (N, 3, 640, 640)).astype(np.float32)
try:
outs = sess.run(out_names, {inp0.name: dummy})
print(f' N={N}: OK output shapes={[list(o.shape) for o in outs]}')
except Exception as e:
print(f' N={N}: FAILED — {e}')
|