Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import time
 import tempfile
-from typing import List, Dict, Optional
+from typing import List, Dict, Optional, Tuple
 import json
 import gradio as gr
 
@@ -12,44 +12,41 @@ SAMPLES_DIR = "samples"
 EMBED_IMG = os.path.join(SAMPLES_DIR, "uav_image.jpg")
 EMBED_VID = os.path.join(SAMPLES_DIR, "uav_video.mp4")
 
-# Optional HF secrets (to force a specific checkpoint)
-HF_MODEL_REPO = os.getenv("HF_MODEL_REPO", "").strip()
-HF_MODEL_FILE = os.getenv("HF_MODEL_FILE", "").strip()
-HF_TOKEN = os.getenv("HF_TOKEN", "").strip()
+HF_TOKEN = os.getenv("HF_TOKEN", "").strip()  # optional (for private/gated repos)
 
-# Public pretrained models (no auth needed)
-MODEL_CANDIDATES = []
-if HF_MODEL_REPO and HF_MODEL_FILE:
-    MODEL_CANDIDATES.append((HF_MODEL_REPO, HF_MODEL_FILE))
-MODEL_CANDIDATES += [
-    ("Javvanny/yolov8m_flying_objects_detection", "yolov8m/weights/best.pt"),
-    ("doguilmak/Drone-Detection-YOLOv8x", "weight/best.pt"),
-]
+# Selectable models (public)
+MODEL_CHOICES: Dict[str, Tuple[str, str]] = {
+    "Multi-class (Drone/Helicopter/Airplane/Bird)":
+        ("Javvanny/yolov8m_flying_objects_detection", "yolov8m/weights/best.pt"),
+    "Drone-only (cleaner, fewer false positives)":
+        ("keremberke/yolov8m-drone-detection", "best.pt"),
+}
 
 # =========================
 # LABELS & THREAT RULES
 # =========================
 LABEL_MAP = {
+    "Airplane": "Airplane",
+    "Bird": "Bird",
+    "Drone": "Drone",
+    "Helicopter": "Helicopter",
+    "UAV": "UAV",
     "БПЛА": "UAV",
     "БПЛА коптер": "Drone",
     "квадрокоптер": "Drone",
     "квадроcамолет": "Drone",
     "самолет": "Airplane",
     "вертолет": "Helicopter",
-    "птица": "Bird",
-    "человек": "Person",
-    "машина": "Car",
     "автомобиль": "Car",
+    "машина": "Car",
+    "БПЛА самелет": "UAV Airplane",
+    "drone": "Drone",  # some checkpoints use lowercase
 }
 THREAT_SET = {"drone", "uav", "airplane", "helicopter"}
 
 def map_label(name: str) -> str:
     if not isinstance(name, str): return name
-    if name in LABEL_MAP: return LABEL_MAP[name]
-    low = name.lower()
-    for ru, en in LABEL_MAP.items():
-        if low == ru.lower(): return en
-    return name
+    return LABEL_MAP.get(name, LABEL_MAP.get(name.lower(), name))
 
 def translate_names_dict(names_dict: Dict[int, str]) -> Dict[int, str]:
     if not isinstance(names_dict, dict): return names_dict
@@ -59,12 +56,12 @@ def is_threat(label_en: str) -> bool:
     return label_en and label_en.lower() in THREAT_SET
 
 # =========================
-# FALSE-POSITIVE FILTERS  (tunable via Space secrets)
+# FALSE-POSITIVE FILTERS (tunable via Space Secrets)
 # =========================
 MIN_CONF      = float(os.getenv("MIN_CONF", 0.60))      # drop low-confidence hits
-MIN_AREA_PCT  = float(os.getenv("MIN_AREA_PCT", 0.004)) # drop tiny boxes (fraction of frame area)
+MIN_AREA_PCT  = float(os.getenv("MIN_AREA_PCT", 0.004)) # drop tiny boxes (fraction of frame)
 SKY_RATIO     = float(os.getenv("SKY_RATIO", 0.65))     # keep boxes whose bottoms are above 65% height
-ALLOWED_CLASSES = {"Drone", "UAV", "Helicopter", "Airplane"}
+ALLOWED_CLASSES = {"Drone", "UAV", "Helicopter", "Airplane"}  # ignored for drone-only automatically
 
 # =========================
 # LAZY GLOBAL STATE
@@ -74,6 +71,7 @@ _model_err = None
 _model_names = None
 _loaded_repo = None
 _loaded_file = None
+_loaded_key = None  # which dropdown choice loaded
 _ffmpeg_status = None
 
 def _lazy_cv2():
@@ -98,34 +96,45 @@ def _download_from_hf(repo_id: str, filename: str) -> str:
         except Exception: pass
     return hf_hub_download(repo_id=repo_id, filename=filename)
 
-def _get_model(conf: float, iou: float):
-    """Load first available YOLO model from HF."""
+def _reset_model_cache():
     global _model, _model_err, _model_names, _loaded_repo, _loaded_file
+    _model = None
+    _model_err = None
+    _model_names = None
+    _loaded_repo = None
+    _loaded_file = None
+
+def _get_model(model_key: str, conf: float, iou: float):
+    """Load the YOLO model selected in the dropdown."""
+    from ultralytics import YOLO
+    global _model, _model_err, _model_names, _loaded_repo, _loaded_file, _loaded_key
+    if _loaded_key != model_key:
+        _reset_model_cache()
+        _loaded_key = model_key
+
     if _model is None and _model_err is None:
-        from ultralytics import YOLO
+        repo, file = MODEL_CHOICES[model_key]
         last_err = None
-        for repo, file in MODEL_CANDIDATES:
+        try:
+            weights = _download_from_hf(repo, file)
+            m = YOLO(weights)
+            m.overrides["max_det"] = 300
+            _model = m
+            _loaded_repo, _loaded_file = repo, file
             try:
-                weights = _download_from_hf(repo, file)
-                m = YOLO(weights)
-                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
+                _model_names = m.model.names if hasattr(m, "model") else None
+            except Exception:
+                _model_names = None
+        except Exception as e:
+            last_err = e
+            _model = None
         if _model is None:
-            _model_err = f"Model load failed. Tried {len(MODEL_CANDIDATES)} candidate(s). Last error: {last_err}"
+            _model_err = f"Model load failed for {repo}/{file}. Error: {last_err}"
     if _model_err: raise RuntimeError(_model_err)
     _model.overrides["conf"] = float(conf)
     _model.overrides["iou"]  = float(iou)
     return _model
 
-# -------- live footer helper --------
 def _model_info_text():
     repo = f"{_loaded_repo}/{_loaded_file}" if _loaded_repo else "not loaded"
     try:
@@ -160,8 +169,10 @@ def _results_to_rows(results) -> List[dict]:
         })
     return rows
 
-def _filter_rows_by_geometry(r, rows: List[dict]) -> List[dict]:
-    """Drop low-conf, tiny, ground-region boxes; keep only classes we care about."""
+def _filter_rows_by_geometry(r, rows: List[dict], model_key: str) -> List[dict]:
+    """Drop low-conf, tiny, ground-region boxes; for drone-only, allow Drone/UAV only."""
+    # Restrict classes differently if drone-only
+    allowed = ALLOWED_CLASSES if "Multi-class" in model_key else {"Drone", "UAV"}
     try:
         H, W = r.orig_img.shape[:2]
     except Exception:
@@ -171,7 +182,7 @@ def _filter_rows_by_geometry(r, rows: List[dict]) -> List[dict]:
         if row.get("confidence") is not None and row["confidence"] < MIN_CONF:
             continue
         cls = map_label(str(row.get("class","")))
-        if ALLOWED_CLASSES and cls not in ALLOWED_CLASSES:
+        if allowed and cls not in allowed:
             continue
         if H and W:
             area = row["width"] * row["height"]
@@ -179,7 +190,7 @@ def _filter_rows_by_geometry(r, rows: List[dict]) -> List[dict]:
                 continue
             y_bottom = row["y2"]
             horizon = H * SKY_RATIO
-            if y_bottom > horizon:
+            if y_bottom > horizon:  # below sky line → likely ground/grass noise
                 continue
         kept.append(row)
     return kept
@@ -271,21 +282,6 @@ def _save_pdf_detections(title: str, detections: List[dict], header_note: str =
     c.showPage(); c.save()
     return out_path
 
-def _normalize_rows(table_rows):
-    """Accept DataFrame/list; return list[dict]."""
-    try:
-        import pandas as pd
-        if isinstance(table_rows, pd.DataFrame):
-            return table_rows.to_dict(orient="records")
-    except Exception:
-        pass
-    if isinstance(table_rows, list) and (not table_rows or isinstance(table_rows[0], dict)):
-        return table_rows or []
-    if isinstance(table_rows, list) and table_rows and isinstance(table_rows[0], list):
-        headers = ["class","confidence","x1","y1","x2","y2","width","height"]
-        return [dict(zip(headers,row)) for row in table_rows]
-    return []
-
 def _apply_english_overlay(r):
     try:
         if hasattr(r,"names") and isinstance(r.names, dict):
@@ -296,18 +292,18 @@ def _apply_english_overlay(r):
 # =========================
 # INFERENCE (with filtering + custom draw)
 # =========================
-def detect_image_safe(image, conf: float, iou: float):
+def detect_image_safe(model_key: str, image, conf: float, iou: float):
     try:
         if image is None:
-            return None, [], "⚠️ No image provided.", [], None
+            return None, [], "⚠️ No image provided.", [], None, _model_info_text()
         cv2 = _lazy_cv2()
-        model = _get_model(conf, iou)
+        model = _get_model(model_key, conf, iou)
         results = model.predict(image, imgsz=960, verbose=False)
         r = results[0]
         _apply_english_overlay(r)
 
         rows_raw = _results_to_rows(results)
-        rows = _filter_rows_by_geometry(r, rows_raw)
+        rows = _filter_rows_by_geometry(r, rows_raw, model_key)
 
         annotated_bgr = _draw_annotations_bgr(r.orig_img, rows)
         now_utc = time.strftime("%Y-%m-%d %H:%M:%S UTC", time.gmtime())
@@ -318,28 +314,30 @@ def detect_image_safe(image, conf: float, iou: float):
             "threat": "Threat" if is_threat(map_label(row["class"])) else "Non-threat",
         } for row in rows]
 
-        counts = _count_by_class(rows)
-        summary = "Detections: " + (", ".join(f"{k}: {v}" for k,v in counts.items()) if rows else "none")
+        counts = {}
+        for d in det_records:
+            counts[d["object"]] = counts.get(d["object"], 0) + 1
+        summary = "Detections: " + (", ".join(f"{k}: {v}" for k,v in counts.items()) if counts else "none")
 
         tmp_img = os.path.join(tempfile.gettempdir(), f"annotated_{int(time.time())}.jpg")
         try: cv2.imwrite(tmp_img, annotated_bgr)
         except Exception: tmp_img = None
 
-        annotated_rgb = annotated_bgr[:, :, ::-1]  # BGR→RGB for Gradio
-        return annotated_rgb, rows, summary, det_records, tmp_img
+        annotated_rgb = annotated_bgr[:, :, ::-1]
+        return annotated_rgb, rows, summary, det_records, tmp_img, _model_info_text()
     except Exception as e:
-        return None, [], f"❌ Error during image detection: {e}", [], None
+        return None, [], f"❌ Error during image detection: {e}", [], None, _model_info_text()
 
-def detect_video_safe(video_path: str, conf: float, iou: float, max_frames: int = 300):
+def detect_video_safe(model_key: str, video_path: str, conf: float, iou: float, max_frames: int = 300):
     try:
         if not video_path:
-            return None, "{}", "⚠️ No video provided.", []
+            return None, "{}", "⚠️ No video provided.", [], _model_info_text()
         cv2 = _lazy_cv2()
-        model = _get_model(conf, iou)
+        model = _get_model(model_key, conf, iou)
 
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
-            return None, "{}", "❌ Failed to open video.", []
+            return None, "{}", "❌ Failed to open video.", [], _model_info_text()
 
         fps = cap.get(cv2.CAP_PROP_FPS) or 24.0
         w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 1280)
@@ -348,7 +346,7 @@ def detect_video_safe(video_path: str, conf: float, iou: float, max_frames: int
         out_path = os.path.join(tempfile.gettempdir(), f"annotated_{int(time.time())}.mp4")
         writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
         if not writer or (hasattr(writer,"isOpened") and not writer.isOpened()):
-            return None, "{}", "❌ Video writer could not open. Try another format/resolution.", []
+            return None, "{}", "❌ Video writer could not open. Try another format/resolution.", [], _model_info_text()
 
         det_records: List[dict] = []
         frames = 0
@@ -364,7 +362,7 @@ def detect_video_safe(video_path: str, conf: float, iou: float, max_frames: int
                 _apply_english_overlay(r)
 
                 rows_raw = _results_to_rows(results)
-                rows = _filter_rows_by_geometry(r, rows_raw)
+                rows = _filter_rows_by_geometry(r, rows_raw, model_key)
 
                 t_sec = frames / float(fps if fps > 0 else 24.0)
                 for row in rows:
@@ -387,9 +385,9 @@ def detect_video_safe(video_path: str, conf: float, iou: float, max_frames: int
             counts[d["object"]] = counts.get(d["object"], 0) + 1
         summary = "Detections (tally): " + (", ".join(f"{k}: {v}" for k,v in sorted(counts.items())) if counts else "none")
         detections_json = json.dumps(det_records[:200], ensure_ascii=False, indent=2)
-        return out_path, detections_json, summary, det_records
+        return out_path, detections_json, summary, det_records, _model_info_text()
     except Exception as e:
-        return None, "{}", f"❌ Error during video detection: {e}", []
+        return None, "{}", f"❌ Error during video detection: {e}", [], _model_info_text()
 
 # ---------- PDF export ----------
 def export_pdf_img(det_records: List[dict], summary: str, annotated_tmp_jpg: Optional[str]):
@@ -414,13 +412,18 @@ def export_pdf_vid(det_records: List[dict], summary: str):
 # =========================
 NOTE = (
     "Detections include timestamp, object, confidence, and Threat/Non-threat. "
-    "Anti-noise filters active: min conf, min box area, sky-only region."
+    "Anti-noise filters: min conf, min box area, sky-only region. Choose model below."
 )
 
 with gr.Blocks(title="UAV / Drone Detector (YOLO)") as demo:
     gr.Markdown("# UAV / Drone Detection (Pretrained YOLO)")
     gr.Markdown("Embedded samples (optional): `samples/uav_image.jpg`, `samples/uav_video.mp4`.")
-    model_info_md = gr.Markdown(value=_model_info_text())
+
+    with gr.Row():
+        model_key = gr.Dropdown(choices=list(MODEL_CHOICES.keys()),
+                                value=list(MODEL_CHOICES.keys())[0],
+                                label="Model")
+        model_info_md = gr.Markdown(value=_model_info_text())
 
     with gr.Tabs():
         # IMAGE
@@ -445,13 +448,12 @@ with gr.Blocks(title="UAV / Drone Detector (YOLO)") as demo:
             annotated_tmp_img_path = gr.State(value=None)
             image_det_state = gr.State(value=[])
 
-            def _run_img(image, conf, iou):
-                result = detect_image_safe(image, conf, iou)
-                return (*result, _model_info_text())
+            def _run_img(mkey, image, conf, iou):
+                return detect_image_safe(mkey, image, conf, iou)
 
             run_img.click(
                 fn=_run_img,
-                inputs=[image_in, conf_img, iou_img],
+                inputs=[model_key, image_in, conf_img, iou_img],
                 outputs=[image_out, table_out, msg_img, image_det_state, annotated_tmp_img_path, model_info_md],
             )
 
@@ -482,13 +484,12 @@ with gr.Blocks(title="UAV / Drone Detector (YOLO)") as demo:
             pdf_vid_path = gr.File(label="PDF Report", interactive=False)
             video_det_state = gr.State(value=[])
 
-            def _run_vid(vpath, conf, iou, maxf):
-                result = detect_video_safe(vpath, conf, iou, int(maxf))
-                return (*result, _model_info_text())
+            def _run_vid(mkey, vpath, conf, iou, maxf):
+                return detect_video_safe(mkey, vpath, conf, iou, int(maxf))
 
             run_vid.click(
                 fn=_run_vid,
-                inputs=[video_in, conf_vid, iou_vid, max_frames],
+                inputs=[model_key, video_in, conf_vid, iou_vid, max_frames],
                 outputs=[video_out, detections_json_text, msg_vid, video_det_state, model_info_md],
             )