"""Minimal RF-DETR ONNX license-plate detector. License: Apache 2.0."""
import argparse
import numpy as np
import onnxruntime as ort
from PIL import Image

RES = 768
MEAN = np.array([0.485, 0.456, 0.406], np.float32)
STD = np.array([0.229, 0.224, 0.225], np.float32)


def detect(sess, in_name, pil, conf=0.5):
    w0, h0 = pil.size
    img = pil.convert("RGB").resize((RES, RES), Image.BILINEAR)
    x = (np.asarray(img, np.float32) / 255.0 - MEAN) / STD
    x = np.ascontiguousarray(x.transpose(2, 0, 1)[None], np.float32)
    dets, labels = sess.run(["dets", "labels"], {in_name: x})
    d = dets[0] if dets.ndim == 3 else dets
    l = labels[0] if labels.ndim == 3 else labels
    scores = l.max(axis=-1) if l.ndim > 1 else l
    if scores.max() > 1 or scores.min() < 0:
        scores = 1 / (1 + np.exp(-scores))
    out = []
    for (cx, cy, bw, bh), s in zip(d, scores):
        if s < conf:
            continue
        out.append((max(0, (cx - bw / 2) * w0), max(0, (cy - bh / 2) * h0),
                    min(w0, (cx + bw / 2) * w0), min(h0, (cy + bh / 2) * h0), float(s)))
    return out


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("image")
    ap.add_argument("--onnx", default="rfdetr-large.onnx")
    ap.add_argument("--conf", type=float, default=0.5)
    args = ap.parse_args()
    sess = ort.InferenceSession(args.onnx, providers=ort.get_available_providers())
    for box in detect(sess, sess.get_inputs()[0].name, Image.open(args.image), args.conf):
        print(f"x1={box[0]:.0f} y1={box[1]:.0f} x2={box[2]:.0f} y2={box[3]:.0f} score={box[4]:.2f}")