Update app.py

app.py

@@ -3,26 +3,29 @@ import io
 import time
 import json
 import tempfile
-from typing import List, Dict, Tuple
+from typing import List, Dict, Tuple, Optional
 
 import cv2
 import gradio as gr
 import numpy as np
 import pandas as pd
+import requests
 
+# YOLO wrapper (pulls pretrained model from Hugging Face by ID)
 from ultralyticsplus import YOLO, render_result
 
 # =========================
 # CONFIG
 # =========================
 MODEL_ID = "mshamrai/yolov8s-visdrone"
+
 SAMPLES_DIR = "samples"
 SAMPLE_IMAGE = os.path.join(SAMPLES_DIR, "drone_sample.jpg")
 SAMPLE_VIDEO = os.path.join(SAMPLES_DIR, "airspace_sample.mp4")
 
+# Small public files for smoke testing (replace with your own if desired)
 SAMPLE_URLS = {
     SAMPLE_IMAGE: "https://huggingface.co/datasets/hf-internal-testing/example-documents/resolve/main/airplane.jpg",
-    # This is a small demo clip just to validate pipeline; replace with your own short UAV/airspace clip if you prefer.
     SAMPLE_VIDEO: "https://huggingface.co/datasets/hf-internal-testing/example-documents/resolve/main/short_harvard_bridge.mp4",
 }
 
@@ -35,35 +38,56 @@ def _ensure_samples():
         if os.path.exists(local_path):
             continue
         try:
-            import requests
-            r = requests.get(url, timeout=15)
+            r = requests.get(url, timeout=20)
             r.raise_for_status()
             with open(local_path, "wb") as f:
                 f.write(r.content)
         except Exception:
-            # If download fails (e.g., no internet policy), we just skip; UI still works with user uploads.
+            # If download fails (e.g., offline build), UI still works with user uploads
             pass
 
 _ensure_samples()
 
 # =========================
-# MODEL
+# DIAGNOSTICS
+# =========================
+def _ffmpeg_ok() -> bool:
+    try:
+        v = cv2.getBuildInformation()
+        return ("FFMPEG:YES" in v) or ("FFMPEG:                      YES" in v)
+    except Exception:
+        return False
+
+# =========================
+# MODEL (robust lazy loader)
 # =========================
-_model = None
-def load_model(conf: float, iou: float):
-    global _model
-    if _model is None:
-        _model = YOLO(MODEL_ID)
+_model: Optional[YOLO] = None
+_model_error: Optional[str] = None
+
+def load_model(conf: float, iou: float) -> YOLO:
+    """
+    Load the pretrained YOLO model once and set runtime thresholds.
+    Raises RuntimeError if loading previously failed.
+    """
+    global _model, _model_error
+    if _model is None and _model_error is None:
+        try:
+            m = YOLO(MODEL_ID)     # pulls weights from HF
+            m.overrides["max_det"] = 300
+            _model = m
+        except Exception as e:
+            _model_error = f"Model load failed: {e}"
+    if _model_error:
+        raise RuntimeError(_model_error)
     _model.overrides["conf"] = float(conf)
     _model.overrides["iou"] = float(iou)
-    _model.overrides["max_det"] = 300
     return _model
 
 # =========================
 # UTILS
 # =========================
 def results_to_rows(results) -> List[dict]:
-    rows = []
+    rows: List[dict] = []
     if not results:
         return rows
     r = results[0]
@@ -86,18 +110,28 @@ def results_to_rows(results) -> List[dict]:
     return rows
 
 def dict_count_by_class(rows: List[dict]) -> Dict[str, int]:
-    tally = {}
+    tally: Dict[str, int] = {}
     for r in rows:
         tally[r["class"]] = tally.get(r["class"], 0) + 1
     return tally
 
-def write_video(path: str, fps: float, w: int, h: int):
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    return cv2.VideoWriter(path, fourcc, fps, (w, h))
+def write_video(base_path: str, fps: float, w: int, h: int) -> Tuple[cv2.VideoWriter, str]:
+    """
+    Try MP4 first; if it fails (codec not available), fall back to AVI/MJPG.
+    Returns (writer, output_path).
+    """
+    # MP4
+    mp4_path = base_path if base_path.endswith(".mp4") else base_path + ".mp4"
+    writer = cv2.VideoWriter(mp4_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
+    if writer is not None and getattr(writer, "isOpened", lambda: False)():
+        return writer, mp4_path
+    # Fallback AVI
+    avi_path = os.path.splitext(mp4_path)[0] + ".avi"
+    writer = cv2.VideoWriter(avi_path, cv2.VideoWriter_fourcc(*"MJPG"), fps, (w, h))
+    return writer, avi_path
 
 def save_dataframe_to_csv(rows: List[dict]) -> str:
     if not rows:
-        # create an empty CSV for consistency
         df = pd.DataFrame(columns=["class","confidence","x1","y1","x2","y2","width","height"])
     else:
         df = pd.DataFrame(rows)
@@ -108,8 +142,8 @@ def save_dataframe_to_csv(rows: List[dict]) -> str:
 def save_pdf_report(title: str,
                     summary_text: str,
                     counts: Dict[str, int],
-                    annotated_image_path: str | None = None) -> str:
-    # Light-weight PDF (no external dependencies besides reportlab)
+                    annotated_image_path: Optional[str] = None) -> str:
+    # Lightweight PDF via reportlab
     from reportlab.lib.pagesizes import A4
     from reportlab.pdfgen import canvas
     from reportlab.lib.units import cm
@@ -125,7 +159,7 @@ def save_pdf_report(title: str,
     y -= 1.2*cm
 
     c.setFont("Helvetica", 11)
-    for line in summary_text.splitlines():
+    for line in (summary_text or "").splitlines():
         c.drawString(2*cm, y, line[:110])
         y -= 0.7*cm
 
@@ -166,17 +200,21 @@ def detect_on_image(image: np.ndarray, conf: float, iou: float):
     model = load_model(conf, iou)
     results = model.predict(image, imgsz=960, verbose=False)
     rows = results_to_rows(results)
-    annotated = render_result(image, results[0])
+    annotated = render_result(image, results[0])  # returns np.ndarray in BGR
 
     counts = dict_count_by_class(rows)
     summary = "Detections: " + ", ".join(f"{k}: {v}" for k, v in counts.items()) if rows else "No objects detected."
 
-    # Save a temp annotated image for PDF export convenience
+    # Save annotated image (ensure correct color order for disk write)
     tmp_img = os.path.join(tempfile.gettempdir(), f"annotated_{int(time.time())}.jpg")
-    cv2.imwrite(tmp_img, annotated[:, :, ::-1])  # BGR->RGB guard if needed
-    csv_path = save_dataframe_to_csv(rows)
+    try:
+        # render_result returns BGR; cv2.imwrite expects BGR, so write directly
+        cv2.imwrite(tmp_img, annotated)
+    except Exception:
+        tmp_img = None
 
-    return annotated, rows, summary, csv_path, tmp_img
+    csv_path = save_dataframe_to_csv(rows)
+    return annotated[:, :, ::-1], rows, summary, csv_path, tmp_img  # Convert to RGB for Gradio Image
 
 def detect_on_video(video_path: str, conf: float, iou: float, max_frames: int = 300):
     if not video_path:
@@ -191,37 +229,40 @@ def detect_on_video(video_path: str, conf: float, iou: float, max_frames: int =
     w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 1280)
     h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 720)
 
-    out_path = os.path.join(tempfile.gettempdir(), f"annotated_{int(time.time())}.mp4")
-    writer = write_video(out_path, fps, w, h)
+    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, None, "Video writer could not open. Try another format/resolution.", None
 
-    total_counts = {}
+    total_counts: Dict[str, int] = {}
     frame_idx = 0
-    while True:
-        ok, frame = cap.read()
-        if not ok:
-            break
-        frame_idx += 1
-        if frame_idx > int(max_frames):
-            break
-
-        results = model.predict(frame, imgsz=960, verbose=False)
-        for row in results_to_rows(results):
-            total_counts[row["class"]] = total_counts.get(row["class"], 0) + 1
+    try:
+        while True:
+            ok, frame = cap.read()
+            if not ok:
+                break
+            frame_idx += 1
+            if frame_idx > int(max_frames):
+                break
 
-        annotated = render_result(frame, results[0])
-        writer.write(annotated)
+            results = model.predict(frame, imgsz=960, verbose=False)
+            for row in results_to_rows(results):
+                total_counts[row["class"]] = total_counts.get(row["class"], 0) + 1
 
-    cap.release()
-    writer.release()
+            annotated = render_result(frame, results[0])
+            writer.write(annotated)
+    finally:
+        cap.release()
+        writer.release()
 
     summary = "Detections (frame-wise tallies): " + ", ".join(f"{k}: {v}" for k, v in total_counts.items()) if total_counts else "No objects detected."
-    # For videos, CSV is a tally (not per-box) to keep file small
+    # For videos, export a compact CSV tally
     rows = [{"class": k, "count": v} for k, v in sorted(total_counts.items())]
     csv_path = save_dataframe_to_csv(rows)
 
     return out_path, total_counts, summary, csv_path
 
-def export_pdf_image(summary: str, table_rows: List[dict], annotated_tmp_jpg: str):
+def export_pdf_image(summary: str, table_rows: List[dict], annotated_tmp_jpg: Optional[str]):
     counts = dict_count_by_class(table_rows or [])
     pdf_path = save_pdf_report(
         title="Airspace Drone Detector — Image Report",
@@ -281,12 +322,12 @@ No dataset or training required — just run it.
                 pdf_img_btn = gr.Button("Generate PDF Report")
                 pdf_img_path = gr.File(label="PDF Report", interactive=False)
 
-            # Hidden state for annotated path (for PDF embedding)
+            # Hidden state for annotated-image path (for PDF embedding)
             annotated_tmp_img_path = gr.State(value=None)
 
             def _run_img(image, conf, iou):
-                annotated, rows, summary, csv_path, tmp_img = detect_on_image(image, conf, iou)
-                return annotated, rows, summary, csv_path, tmp_img
+                annotated_rgb, rows, summary, csv_path, tmp_img = detect_on_image(image, conf, iou)
+                return annotated_rgb, rows, summary, csv_path, tmp_img
 
             run_img.click(
                 fn=_run_img,
@@ -300,7 +341,6 @@ No dataset or training required — just run it.
                 outputs=[pdf_img_path],
             )
 
-            # Prefilled example (if sample exists)
             if os.path.exists(SAMPLE_IMAGE):
                 gr.Examples(
                     examples=[[SAMPLE_IMAGE]],
@@ -352,9 +392,13 @@ No dataset or training required — just run it.
                 )
 
     gr.Markdown(
-        """
-**Model:** `mshamrai/yolov8s-visdrone` (pretrained; pulled via `ultralyticsplus`)  
-**Credits:** Ultralytics and VisDrone community weights.  
+        f"""
+**Model:** `{MODEL_ID}` (pretrained; pulled via `ultralyticsplus`)  
+**Diagnostics**  
+- FFmpeg available: {'Yes' if _ffmpeg_ok() else 'No'}  
+- Python: 3.10 (set via runtime.txt)  
+- Torch: 2.3.1 (pinned in requirements)  
+- Ultralytics: 8.3.x  
         """
     )