Update app.py

app.py

@@ -1,188 +1,125 @@
-# app.py
 import gradio as gr
 import onnxruntime as ort
 import numpy as np
 from PIL import Image, ImageDraw, ImageFont
 import time, os
+import pyttsx3  # for optional voice alerts
 
 # ---------------------------
 # 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
+MODEL_PATH = "model.onnx"
+INPUT_SIZE = (640, 640)
+CONF_THRESHOLD_DEFAULT = 0.35
 
-# ---------------------------
-# 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
+# Initialize voice engine
+engine = pyttsx3.init()
+engine.setProperty("rate", 180)
 
-# ---------------------------
-# LOAD MODEL
-# ---------------------------
+# 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(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}")
+    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):
+# Preprocess
+def preprocess_frame(frame_np):
     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
+    img_resized = img.resize(INPUT_SIZE)
+    arr = np.array(img_resized).astype(np.float32) / 255.0
+    arr = np.transpose(arr, (2, 0, 1))[np.newaxis, ...]  # 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]
-
+# Postprocess
+def postprocess_outputs(outputs, orig_w, orig_h, conf_thresh=0.35):
+    outs = [np.array(o) for o in outputs]
+    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:]
+        for row in cand:
+            cx, cy, w, h = row[0], row[1], row[2], row[3]
+            obj_conf = float(row[4])
+            class_probs = row[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
+                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
-# ---------------------------
+# Draw boxes
 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"]]
+        label = f"Class {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):
+# Voice alert
+last_spoken = ""
+def speak_alert(detections):
+    global last_spoken
+    if not detections:
+        return
+    labels_detected = [f"class {d['class']}" for d in detections]
+    msg = ", ".join(set(labels_detected))
+    if msg != last_spoken:
+        engine.say(f"Detected: {msg}")
+        engine.runAndWait()
+        last_spoken = msg
+
+# Main function
+def predict_live(frame, conf_threshold):
     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)
+    orig_h, orig_w = frame.shape[:2]
+    input_tensor = preprocess_frame(frame)
     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}"
+    outputs = sess.run(None, {input_name: input_tensor})
+    detections = postprocess_outputs(outputs, orig_w, orig_h, conf_thresh=conf_threshold)
 
-    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()
+    speak_alert(detections)
+
     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"
+        f"Detections: {len(detections)}"
     )
     return out_img, debug_txt
 
-# ---------------------------
-# GRADIO INTERFACE
-# ---------------------------
+# Gradio interface with webcam + slider
 iface = gr.Interface(
     fn=predict_live,
-    inputs=gr.Image(source="webcam", type="numpy"),
-    outputs=[gr.Image(type="pil"), gr.Textbox(lines=6)],
+    inputs=[
+        gr.Image(sources=["webcam"], type="numpy", label="Live Camera"),
+        gr.Slider(0.05, 0.9, value=CONF_THRESHOLD_DEFAULT, step=0.05, label="Confidence Threshold")
+    ],
+    outputs=[gr.Image(type="pil"), gr.Textbox(lines=4)],
     live=True,
-    title="ONNX Live Detection",
-    description="Real-time detection using your ONNX model. Adjust CONF_THRESHOLD or input size if needed."
+    title="ONNX Live Camera Detection",
+    description="Continuous live detection with bounding boxes + voice alerts"
 )
 
 if __name__ == "__main__":