#!/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}')