Update app.py

app.py

@@ -1,82 +1,189 @@
-# app.py
-import gradio as gr
-import onnxruntime as ort
-import numpy as np
-from PIL import Image, ImageDraw, ImageFont
-import cv2
-import io
-
-# Load labels
-with open("labels.txt", "r") as f:
-    LABELS = [x.strip() for x in f.readlines()]
-
-# Load ONNX model
-sess = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
-
-# Helper: simple preprocess - adapt to your model's needs
-def preprocess(pil_img, input_size=(640,640)):
-    img = pil_img.convert("RGB")
-    img = img.resize(input_size)
-    np_img = np.array(img).astype(np.float32) / 255.0   # normalize if model expected 0-1
-    # Change to channels-first if model expects (1,C,H,W)
-    np_img = np.transpose(np_img, (2,0,1))[np.newaxis, ...]
-    return np_img
-
-# Helper: very basic NMS and postprocess - adapt as necessary
-def postprocess(outputs, orig_w, orig_h, conf_threshold=0.3, iou_threshold=0.45):
-    # This section depends on your model outputs.
-    # Example: suppose outputs[0] -> [N, 6] with (x1,y1,x2,y2,score,class)
-    preds = outputs[0]
-    boxes = []
-    for row in preds:
-        x1,y1,x2,y2,score,cls = row
-        if score < conf_threshold:
-            continue
-        # scale coords back to original image size if your model used 640x640
-        boxes.append({
-            "box": [x1*orig_w, y1*orig_h, x2*orig_w, y2*orig_h],
-            "score": float(score),
-            "class": int(cls)
-        })
-    # (Optional) Apply NMS here if model doesn't already do it
-    return boxes
-
-def draw_boxes(pil_img, boxes):
-    img = pil_img.convert("RGB")
-    draw = ImageDraw.Draw(img)
-    for b in boxes:
-        x1,y1,x2,y2 = b["box"]
-        label = LABELS[b["class"]] if 0 <= b["class"] < len(LABELS) else str(b["class"])
-        draw.rectangle([x1,y1,x2,y2], outline="red", width=3)
-        draw.text((x1, y1-10), f"{label} {b['score']:.2f}", fill="red")
-    return img
-
-def predict(image):
-    if image is None:
-        return None, "No image"
-    pil = Image.fromarray(image.astype('uint8')) if isinstance(image, np.ndarray) else Image.open(io.BytesIO(image.read()))
-    orig_w, orig_h = pil.size
-    input_tensor = preprocess(pil)  # adapt input_size if needed
-
-    # Run ONNX
-    input_name = sess.get_inputs()[0].name
-    outputs = sess.run(None, {input_name: input_tensor})
-    # Postprocess according to your model's output structure
-    boxes = postprocess(outputs, orig_w, orig_h)
-    out_img = draw_boxes(pil, boxes)
-    txt = "\n".join([f"{LABELS[b['class']]}: {b['score']:.2f}" for b in boxes]) if boxes else "No detections"
-    return out_img, txt
-
-# Gradio UI
-title = "ONNX Demo"
-desc = "Upload an image or use webcam. Adapt preprocessing/postprocessing per your model."
-
-iface = gr.Interface(fn=predict,
-                     inputs=gr.Image(source="upload", tool="editor", type="numpy"),
-                     outputs=[gr.Image(type="pil"), gr.Textbox()],
-                     title=title,
-                     description=desc,
-                     examples=None)
-
-if __name__ == "__main__":
-    iface.launch()
+# app.py
+import gradio as gr
+import onnxruntime as ort
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+import time, os
+
+# ---------------------------
+# CONFIG
+# ---------------------------
+MODEL_PATH = "model.onnx"       # Ensure model.onnx is in repo
+LABELS_PATH = "labels.txt"      # Optional: one label per line
+CONF_THRESHOLD = 0.35
+PREVIEW_INPUT_SIZE = (640, 640)  # Change if model expects different input size
+
+# ---------------------------
+# LOAD LABELS
+# ---------------------------
+if os.path.exists(LABELS_PATH):
+    with open(LABELS_PATH, "r") as f:
+        LABELS = [l.strip() for l in f.readlines() if l.strip()]
+else:
+    LABELS = None
+
+# ---------------------------
+# LOAD MODEL
+# ---------------------------
+print(f"Loading ONNX model from: {MODEL_PATH}")
+sess = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])
+
+print("\nONNX Model Inputs:")
+for i, inp in enumerate(sess.get_inputs()):
+    print(f"  input[{i}] name={inp.name}, shape={inp.shape}, dtype={inp.type}")
+
+print("\nONNX Model Outputs:")
+for i, out in enumerate(sess.get_outputs()):
+    print(f"  output[{i}] name={out.name}, shape={out.shape}, dtype={out.type}")
+
+# ---------------------------
+# PREPROCESS FUNCTION
+# ---------------------------
+def preprocess_frame(frame_np, input_size=PREVIEW_INPUT_SIZE):
+    img = Image.fromarray(frame_np.astype("uint8"), "RGB")
+    img_resized = img.resize(input_size)
+    arr = np.array(img_resized).astype(np.float32) / 255.0  # normalize 0..1
+    arr = np.transpose(arr, (2, 0, 1))[np.newaxis, ...]  # to NCHW
+    return arr
+
+# ---------------------------
+# POSTPROCESS FUNCTION (FIXED)
+# ---------------------------
+def postprocess_outputs(outputs, orig_w, orig_h, conf_thresh=0.35, debug=False):
+    outs = [o if isinstance(o, np.ndarray) else np.array(o) for o in outputs]
+    if len(outs) == 0:
+        return []
+
+    cand = None
+    for o in outs:
+        if o.ndim >= 2 and o.shape[-1] >= 4:
+            cand = o
+            break
+    if cand is None:
+        cand = outs[0]
+
+    if cand.ndim == 3 and cand.shape[0] == 1:
+        cand = cand[0]
+
+    detections = []
+
+    if debug:
+        print("Raw chosen output shape:", cand.shape)
+        try:
+            print("Sample rows:", cand.reshape(-1, cand.shape[-1])[:5])
+        except Exception:
+            pass
+
+    # Case 1: Nx6
+    if cand.ndim == 2 and cand.shape[1] == 6:
+        for r in cand:
+            x1, y1, x2, y2, score, cls = r
+            if score < conf_thresh:
+                continue
+            if max(x1, y1, x2, y2) <= 1.0:
+                x1, y1, x2, y2 = x1*orig_w, y1*orig_h, x2*orig_w, y2*orig_h
+            detections.append({"box": [x1, y1, x2, y2], "score": float(score), "class": int(cls)})
+        return detections
+
+    # Case 2: YOLO-style Nx(5+num_classes)
+    if cand.ndim == 2 and cand.shape[1] >= 6:
+        for r in cand:
+            cx, cy, w, h = r[0], r[1], r[2], r[3]
+            obj_conf = float(r[4])
+            class_probs = r[5:]
+            best_idx = int(np.argmax(class_probs)) if class_probs.size > 0 else 0
+            cls_conf = float(class_probs[best_idx]) if class_probs.size > 0 else 1.0
+            score = obj_conf * cls_conf
+            if score < conf_thresh:
+                continue
+            if max(cx, cy, w, h) <= 1.0:
+                x1 = (cx - w/2) * orig_w
+                y1 = (cy - h/2) * orig_h
+                x2 = (cx + w/2) * orig_w
+                y2 = (cy + h/2) * orig_h
+            else:
+                x1, y1, x2, y2 = cx - w/2, cy - h/2, cx + w/2, cy + h/2
+            detections.append({"box": [x1, y1, x2, y2], "score": score, "class": best_idx})
+        return detections
+
+    # Case 3: Separate outputs (boxes, scores, labels)
+    if len(outs) >= 3:
+        boxes_arr = next((o for o in outs if o.ndim == 2 and o.shape[1] == 4), None)
+        scores_arr = next((o for o in outs if o.ndim <= 2 and o.size == boxes_arr.shape[0]), None) if boxes_arr is not None else None
+        labels_arr = next((o for o in outs if o.ndim <= 2 and o.size == boxes_arr.shape[0]), None) if boxes_arr is not None else None
+        if boxes_arr is not None:
+            for i, bx in enumerate(boxes_arr):
+                score = float(scores_arr[i]) if scores_arr is not None else 1.0
+                if score < conf_thresh:
+                    continue
+                if max(bx) <= 1.0:
+                    x1, y1, x2, y2 = bx[0]*orig_w, bx[1]*orig_h, bx[2]*orig_w, bx[3]*orig_h
+                else:
+                    x1, y1, x2, y2 = bx
+                detections.append({"box": [x1, y1, x2, y2], "score": score, "class": int(labels_arr[i]) if labels_arr is not None else 0})
+            return detections
+
+    if debug:
+        print("Could not parse model outputs automatically.")
+    return detections
+
+# ---------------------------
+# DRAW BOXES ON IMAGE
+# ---------------------------
+def draw_boxes_on_image(pil_img, detections):
+    img = pil_img.convert("RGB")
+    draw = ImageDraw.Draw(img)
+    font = ImageFont.load_default()
+    for d in detections:
+        x1, y1, x2, y2 = d["box"]
+        label = str(d["class"])
+        if LABELS and 0 <= d["class"] < len(LABELS):
+            label = LABELS[d["class"]]
+        txt = f"{label} {d['score']:.2f}"
+        draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
+        draw.text((x1, max(0, y1 - 12)), txt, fill="red", font=font)
+    return img
+
+# ---------------------------
+# MAIN PREDICT FUNCTION
+# ---------------------------
+def predict_live(frame):
+    if frame is None:
+        return None, "No frame"
+    t0 = time.time()
+    orig_h, orig_w = frame.shape[0], frame.shape[1]
+    input_tensor = preprocess_frame(frame, PREVIEW_INPUT_SIZE)
+    input_name = sess.get_inputs()[0].name
+    try:
+        outputs = sess.run(None, {input_name: input_tensor})
+    except Exception as e:
+        return None, f"ONNX runtime error: {e}"
+
+    detections = postprocess_outputs(outputs, orig_w, orig_h, conf_thresh=CONF_THRESHOLD, debug=True)
+    pil_img = Image.fromarray(frame.astype("uint8"), "RGB")
+    out_img = draw_boxes_on_image(pil_img, detections)
+
+    t1 = time.time()
+    debug_txt = (
+        f"Model: {os.path.basename(MODEL_PATH)}\n"
+        f"Input shape: {sess.get_inputs()[0].shape}\n"
+        f"Output(s): {[o.shape for o in sess.get_outputs()]}\n"
+        f"Detections: {len(detections)}\n"
+        f"Inference time: {(t1 - t0)*1000:.1f} ms"
+    )
+    return out_img, debug_txt
+
+# ---------------------------
+# GRADIO INTERFACE
+# ---------------------------
+iface = gr.Interface(
+    fn=predict_live,
+    inputs=gr.Image(source="webcam", type="numpy"),
+    outputs=[gr.Image(type="pil"), gr.Textbox(lines=6)],
+    live=True,
+    title="ONNX Live Detection",
+    description="Real-time detection using your ONNX model. Adjust CONF_THRESHOLD or input size if needed."
+)
+
+if __name__ == "__main__":
+    iface.launch(server_name="0.0.0.0", server_port=7860)