Update app.py

app.py

@@ -5,23 +5,41 @@ from typing import List, Dict, Tuple, Optional
 import json
 import gradio as gr
 
-# -------------------
-# Config
-# -------------------
-REPO_ID = "mshamrai/yolov8s-visdrone"
-FILENAME = "weights/best.pt"
-
+# =========================================================
+# Config — you can override these via Space Secrets / Env
+# =========================================================
+# If you know the exact HF repo + file you want, set:
+#   HF_MODEL_REPO = "owner/repo"
+#   HF_MODEL_FILE = "path/to/weights.pt"
+HF_MODEL_REPO = os.getenv("HF_MODEL_REPO", "").strip()
+HF_MODEL_FILE = os.getenv("HF_MODEL_FILE", "").strip()
+
+# Fallback candidates (tried in order) — real drone/UAV detectors
+MODEL_CANDIDATES = []
+if HF_MODEL_REPO and HF_MODEL_FILE:
+    MODEL_CANDIDATES.append((HF_MODEL_REPO, HF_MODEL_FILE))
+
+# A couple of known community models. If one is unavailable, the next is tried.
+MODEL_CANDIDATES += [
+    ("keremberke/yolov8n-drone-detection", "best.pt"),  # small, fast
+    ("keremberke/yolov8m-drone-detection", "best.pt"),  # larger, more accurate
+]
+
+# Embedded samples (we’ll download a short drone clip and auto‑extract a frame as the image)
 SAMPLES_DIR = "samples"
-EMBED_IMG = os.path.join(SAMPLES_DIR, "aerial_image.jpg")
-EMBED_VID = os.path.join(SAMPLES_DIR, "aerial_video.mp4")
+EMBED_VID = os.path.join(SAMPLES_DIR, "uav_sample.mp4")
+EMBED_IMG = os.path.join(SAMPLES_DIR, "uav_sample_frame.jpg")
+DRONE_VIDEO_URL = "https://github.com/ultralytics/assets/releases/download/v0.0.0/drone.mp4"
 
-# -------------------
+# =========================================================
 # Lazy state
-# -------------------
+# =========================================================
 _model = None
 _model_err = None
 _model_names = None
 _ffmpeg_status = None
+_loaded_repo = None
+_loaded_file = None
 
 def _lazy_cv2():
     import cv2
@@ -39,35 +57,71 @@ def _ffmpeg_ok() -> bool:
         _ffmpeg_status = False
     return _ffmpeg_status
 
-def _lazy_hf_download() -> str:
+def _ensure_samples():
+    os.makedirs(SAMPLES_DIR, exist_ok=True)
+    # Download drone video if missing
+    if not os.path.exists(EMBED_VID):
+        try:
+            import requests
+            r = requests.get(DRONE_VIDEO_URL, timeout=30)
+            r.raise_for_status()
+            with open(EMBED_VID, "wb") as f:
+                f.write(r.content)
+        except Exception:
+            pass
+    # Extract one frame from the video as the image sample
+    if os.path.exists(EMBED_VID) and not os.path.exists(EMBED_IMG):
+        try:
+            cv2 = _lazy_cv2()
+            cap = cv2.VideoCapture(EMBED_VID)
+            # Skip a few frames so the drone is centered
+            frame_no = 15
+            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_no)
+            ok, frame = cap.read()
+            cap.release()
+            if ok and frame is not None:
+                cv2.imwrite(EMBED_IMG, frame)
+        except Exception:
+            pass
+
+_ensure_samples()
+
+def _download_from_hf(repo_id: str, filename: str) -> str:
     from huggingface_hub import hf_hub_download
-    return hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
+    return hf_hub_download(repo_id=repo_id, filename=filename)
 
 def _get_model(conf: float, iou: float):
-    """Load YOLO weights from HF on first use."""
-    global _model, _model_err, _model_names
+    """Try to load a UAV-specific YOLO model from the candidate list."""
+    global _model, _model_err, _model_names, _loaded_repo, _loaded_file
     if _model is None and _model_err is None:
-        try:
-            from ultralytics import YOLO
-            weights_path = _lazy_hf_download()
-            m = YOLO(weights_path)
-            m.overrides["max_det"] = 300
-            _model = m
+        from ultralytics import YOLO
+        last_err = None
+        for repo, file in MODEL_CANDIDATES:
             try:
-                _model_names = m.model.names if hasattr(m, "model") else None
-            except Exception:
-                _model_names = None
-        except Exception as e:
-            _model_err = f"Model load failed: {e}"
+                weights_path = _download_from_hf(repo, file)
+                m = YOLO(weights_path)
+                m.overrides["max_det"] = 300
+                _model = m
+                _loaded_repo, _loaded_file = repo, file
+                try:
+                    _model_names = m.model.names if hasattr(m, "model") else None
+                except Exception:
+                    _model_names = None
+                break
+            except Exception as e:
+                last_err = e
+                continue
+        if _model is None and last_err is not None:
+            _model_err = f"Model load failed. Tried: {MODEL_CANDIDATES}. Last error: {last_err}"
     if _model_err:
         raise RuntimeError(_model_err)
     _model.overrides["conf"] = float(conf)
     _model.overrides["iou"] = float(iou)
     return _model
 
-# -------------------
+# =========================================================
 # Helpers
-# -------------------
+# =========================================================
 def _results_to_rows(results) -> List[dict]:
     rows: List[dict] = []
     if not results:
@@ -147,9 +201,9 @@ def _save_pdf(title: str, summary: str, counts: Dict[str, int], annotated_image_
     c.showPage(); c.save()
     return out_path
 
-# -------------------
+# =========================================================
 # Inference
-# -------------------
+# =========================================================
 def detect_image(image, conf: float, iou: float):
     if image is None:
         return None, [], "No image provided.", None, None
@@ -211,26 +265,26 @@ def detect_video(video_path: str, conf: float, iou: float, max_frames: int = 300
 
 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("Airspace Drone Detector — Image Report", summary or "No summary.", counts,
+    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)
 
 def export_pdf_vid(summary: str, counts: dict):
-    return _save_pdf("Airspace Drone Detector — Video Report", summary or "No summary.", counts or {}, None)
+    return _save_pdf("UAV Detector — Video Report", summary or "No summary.", counts or {}, None)
 
-# -------------------
-# UI (embedded-local samples + uploads)
-# -------------------
+# =========================================================
+# UI (embedded UAV samples + uploads)
+# =========================================================
 NOTE = (
-    "Model: VisDrone (aerial **cars/pedestrians/vehicles**). It does **not** include a 'drone' class. "
-    "Use top‑down scenes with people/traffic for best results."
+    "UAV model: detects drones (class names vary per checkpoint, e.g., 'drone', 'uav'). "
+    "Use scenes where the drone occupies enough pixels (≥ 30–40 px on the short side)."
 )
 
-with gr.Blocks(title="Aerial Object Detector (VisDrone)") as demo:
+with gr.Blocks(title="UAV / Drone Detector (YOLO)") as demo:
     gr.Markdown(
         """
-# Aerial Object Detector (Pretrained on VisDrone)
-Use the **embedded samples** or your own uploads.  
-Exports: **CSV** and **PDF** reports.
+# UAV / Drone Detector (Pretrained YOLO)
+We embedded a **drone video** and auto‑extracted an **image frame** so you can test immediately.
+Use your own uploads too. Exports: **CSV** and **PDF**.
         """
     )
 
@@ -246,7 +300,7 @@ Exports: **CSV** and **PDF** reports.
                 with gr.Column():
                     conf_img = gr.Slider(0.05, 0.8, 0.35, step=0.05, label="Confidence")
                     iou_img = gr.Slider(0.1, 0.9, 0.45, step=0.05, label="NMS IoU")
-                    load_embed_img = gr.Button("Load Embedded Sample Image")
+                    load_embed_img = gr.Button("Load Embedded UAV Image")
                     run_img = gr.Button("Run Detection")
                     gr.Markdown(NOTE)
 
@@ -262,9 +316,7 @@ Exports: **CSV** and **PDF** reports.
             annotated_tmp_img_path = gr.State(value=None)
 
             def _load_embed_img():
-                if os.path.exists(EMBED_IMG):
-                    return EMBED_IMG
-                return None
+                return EMBED_IMG if os.path.exists(EMBED_IMG) else None
 
             load_embed_img.click(fn=_load_embed_img, outputs=[image_in])
 
@@ -294,7 +346,7 @@ Exports: **CSV** and **PDF** reports.
                     conf_vid = gr.Slider(0.05, 0.8, 0.35, step=0.05, label="Confidence")
                     iou_vid = gr.Slider(0.1, 0.9, 0.45, step=0.05, label="NMS IoU")
                     max_frames = gr.Slider(60, 2000, 300, step=10, label="Max frames to process")
-                    load_embed_vid = gr.Button("Load Embedded Sample Video")
+                    load_embed_vid = gr.Button("Load Embedded UAV Video")
                     run_vid = gr.Button("Run Detection")
                     gr.Markdown(NOTE)
 
@@ -308,9 +360,7 @@ Exports: **CSV** and **PDF** reports.
                 pdf_vid_path = gr.File(label="PDF Report", interactive=False)
 
             def _load_embed_vid():
-                if os.path.exists(EMBED_VID):
-                    return EMBED_VID
-                return None
+                return EMBED_VID if os.path.exists(EMBED_VID) else None
 
             load_embed_vid.click(fn=_load_embed_vid, outputs=[video_in])
 
@@ -338,11 +388,13 @@ Exports: **CSV** and **PDF** reports.
                 outputs=[pdf_vid_path],
             )
 
+    # Footer / diagnostics
+    model_str = f"{_loaded_repo}/{_loaded_file}" if _loaded_repo else "loading on first run"
     gr.Markdown(
         f"""
-**Weights:** `{REPO_ID}/{FILENAME}` (downloaded lazily)  
-**Diagnostics** — FFmpeg: {'Yes' if _ffmpeg_ok() else 'No'} • Python: 3.10  
-**Tip:** For true *drone* detection, I can swap in a UAV‑specific model. Say the word and I’ll rewire it.
+**Model:** {model_str}  
+**Diagnostics:** FFmpeg: {'Yes' if _ffmpeg_ok() else 'No'} • Python: 3.10  
+If loading fails, set Space Secrets `HF_MODEL_REPO` and `HF_MODEL_FILE` to a known drone checkpoint.
         """
     )