Update app.py

app.py

@@ -1,19 +1,17 @@
-""""
+"""
 =================================================================
-Disaster AI_HuggingFace Spaces API
-Final version_all fixes applied
+Disaster AI - HuggingFace Spaces API
+Final version - all models integrated
 =================================================================
 """
 
 import os
 import io
-import json
 import time
 import base64
 import threading
 import traceback
 import numpy as np
-from pathlib import Path
 from PIL import Image
 import cv2
 import torch
@@ -21,7 +19,6 @@ import requests
 
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
 from huggingface_hub import hf_hub_download
 
 # ════════════════════════════════
@@ -30,7 +27,7 @@ from huggingface_hub import hf_hub_download
 app = FastAPI(
     title="Disaster AI Inference API",
     description="Multi-model disaster scene analysis for Dokai / RoboXavier",
-    version="2.0.0",
+    version="3.0.0",
 )
 
 app.add_middleware(
@@ -41,13 +38,17 @@ app.add_middleware(
 )
 
 # ════════════════════════════════
-# Configuration
+# Configuration — all from secrets
 # ════════════════════════════════
-HF_VICTIM_MODEL_REPO = os.getenv("HF_VICTIM_MODEL_REPO", "")
+HF_TOKEN             = os.getenv("HF_TOKEN", None)
+HF_VICTIM_MODEL_REPO = os.getenv("HF_VICTIM_MODEL_REPO", "EgoisticCoderX/dokai-victim-detection")
+HF_XVIEW2_MODEL_REPO = os.getenv("HF_XVIEW2_MODEL_REPO", "EgoisticCoderX/dokai-xview2-damage")
 ROBOFLOW_API_KEY     = os.getenv("ROBOFLOW_API_KEY", "rltTa8UANpettqj6aHJG")
 MODEL_CACHE_DIR      = "/tmp/model_cache"
+
 os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
 
+# ── Victim detection class map ──
 TARGET_CLASSES = {
     0: "injured_civilian",
     1: "trapped_civilian",
@@ -62,14 +63,22 @@ CLASS_PRIORITY = {
     "rescue_personnel":  0.0,
 }
 
+# ── xView2 damage severity map ──
+DAMAGE_SEVERITY_ORDER = {
+    "destroyed":    0,
+    "major_damage": 1,
+    "minor_damage": 2,
+    "no_damage":    3,
+}
+
 # ════════════════════════════════
 # Model Registry
 # ════════════════════════════════
 class ModelRegistry:
     def __init__(self):
-        self._models  = {}
-        self._errors  = {}
-        self._lock    = threading.Lock()
+        self._models = {}
+        self._errors = {}
+        self._lock   = threading.Lock()
 
     def get(self, name):
         return self._models.get(name)
@@ -103,20 +112,17 @@ registry = ModelRegistry()
 # ════════════════════════════════
 
 def load_ladi_model():
-    """Load LADI-v2 classifier from HuggingFace Hub."""
+    """Load LADI-v2 scene classifier from HuggingFace Hub."""
     if registry.is_loaded("ladi"):
         return registry.get("ladi")
-
     try:
         from transformers import AutoImageProcessor, AutoModelForImageClassification
 
         print("⬇️  Loading MITLL/LADI-v2-classifier-small ...")
-
         processor = AutoImageProcessor.from_pretrained(
             "MITLL/LADI-v2-classifier-small",
             cache_dir=MODEL_CACHE_DIR,
         )
-
         model = AutoModelForImageClassification.from_pretrained(
             "MITLL/LADI-v2-classifier-small",
             cache_dir=MODEL_CACHE_DIR,
@@ -124,8 +130,6 @@ def load_ladi_model():
             ignore_mismatched_sizes=True,
         )
         model.eval()
-
-        # CPU only on HF free tier
         registry.register("ladi", {"model": model, "processor": processor})
         print("✅ LADI-v2 ready")
         return registry.get("ladi")
@@ -142,7 +146,7 @@ def load_victim_model():
         return registry.get("victim")
 
     if not HF_VICTIM_MODEL_REPO:
-        registry.set_error("victim", "HF_VICTIM_MODEL_REPO secret not set — train the model first")
+        registry.set_error("victim", "HF_VICTIM_MODEL_REPO secret not set")
         return None
 
     try:
@@ -153,6 +157,7 @@ def load_victim_model():
             repo_id=HF_VICTIM_MODEL_REPO,
             filename="best.pt",
             cache_dir=MODEL_CACHE_DIR,
+            token=HF_TOKEN,
         )
         model = YOLO(model_path)
         registry.register("victim", model)
@@ -165,33 +170,69 @@ def load_victim_model():
         return None
 
 
+def load_xview2_model():
+    """Load xView2 building damage YOLOv8 model from HuggingFace Hub."""
+    if registry.is_loaded("xview2"):
+        return registry.get("xview2")
+
+    if not HF_XVIEW2_MODEL_REPO:
+        registry.set_error("xview2", "HF_XVIEW2_MODEL_REPO secret not set")
+        return None
+
+    try:
+        from ultralytics import YOLO
+
+        print(f"⬇️  Loading xView2 model from {HF_XVIEW2_MODEL_REPO} ...")
+        model_path = hf_hub_download(
+            repo_id=HF_XVIEW2_MODEL_REPO,
+            filename="best.pt",
+            cache_dir=MODEL_CACHE_DIR,
+            token=HF_TOKEN,
+        )
+        model = YOLO(model_path)
+        registry.register("xview2", model)
+        print("✅ xView2 damage model ready")
+        return model
+
+    except Exception as e:
+        print(f"❌ xView2 model load failed:\n{traceback.format_exc()}")
+        registry.set_error("xview2", e)
+        return None
+
+
 # ════════════════════════════════
-# Startup — preload everything
+# Startup
 # ════════════════════════════════
 @app.on_event("startup")
 async def startup_event():
-    print("\n" + "="*50)
-    print("🚀 Disaster AI API starting up...")
-    print("="*50)
+    print("\n" + "=" * 55)
+    print("🚀 Disaster AI API v3.0 starting up...")
+    print("=" * 55)
 
-    # Always load LADI — it's a public HF model
+    # LADI always loads — public model
     load_ladi_model()
 
-    # Only load victim model if repo is configured
+    # Victim model — needs secret
     if HF_VICTIM_MODEL_REPO:
         load_victim_model()
     else:
         print("⚠️  Victim model skipped — HF_VICTIM_MODEL_REPO not set")
-        print("   Train the model first, then add the secret to this Space")
 
-    print("="*50)
-    print(f"📊 Registry status: {registry.status()}")
-    print("="*50 + "\n")
+    # xView2 model — needs secret
+    if HF_XVIEW2_MODEL_REPO:
+        load_xview2_model()
+    else:
+        print("⚠️  xView2 model skipped — HF_XVIEW2_MODEL_REPO not set")
+
+    print("=" * 55)
+    print(f"📊 Registry: {registry.status()}")
+    print("=" * 55 + "\n")
 
 
 # ════════════════════════════════
-# Utility
+# Utilities
 # ════════════════════════════════
+
 def read_image(file_bytes: bytes) -> np.ndarray:
     nparr = np.frombuffer(file_bytes, np.uint8)
     img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
@@ -239,7 +280,7 @@ def compute_triage(detections: list) -> dict:
             "total": 0, "critical": 0, "high": 0,
             "moderate": 0, "low": 0,
             "highest_score": 0.0,
-            "action": "✅ No victims detected",
+            "action": "No victims detected",
             "ranked_victims": [],
         }
 
@@ -265,10 +306,10 @@ def compute_triage(detections: list) -> dict:
     low      = sum(1 for d in scored if d["priority_rank"] == "LOW")
 
     action = (
-        "⚠️ IMMEDIATE RESCUE — Critical victims present"  if critical else
-        "🔴 Deploy rescue team — High priority victims"   if high     else
-        "🟡 Assess and triage — Moderate victims present" if moderate else
-        "🟢 Low priority — Monitor the area"
+        "IMMEDIATE RESCUE - Critical victims present"  if critical else
+        "Deploy rescue team - High priority victims"   if high     else
+        "Assess and triage - Moderate victims present" if moderate else
+        "Low priority - Monitor the area"
     )
 
     return {
@@ -283,6 +324,31 @@ def compute_triage(detections: list) -> dict:
     }
 
 
+def compute_zone_color(triage_data: dict, damage_counts: dict, top_class: str) -> str:
+    """
+    Unified zone color logic combining victim triage + building damage + scene class.
+    red > orange > yellow > green
+    """
+    critical        = triage_data.get("critical", 0)
+    high            = triage_data.get("high", 0)
+    destroyed       = damage_counts.get("destroyed", 0)
+    major_damage    = damage_counts.get("major_damage", 0)
+    minor_damage    = damage_counts.get("minor_damage", 0)
+    victim_total    = triage_data.get("total", 0)
+
+    scene_critical  = any(w in top_class for w in ["destroy", "collapse", "major"])
+    scene_moderate  = "minor_damage" in top_class
+
+    if critical > 0 or destroyed > 0 or scene_critical:
+        return "red"
+    elif high > 0 or major_damage > 0 or scene_moderate:
+        return "orange"
+    elif victim_total > 0 or minor_damage > 0:
+        return "yellow"
+    else:
+        return "green"
+
+
 # ════════════════════════════════
 # Routes
 # ════════════════════════════════
@@ -290,16 +356,17 @@ def compute_triage(detections: list) -> dict:
 @app.get("/")
 def root():
     return {
-        "service":   "Disaster AI Inference API",
-        "version":   "2.0.0",
-        "status":    registry.status(),
+        "service": "Disaster AI Inference API",
+        "version": "3.0.0",
+        "status":  registry.status(),
         "endpoints": {
             "GET  /health":          "Health check + model status",
-            "POST /classify":        "LADI-v2 scene classification",
+            "POST /classify":        "LADI-v2 disaster scene classification",
             "POST /detect/victims":  "Victim detection + triage priority",
-            "POST /detect/vehicles": "Emergency vehicle detection",
-            "POST /analyze/full":    "All models in one call",
-        }
+            "POST /detect/vehicles": "Emergency vehicle detection (Roboflow)",
+            "POST /detect/damage":   "xView2 building damage assessment",
+            "POST /analyze/full":    "All models in one call (main endpoint)",
+        },
     }
 
 
@@ -309,12 +376,18 @@ def health():
         "status":        "ok",
         "registry":      registry.status(),
         "gpu_available": torch.cuda.is_available(),
-        "timestamp":     time.time(),
+        "secrets_set": {
+            "HF_TOKEN":             HF_TOKEN is not None,
+            "HF_VICTIM_MODEL_REPO": bool(HF_VICTIM_MODEL_REPO),
+            "HF_XVIEW2_MODEL_REPO": bool(HF_XVIEW2_MODEL_REPO),
+            "ROBOFLOW_API_KEY":     bool(ROBOFLOW_API_KEY),
+        },
+        "timestamp": time.time(),
     }
 
 
 # ─────────────────────────────────────────────
-# LADI-v2 Classification
+# 1. LADI-v2 Scene Classification
 # ─────────────────────────────────────────────
 @app.post("/classify")
 async def classify_scene(
@@ -322,8 +395,8 @@ async def classify_scene(
     top_k: int = 5,
 ):
     """
-    Classify disaster scene using LADI-v2.
-    Returns top-k predicted damage categories with confidence scores.
+    Classify disaster scene using LADI-v2 (aerial damage categories).
+    Returns top-k predicted labels with confidence scores.
     """
     ladi = load_ladi_model()
     if ladi is None:
@@ -350,9 +423,9 @@ async def classify_scene(
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Inference failed: {e}")
 
-    elapsed = round((time.time() - t0) * 1000, 2)
+    elapsed  = round((time.time() - t0) * 1000, 2)
+    id2label = model.config.id2label
 
-    id2label  = model.config.id2label
     all_scores = sorted(
         [
             {
@@ -365,7 +438,6 @@ async def classify_scene(
         reverse=True,
     )
 
-    # Exclude water/flood from top predictions (not relevant for rover)
     relevant = [
         s for s in all_scores
         if not any(w in s["class"] for w in ["water", "flood"])
@@ -380,7 +452,7 @@ async def classify_scene(
 
 
 # ─────────────────────────────────────────────
-# Victim Detection
+# 2. Victim Detection + Triage
 # ─────────────────────────────────────────────
 @app.post("/detect/victims")
 async def detect_victims(
@@ -388,8 +460,9 @@ async def detect_victims(
     confidence: float = 0.35,
 ):
     """
-    Detect victims and classify by triage priority.
-    Returns CRITICAL / HIGH / MODERATE / LOW ranked detections.
+    Detect victims and rank by triage priority.
+    Classes: injured_civilian, trapped_civilian, safe_civilian, rescue_personnel.
+    Priority ranks: CRITICAL / HIGH / MODERATE / LOW
     """
     model = load_victim_model()
     if model is None:
@@ -418,11 +491,11 @@ async def detect_victims(
                 "class":      TARGET_CLASSES.get(cls_id, "unknown"),
                 "class_id":   cls_id,
                 "confidence": round(conf_val, 4),
-                "box": {"xmin": x1, "ymin": y1, "xmax": x2, "ymax": y2},
+                "box":        {"xmin": x1, "ymin": y1, "xmax": x2, "ymax": y2},
             })
 
-    triage   = compute_triage(raw)
-    victims  = triage.pop("ranked_victims", raw)
+    triage  = compute_triage(raw)
+    victims = triage.pop("ranked_victims", raw)
 
     return {
         "detections":        victims,
@@ -432,26 +505,25 @@ async def detect_victims(
 
 
 # ─────────────────────────────────────────────
-# Emergency Vehicle Detection
+# 3. Emergency Vehicle Detection (Roboflow)
 # ─────────────────────────────────────────────
 @app.post("/detect/vehicles")
 async def detect_vehicles(file: UploadFile = File(...)):
     """
-    Detect emergency vehicles using Roboflow.
-    Returns ambulance / fire truck / rescue vehicle detections.
+    Detect emergency vehicles via Roboflow hosted model.
+    Returns ambulance / fire truck detections and rescue_arrived flag.
     """
     if not ROBOFLOW_API_KEY:
         raise HTTPException(status_code=503, detail="ROBOFLOW_API_KEY secret not set")
 
-    contents = await file.read()
-    img = read_image(contents)
-
+    contents   = await file.read()
+    img        = read_image(contents)
     t0         = time.time()
     detections = call_roboflow(img, "ambulance-4bova/1", confidence=40)
     elapsed    = round((time.time() - t0) * 1000, 2)
 
-    has_ambulance  = any("ambulance"  in d["class"].lower() for d in detections)
-    has_fire_truck = any("fire"       in d["class"].lower() for d in detections)
+    has_ambulance  = any("ambulance" in d["class"].lower() for d in detections)
+    has_fire_truck = any("fire"      in d["class"].lower() for d in detections)
 
     return {
         "detections": detections,
@@ -465,30 +537,106 @@ async def detect_vehicles(file: UploadFile = File(...)):
 
 
 # ─────────────────────────────────────────────
-# Full Analysis — all models in one call
+# 4. xView2 Building Damage Assessment
+# ─────────────────────────────────────────────
+@app.post("/detect/damage")
+async def detect_building_damage(
+    file:       UploadFile = File(...),
+    confidence: float = 0.30,
+):
+    """
+    Assess building damage using xView2-trained YOLOv8.
+    Classes: destroyed, major_damage, minor_damage, no_damage.
+    Returns per-building detections, counts, and zone color.
+    """
+    model = load_xview2_model()
+    if model is None:
+        raise HTTPException(
+            status_code=503,
+            detail=f"xView2 model unavailable: {registry.get_error('xview2')}"
+        )
+
+    contents = await file.read()
+    img      = read_image(contents)
+
+    t0 = time.time()
+    try:
+        results = model.predict(source=img, conf=confidence, verbose=False)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Inference failed: {e}")
+    elapsed = round((time.time() - t0) * 1000, 2)
+
+    detections = []
+    counts     = {"destroyed": 0, "major_damage": 0, "minor_damage": 0, "no_damage": 0}
+
+    for r in results:
+        for box in r.boxes:
+            x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
+            conf_val   = float(box.conf[0])
+            cls_id     = int(box.cls[0])
+            class_name = model.names[cls_id].lower().replace(" ", "_")
+
+            # Map raw class name to standard severity key
+            matched_key = next(
+                (k for k in counts if k in class_name),
+                "no_damage"
+            )
+            counts[matched_key] += 1
+
+            detections.append({
+                "class":      class_name,
+                "confidence": round(conf_val, 4),
+                "severity":   matched_key,
+                "box":        {"xmin": x1, "ymin": y1, "xmax": x2, "ymax": y2},
+            })
+
+    # Sort: destroyed first, no_damage last
+    detections.sort(key=lambda d: DAMAGE_SEVERITY_ORDER.get(d["severity"], 9))
+
+    if counts["destroyed"] > 0:
+        zone_color = "red"
+    elif counts["major_damage"] > 0:
+        zone_color = "orange"
+    elif counts["minor_damage"] > 0:
+        zone_color = "yellow"
+    else:
+        zone_color = "green"
+
+    return {
+        "detections":        detections,
+        "summary":           counts,
+        "total_buildings":   sum(counts.values()),
+        "zone_color":        zone_color,
+        "inference_time_ms": elapsed,
+    }
+
+
+# ─────────────────────────────────────────────
+# 5. Full Analysis — all models in one call
 # ─────────────────────────────────────────────
 @app.post("/analyze/full")
 async def full_analysis(
     file:         UploadFile = File(...),
+    run_classify: bool = True,
     run_victims:  bool = True,
     run_vehicles: bool = True,
-    run_classify: bool = True,
+    run_damage:   bool = True,
 ):
     """
     Run all available models on one image.
-    This is the main endpoint your rover Flask app should call.
-
-    Returns unified JSON with zone color, all detections, triage summary.
+    Main endpoint for the RoboXavier rover Flask app.
+    Returns unified zone_color, all detections, and triage/damage summaries.
     """
     contents = await file.read()
     t_total  = time.time()
     output   = {}
 
-    # ── LADI classification ──
+    # ── LADI scene classification ──
     if run_classify:
         try:
-            fake_file = UploadFile(filename="f.jpg", file=io.BytesIO(contents))
-            output["classification"] = await classify_scene(fake_file)
+            output["classification"] = await classify_scene(
+                UploadFile(filename="f.jpg", file=io.BytesIO(contents))
+            )
         except HTTPException as e:
             output["classification"] = {"error": e.detail}
         except Exception as e:
@@ -497,39 +645,43 @@ async def full_analysis(
     # ── Victim detection ──
     if run_victims:
         try:
-            fake_file = UploadFile(filename="f.jpg", file=io.BytesIO(contents))
-            output["victims"] = await detect_victims(fake_file)
+            output["victims"] = await detect_victims(
+                UploadFile(filename="f.jpg", file=io.BytesIO(contents))
+            )
         except HTTPException as e:
             output["victims"] = {"error": e.detail}
         except Exception as e:
             output["victims"] = {"error": str(e)}
 
-    # ── Vehicle detection ──
+    # ── Emergency vehicle detection ──
     if run_vehicles:
         try:
-            fake_file = UploadFile(filename="f.jpg", file=io.BytesIO(contents))
-            output["vehicles"] = await detect_vehicles(fake_file)
+            output["vehicles"] = await detect_vehicles(
+                UploadFile(filename="f.jpg", file=io.BytesIO(contents))
+            )
         except HTTPException as e:
             output["vehicles"] = {"error": e.detail}
         except Exception as e:
             output["vehicles"] = {"error": str(e)}
 
-    # ── Zone color ──
-    triage_data  = output.get("victims",        {}).get("triage_summary", {})
-    classify_top = output.get("classification", {}).get("top_predictions", [{}])
-    top_class    = classify_top[0].get("class", "") if classify_top else ""
+    # ── xView2 building damage ──
+    if run_damage:
+        try:
+            output["building_damage"] = await detect_building_damage(
+                UploadFile(filename="f.jpg", file=io.BytesIO(contents))
+            )
+        except HTTPException as e:
+            output["building_damage"] = {"error": e.detail}
+        except Exception as e:
+            output["building_damage"] = {"error": str(e)}
 
-    critical = triage_data.get("critical", 0)
-    high     = triage_data.get("high",     0)
+    # ── Unified zone color (all signals combined) ──
+    triage_data   = output.get("victims",         {}).get("triage_summary", {})
+    damage_counts = output.get("building_damage", {}).get("summary", {})
+    classify_top  = output.get("classification",  {}).get("top_predictions", [{}])
+    top_class     = classify_top[0].get("class", "") if classify_top else ""
 
-    if critical > 0 or any(w in top_class for w in ["destroy", "collapse", "major"]):
-        zone_color = "red"
-    elif high > 0 or "minor_damage" in top_class:
-        zone_color = "orange"
-    elif triage_data.get("total", 0) > 0:
-        zone_color = "yellow"
-    else:
-        zone_color = "green"
+    zone_color = compute_zone_color(triage_data, damage_counts, top_class)
 
     return {
         "zone_color":    zone_color,