Add inference script with decoder

inference.py

@@ -0,0 +1,173 @@
+"""
+05_inference_test.py — Inference for YOLOv4-tiny raw ONNX (float or INT8).
+
+Reuses the decoder pattern from clean_yolov4/05b_inference_int8_raw.py
+but with YOLOv4-tiny's 2-head config.
+"""
+import os, sys, time, argparse
+import numpy as np, cv2, requests
+from io import BytesIO
+from PIL import Image, ImageDraw, ImageFont
+import onnxruntime as ort
+
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+
+# YOLOv4-tiny constants
+ANCHORS = [(10,14),(23,27),(37,58),(81,82),(135,169),(344,319)]
+# (stride, anchor_indices, scale_xy) — order matches DarknetRaw export
+CFG_HEADS = [
+    (32, [3, 4, 5], 1.05),   # out0 = stride 32 (13x13)
+    (16, [1, 2, 3], 1.05),   # out1 = stride 16 (26x26)
+]
+
+INPUT_SIZE = 416
+SCORE_THR  = 0.30
+NMS_THR    = 0.45
+
+COCO = ["person","bicycle","car","motorcycle","airplane","bus","train","truck","boat",
+        "traffic light","fire hydrant","stop sign","parking meter","bench","bird","cat",
+        "dog","horse","sheep","cow","elephant","bear","zebra","giraffe","backpack","umbrella",
+        "handbag","tie","suitcase","frisbee","skis","snowboard","sports ball","kite",
+        "baseball bat","baseball glove","skateboard","surfboard","tennis racket","bottle",
+        "wine glass","cup","fork","knife","spoon","bowl","banana","apple","sandwich",
+        "orange","broccoli","carrot","hot dog","pizza","donut","cake","chair","couch",
+        "potted plant","bed","dining table","toilet","tv","laptop","mouse","remote",
+        "keyboard","cell phone","microwave","oven","toaster","sink","refrigerator","book",
+        "clock","vase","scissors","teddy bear","hair drier","toothbrush"]
+
+np.random.seed(42)
+PALETTE = [(int(r), int(g), int(b)) for r, g, b in np.random.randint(60, 255, (80, 3))]
+
+
+def get_font(size):
+    for f in ("arialbd.ttf", "arial.ttf", "segoeui.ttf"):
+        try: return ImageFont.truetype(f, size)
+        except Exception: continue
+    return ImageFont.load_default()
+
+
+def sigmoid(x):
+    return 1.0 / (1.0 + np.exp(-np.clip(x, -50, 50)))
+
+
+def letterbox_nchw(rgb, size=INPUT_SIZE):
+    h, w = rgb.shape[:2]; s = min(size/h, size/w)
+    nh, nw = int(round(h*s)), int(round(w*s))
+    resized = cv2.resize(rgb, (nw, nh))
+    pad = np.full((size, size, 3), 114, np.uint8); pad[:nh,:nw] = resized
+    chw = pad.astype(np.float32).transpose(2,0,1) / 255.0
+    return np.expand_dims(chw, 0), s
+
+
+def decode_one_head(raw, stride, anchor_idxs, scale_xy):
+    _, ch, H, W = raw.shape
+    n_anchors = len(anchor_idxs)
+    n_classes = ch // n_anchors - 5
+    x = raw.reshape(1, n_anchors, 5 + n_classes, H, W).transpose(0, 3, 4, 1, 2)[0]
+    txty = sigmoid(x[..., 0:2]) * scale_xy - (scale_xy - 1) / 2
+    twth = np.exp(np.clip(x[..., 2:4], -10, 10))
+    obj  = sigmoid(x[..., 4:5])
+    cls  = sigmoid(x[..., 5:])
+    anc = np.array([[ANCHORS[i][0]/stride, ANCHORS[i][1]/stride] for i in anchor_idxs], dtype=np.float32)
+    twth = twth * anc[None, None, :, :]
+    yy, xx = np.meshgrid(np.arange(H), np.arange(W), indexing="ij")
+    grid = np.stack([xx, yy], axis=-1).astype(np.float32)
+    txty = txty + grid[:, :, None, :]
+    txty *= stride; twth *= stride
+    pred = np.concatenate([txty, twth, obj, cls], axis=-1).reshape(-1, 5 + n_classes)
+    return pred
+
+
+def decode_all(raws, ratio):
+    all_pred = [decode_one_head(r, s, a, sxy) for r, (s, a, sxy) in zip(raws, CFG_HEADS)]
+    pred = np.concatenate(all_pred, axis=0)
+    obj = pred[:, 4]; cls = pred[:, 5:]
+    cls_id = np.argmax(cls, axis=1)
+    cls_score = cls[np.arange(len(cls)), cls_id]
+    score = obj * cls_score
+    keep = score > SCORE_THR
+    pred = pred[keep]; score = score[keep]; cls_id = cls_id[keep]
+    if len(pred) == 0: return []
+    cx, cy, w, h = pred[:,0], pred[:,1], pred[:,2], pred[:,3]
+    x1 = (cx-w/2)/ratio; y1 = (cy-h/2)/ratio; x2 = (cx+w/2)/ratio; y2 = (cy+h/2)/ratio
+    dets = [{"class": int(cls_id[i]), "score": float(score[i]),
+             "bbox":[float(x1[i]), float(y1[i]), float(x2[i]), float(y2[i])]}
+            for i in range(len(pred))]
+    dets.sort(key=lambda d: -d["score"])
+    keep_d = []
+    while dets:
+        keep_d.append(dets[0]); rest = []
+        for d in dets[1:]:
+            if d["class"] != keep_d[-1]["class"]:
+                rest.append(d); continue
+            ax1,ay1,ax2,ay2 = keep_d[-1]["bbox"]; bx1,by1,bx2,by2 = d["bbox"]
+            iw = max(0,min(ax2,bx2)-max(ax1,bx1)); ih = max(0,min(ay2,by2)-max(ay1,by1))
+            inter = iw*ih; aa=max(0,(ax2-ax1)*(ay2-ay1)); ab=max(0,(bx2-bx1)*(by2-by1))
+            iou = inter/(aa+ab-inter+1e-9)
+            if iou < NMS_THR: rest.append(d)
+        dets = rest
+    return keep_d
+
+
+def draw(pil, dets):
+    img = pil.copy(); d = ImageDraw.Draw(img)
+    W, H = img.size
+    th = max(3, min(W,H)//200); font = get_font(max(14, min(W,H)//40))
+    for x in dets:
+        x1, y1, x2, y2 = x["bbox"]
+        x1 = max(0, min(x1,W-1)); y1 = max(0, min(y1,H-1))
+        x2 = max(0, min(x2,W-1)); y2 = max(0, min(y2,H-1))
+        cls = x["class"]; cname = COCO[cls]; color = PALETTE[cls % len(PALETTE)]
+        for t in range(th):
+            d.rectangle([x1-t, y1-t, x2+t, y2+t], outline=color)
+        label = f"{cname} {x['score']*100:.0f}%"
+        bb = d.textbbox((x1, y1-18), label, font=font)
+        d.rectangle(bb, fill=color); d.text((bb[0], bb[1]), label, fill=(0,0,0), font=font)
+    return img
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--onnx", default=os.path.join(SCRIPT_DIR, "out_onnx", "yolov4-tiny-416_float_raw.onnx"))
+    ap.add_argument("--out-dir", default=os.path.join(SCRIPT_DIR, "inference"))
+    args = ap.parse_args()
+    os.makedirs(args.out_dir, exist_ok=True)
+    if not os.path.isfile(args.onnx):
+        print(f"[FAIL] no existe: {args.onnx}"); return 1
+    print(f"Loading: {args.onnx}  ({os.path.getsize(args.onnx)/1e6:.1f} MB)")
+    sess = ort.InferenceSession(args.onnx, providers=["CPUExecutionProvider"])
+    inp = sess.get_inputs()[0].name
+    print(f"  inputs:  {[i.name for i in sess.get_inputs()]}")
+    print(f"  outputs: {[(o.name, o.shape) for o in sess.get_outputs()]}")
+    tests = [
+        ("dog",     "https://raw.githubusercontent.com/pjreddie/darknet/master/data/dog.jpg"),
+        ("traffic", "http://images.cocodataset.org/val2017/000000011197.jpg"),
+        ("skaters", "http://images.cocodataset.org/val2017/000000087038.jpg"),
+        ("kitchen", "http://images.cocodataset.org/val2017/000000037777.jpg"),
+        ("market",  "http://images.cocodataset.org/val2017/000000289343.jpg"),
+        ("parking", "http://images.cocodataset.org/val2017/000000017627.jpg"),
+        ("dining",  "http://images.cocodataset.org/val2017/000000080340.jpg"),
+        ("bus",     "http://images.cocodataset.org/val2017/000000000785.jpg"),
+    ]
+    for name, url in tests:
+        try:
+            r = requests.get(url, timeout=30); r.raise_for_status()
+            pil = Image.open(BytesIO(r.content)).convert("RGB")
+        except Exception as e:
+            print(f"[skip {name}] {e}"); continue
+        rgb = np.array(pil)
+        blob, ratio = letterbox_nchw(rgb, INPUT_SIZE)
+        t0 = time.time()
+        outs = sess.run(None, {inp: blob})
+        t = (time.time()-t0)*1000
+        dets = decode_all(outs, ratio)
+        annotated = draw(pil, dets)
+        out_path = os.path.join(args.out_dir, f"{name}.png")
+        annotated.save(out_path)
+        print(f"  {name:>10s}: {len(dets):>2d} dets in {t:6.1f} ms")
+        for d in dets[:8]:
+            print(f"      {COCO[d['class']]:>16s}  {d['score']*100:5.1f}%")
+
+
+if __name__ == "__main__":
+    sys.exit(main() or 0)