Update app.py

app.py

@@ -76,7 +76,7 @@ def _get_model(conf: float, iou: float):
             except Exception as e:
                 last_err = e
         if _model is None:
-            _model_err = f"Model load failed. Last error: {last_err}"
+            _model_err = f"Model load failed. Tried {len(MODEL_CANDIDATES)} candidate(s). Last error: {last_err}"
     if _model_err:
         raise RuntimeError(_model_err)
     _model.overrides["conf"] = float(conf)
@@ -166,74 +166,118 @@ def _save_pdf(title: str, summary: str, counts: Dict[str, int], annotated_image_
     return out_path
 
 # =========================
-# INFERENCE
+# INFERENCE (SAFE WRAPPERS)
 # =========================
-def detect_image(image, conf: float, iou: float):
-    if image is None:
-        return None, [], "No image provided.", None, None
-    cv2 = _lazy_cv2()
-    model = _get_model(conf, iou)
-    results = model.predict(image, imgsz=960, verbose=False)
-    r = results[0]
-    rows = _results_to_rows(results)
-    annotated = r.plot()  # BGR ndarray
-    counts = _count_by_class(rows)
-    summary = "Detections: " + (", ".join(f"{k}: {v}" for k, v in counts.items()) if rows else "none")
-    tmp_img = os.path.join(tempfile.gettempdir(), f"annotated_{int(time.time())}.jpg")
+def detect_image_safe(image, conf: float, iou: float):
     try:
-        cv2.imwrite(tmp_img, annotated)
-    except Exception:
-        tmp_img = None
-    csv_path = _save_csv(rows)
-    annotated_rgb = annotated[:, :, ::-1] if hasattr(annotated, "shape") and len(annotated.shape) == 3 else annotated
-    return annotated_rgb, rows, summary, csv_path, tmp_img
-
-def detect_video(video_path: str, conf: float, iou: float, max_frames: int = 300):
-    if not video_path:
-        return None, {}, "No video provided.", None
-    cv2 = _lazy_cv2()
-    model = _get_model(conf, iou)
-    cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        return None, {}, "Failed to open video.", None
-    fps = cap.get(cv2.CAP_PROP_FPS) or 24.0
-    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 1280)
-    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 720)
-    writer, out_path = _write_video(os.path.join(tempfile.gettempdir(), f"annotated_{int(time.time())}"), fps, w, h)
-    if writer is None or (hasattr(writer, "isOpened") and not writer.isOpened()):
-        cap.release()
-        return None, {}, "Video writer could not open. Try another format/resolution.", None
-    totals: Dict[str, int] = {}
-    frames = 0
+        if image is None:
+            return None, [], "⚠️ No image provided.", None, None
+        cv2 = _lazy_cv2()
+
+        # If the component is type="filepath", it sends a string path.
+        # If type="numpy", it sends an array. Ultralytics handles both,
+        # but we’ll keep it robust either way.
+        model = _get_model(conf, iou)
+        results = model.predict(image, imgsz=960, verbose=False)
+        r = results[0]
+        rows = _results_to_rows(results)
+        annotated = r.plot()  # BGR ndarray
+        counts = _count_by_class(rows)
+        summary = "Detections: " + (", ".join(f"{k}: {v}" for k, v in counts.items()) if rows else "none")
+
+        tmp_img = os.path.join(tempfile.gettempdir(), f"annotated_{int(time.time())}.jpg")
+        try:
+            cv2.imwrite(tmp_img, annotated)
+        except Exception:
+            tmp_img = None
+
+        csv_path = _save_csv(rows)
+
+        # Convert BGR->RGB for display if ndarray
+        try:
+            if hasattr(annotated, "shape") and len(annotated.shape) == 3:
+                annotated = annotated[:, :, ::-1]
+        except Exception:
+            pass
+
+        return annotated, rows, summary, csv_path, tmp_img
+
+    except Exception as e:
+        # Return error in the summary field instead of crashing
+        return None, [], f"❌ Error during image detection: {e}", None, None
+
+def detect_video_safe(video_path: str, conf: float, iou: float, max_frames: int = 300):
     try:
-        while True:
-            ok, frame = cap.read()
-            if not ok:
-                break
-            frames += 1
-            if frames > int(max_frames):
-                break
-            results = model.predict(frame, imgsz=960, verbose=False)
-            r = results[0]
-            for row in _results_to_rows(results):
-                totals[row["class"]] = totals.get(row["class"], 0) + 1
-            annotated = r.plot()
-            writer.write(annotated)
-    finally:
-        cap.release()
-        writer.release()
-    summary = "Detections (frame-wise tallies): " + (", ".join(f"{k}: {v}" for k, v in totals.items()) if totals else "none")
-    tally_rows = [{"class": k, "count": v} for k, v in sorted(totals.items())]
-    csv_path = _save_csv(tally_rows)
-    return out_path, totals, summary, csv_path
+        if not video_path:
+            return None, "{}", "⚠️ No video provided.", None
 
-def export_pdf_img(summary: str, table_rows: List[dict], annotated_tmp_jpg: Optional[str]):
-    counts = _count_by_class(table_rows or [])
-    return _save_pdf("UAV Detector — Image Report", summary or "No summary.", counts,
-                     annotated_tmp_jpg if annotated_tmp_jpg and os.path.exists(annotated_tmp_jpg) else None)
+        cv2 = _lazy_cv2()
+        model = _get_model(conf, iou)
+
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            return None, "{}", "❌ Failed to open video.", None
 
-def export_pdf_vid(summary: str, counts: dict):
-    return _save_pdf("UAV Detector — Video Report", summary or "No summary.", counts or {}, None)
+        fps = cap.get(cv2.CAP_PROP_FPS) or 24.0
+        w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 1280)
+        h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 720)
+
+        writer, out_path = _write_video(os.path.join(tempfile.gettempdir(), f"annotated_{int(time.time())}"), fps, w, h)
+        if writer is None or (hasattr(writer, "isOpened") and not writer.isOpened()):
+            cap.release()
+            return None, "{}", "❌ Video writer could not open. Try another format/resolution.", None
+
+        totals: Dict[str, int] = {}
+        frames = 0
+        try:
+            while True:
+                ok, frame = cap.read()
+                if not ok:
+                    break
+                frames += 1
+                if frames > int(max_frames):
+                    break
+
+                results = model.predict(frame, imgsz=960, verbose=False)
+                r = results[0]
+
+                for row in _results_to_rows(results):
+                    totals[row["class"]] = totals.get(row["class"], 0) + 1
+
+                annotated = r.plot()
+                writer.write(annotated)
+        finally:
+            cap.release()
+            writer.release()
+
+        summary = "Detections (frame-wise tallies): " + (", ".join(f"{k}: {v}" for k, v in totals.items()) if totals else "none")
+        tally_rows = [{"class": k, "count": v} for k, v in sorted(totals.items())]
+        csv_path = _save_csv(tally_rows)
+        counts_json = json.dumps(totals or {}, ensure_ascii=False, indent=2)
+
+        return out_path, counts_json, summary, csv_path
+
+    except Exception as e:
+        return None, "{}", f"❌ Error during video detection: {e}", None
+
+def export_pdf_img(summary: str, table_rows: List[dict], annotated_tmp_jpg: Optional[str]):
+    try:
+        counts = _count_by_class(table_rows or [])
+        return _save_pdf("UAV Detector — Image Report", summary or "No summary.", counts,
+                         annotated_tmp_jpg if annotated_tmp_jpg and os.path.exists(annotated_tmp_jpg) else None)
+    except Exception as e:
+        # Create a tiny report that just contains the error
+        return _save_pdf("UAV Detector — Image Report", f"❌ PDF export error: {e}", {}, None)
+
+def export_pdf_vid(summary: str, counts_json: str):
+    try:
+        counts = json.loads(counts_json) if counts_json else {}
+    except Exception:
+        counts = {}
+    try:
+        return _save_pdf("UAV Detector — Video Report", summary or "No summary.", counts or {}, None)
+    except Exception as e:
+        return _save_pdf("UAV Detector — Video Report", f"❌ PDF export error: {e}", {}, None)
 
 # =========================
 # UI (local embedded samples)
@@ -259,9 +303,10 @@ If they’re missing, you can still upload your own.
         # ---------- IMAGE ----------
         with gr.TabItem("Image"):
             with gr.Row():
+                # Use type="filepath" so embedded path loads directly. Uploads also pass a path.
                 image_in = gr.Image(
                     value=EMBED_IMG if os.path.exists(EMBED_IMG) else None,
-                    type="filepath",  # load directly from path
+                    type="filepath",
                     label="Input Image (embedded or upload)"
                 )
                 with gr.Column():
@@ -285,11 +330,8 @@ If they’re missing, you can still upload your own.
                 pdf_img_path = gr.File(label="PDF Report", interactive=False)
             annotated_tmp_img_path = gr.State(value=None)
 
-            def _run_img(image, conf, iou):
-                return detect_image(image, conf, iou)
-
             run_img.click(
-                fn=_run_img,
+                fn=detect_image_safe,
                 inputs=[image_in, conf_img, iou_img],
                 outputs=[image_out, table_out, msg_img, csv_img_path, annotated_tmp_img_path],
             )
@@ -328,25 +370,14 @@ If they’re missing, you can still upload your own.
                 pdf_vid_btn = gr.Button("Generate PDF Report")
                 pdf_vid_path = gr.File(label="PDF Report", interactive=False)
 
-            def _run_vid(vpath, conf, iou, maxf):
-                out_path, counts, summary, csv_path = detect_video(vpath, conf, iou, int(maxf))
-                return out_path, json.dumps(counts or {}, ensure_ascii=False, indent=2), summary, csv_path
-
             run_vid.click(
-                fn=_run_vid,
+                fn=detect_video_safe,
                 inputs=[video_in, conf_vid, iou_vid, max_frames],
                 outputs=[video_out, counts_text, msg_vid, csv_vid_path],
             )
 
-            def _export_pdf_vid(summary: str, counts_json_str: str):
-                try:
-                    counts = json.loads(counts_json_str) if counts_json_str else {}
-                except Exception:
-                    counts = {}
-                return export_pdf_vid(summary, counts)
-
             pdf_vid_btn.click(
-                fn=_export_pdf_vid,
+                fn=export_pdf_vid,
                 inputs=[msg_vid, counts_text],
                 outputs=[pdf_vid_path],
             )