| """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}") |
|
|
