Create app.py

app.py

@@ -0,0 +1,487 @@
+"""
+=================================================================
+DISASTER AI — HuggingFace Spaces API (Permanent Free Server)
+=================================================================
+Deploy this to HuggingFace Spaces for a permanent, always-on
+API that your friends can call without knowing any AI.
+
+Setup:
+1. Go to https://huggingface.co/spaces
+2. Create New Space → SDK: Docker (or Gradio)
+3. Repo name: EgoisticCoderX/dokai-inference-api
+4. Upload this file as app.py
+5. Upload requirements.txt
+6. Add secrets in Space Settings:
+   - HF_VICTIM_MODEL_REPO = EgoisticCoderX/dokai-victim-detection
+   - ROBOFLOW_API_KEY = your_key
+
+Your API will be at:
+  https://egoisticcoderx-dokai-inference-api.hf.space/
+
+ZeroGPU Note: HF Spaces has ZeroGPU (A10G, free) for inference.
+Add @spaces.GPU decorator for GPU-accelerated endpoints.
+=================================================================
+"""
+
+import os
+import io
+import json
+import time
+import base64
+import threading
+import numpy as np
+from pathlib import Path
+from PIL import Image
+import cv2
+import torch
+import requests
+
+from fastapi import FastAPI, File, UploadFile, HTTPException, Depends
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+from huggingface_hub import hf_hub_download
+
+# ── ZeroGPU support (HuggingFace Spaces) ──
+try:
+    import spaces
+    HAS_ZERO_GPU = True
+    print("✅ ZeroGPU available")
+except ImportError:
+    HAS_ZERO_GPU = False
+    print("ℹ️  ZeroGPU not available (running locally or non-GPU space)")
+
+
+# ════════════════════════════════
+# App Setup
+# ════════════════════════════════
+app = FastAPI(
+    title="Disaster AI Inference API",
+    description="Multi-model disaster scene analysis API for Dokai/RoboXavier",
+    version="1.0.0",
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ════════════════════════════════
+# Configuration
+# ════════════════════════════════
+HF_VICTIM_MODEL_REPO = os.getenv("HF_VICTIM_MODEL_REPO", "EgoisticCoderX/dokai-victim-detection")
+ROBOFLOW_API_KEY     = os.getenv("ROBOFLOW_API_KEY", "")
+MODEL_CACHE_DIR      = "/tmp/model_cache"
+os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
+
+TARGET_CLASSES = {
+    0: "injured_civilian",
+    1: "trapped_civilian",
+    2: "safe_civilian",
+    3: "rescue_personnel",
+}
+
+CLASS_PRIORITY = {
+    "injured_civilian":  1.0,
+    "trapped_civilian":  0.95,
+    "safe_civilian":     0.3,
+    "rescue_personnel":  0.0,
+}
+
+
+# ════════════════════════════════
+# Model Registry — lazy loading
+# ════════════════════════════════
+class ModelRegistry:
+    """
+    Lazy-loads models on first request.
+    Prevents OOM on Space startup.
+    """
+    def __init__(self):
+        self._models    = {}
+        self._lock      = threading.Lock()
+        self._loading   = set()
+
+    def get(self, name: str):
+        if name in self._models:
+            return self._models[name]
+        return None
+
+    def register(self, name: str, model):
+        with self._lock:
+            self._models[name] = model
+        print(f"✅ Model registered: {name}")
+
+    def is_loaded(self, name: str) -> bool:
+        return name in self._models
+
+
+registry = ModelRegistry()
+
+
+def load_victim_model():
+    """Download + load YOLOv8 victim detection model from HuggingFace."""
+    if registry.is_loaded("victim"):
+        return registry.get("victim")
+
+    try:
+        from ultralytics import YOLO
+        model_path = hf_hub_download(
+            repo_id=HF_VICTIM_MODEL_REPO,
+            filename="best.pt",
+            cache_dir=MODEL_CACHE_DIR,
+        )
+        model = YOLO(model_path)
+        registry.register("victim", model)
+        return model
+    except Exception as e:
+        print(f"❌ Failed to load victim model: {e}")
+        return None
+
+
+def load_ladi_model():
+    """Load LADI-v2 classifier from HuggingFace."""
+    if registry.is_loaded("ladi"):
+        return registry.get("ladi")
+
+    try:
+        from transformers import AutoImageProcessor, AutoModelForImageClassification
+        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,
+        )
+        model.eval()
+        registry.register("ladi", {"model": model, "processor": processor})
+        return registry.get("ladi")
+    except Exception as e:
+        print(f"❌ Failed to load LADI model: {e}")
+        return None
+
+
+# ════════════════════════════════
+# Utility Functions
+# ════════════════════════════════
+def read_image_from_upload(file_bytes: bytes) -> np.ndarray:
+    nparr = np.frombuffer(file_bytes, np.uint8)
+    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+    if img is None:
+        raise HTTPException(status_code=400, detail="Invalid image — cannot decode")
+    return img
+
+
+def call_roboflow(image: np.ndarray, model_id: str, confidence: int = 40) -> list:
+    """Call Roboflow hosted model (ambulance, vest detection)."""
+    if not ROBOFLOW_API_KEY:
+        return []
+    try:
+        _, buffer = cv2.imencode('.jpg', image, [cv2.IMWRITE_JPEG_QUALITY, 80])
+        img_b64 = base64.b64encode(buffer)
+        url = f"https://detect.roboflow.com/{model_id}?api_key={ROBOFLOW_API_KEY}&confidence={confidence}"
+        res = requests.post(
+            url,
+            data=img_b64,
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+            timeout=8,
+        )
+        res.raise_for_status()
+        preds = res.json().get("predictions", [])
+        return [
+            {
+                "class": p["class"],
+                "confidence": round(p["confidence"], 4),
+                "box": {
+                    "xmin": int(p["x"] - p["width"] / 2),
+                    "ymin": int(p["y"] - p["height"] / 2),
+                    "xmax": int(p["x"] + p["width"] / 2),
+                    "ymax": int(p["y"] + p["height"] / 2),
+                },
+            }
+            for p in preds
+        ]
+    except Exception as e:
+        print(f"Roboflow error ({model_id}): {e}")
+        return []
+
+
+def compute_triage(detections: list) -> dict:
+    """Compute triage priority summary from victim detections."""
+    if not detections:
+        return {
+            "total": 0,
+            "critical": 0,
+            "high": 0,
+            "moderate": 0,
+            "low": 0,
+            "highest_score": 0.0,
+            "action": "✅ No victims detected",
+        }
+
+    scored = []
+    for d in detections:
+        cls_name = d.get("class", "")
+        conf = d.get("confidence", 0.5)
+        weight = CLASS_PRIORITY.get(cls_name, 0.5)
+        score = conf * weight
+        rank = (
+            "CRITICAL" if score >= 0.7
+            else "HIGH" if score >= 0.4
+            else "MODERATE" if score >= 0.2
+            else "LOW"
+        )
+        scored.append({**d, "priority_score": round(score, 4), "priority_rank": rank})
+
+    scored.sort(key=lambda x: x["priority_score"], reverse=True)
+
+    critical = sum(1 for d in scored if d["priority_rank"] == "CRITICAL")
+    high     = sum(1 for d in scored if d["priority_rank"] == "HIGH")
+    moderate = sum(1 for d in scored if d["priority_rank"] == "MODERATE")
+    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"
+    )
+
+    return {
+        "total":          len(scored),
+        "critical":       critical,
+        "high":           high,
+        "moderate":       moderate,
+        "low":            low,
+        "highest_score":  scored[0]["priority_score"] if scored else 0.0,
+        "action":         action,
+        "ranked_victims": scored,
+    }
+
+
+# ════════════════════════════════
+# Routes
+# ════════════════════════════════
+
+@app.get("/")
+def root():
+    return {
+        "service": "Disaster AI Inference API",
+        "version": "1.0.0",
+        "endpoints": {
+            "/health":          "GET  — Service health check",
+            "/detect/victims":  "POST — Victim detection + triage",
+            "/detect/hazards":  "POST — Fire, smoke, building damage",
+            "/detect/vehicles": "POST — Emergency vehicle detection",
+            "/classify":        "POST — LADI-v2 scene classification",
+            "/analyze/full":    "POST — All models in parallel (full analysis)",
+        }
+    }
+
+
+@app.get("/health")
+def health():
+    return {
+        "status": "ok",
+        "models_loaded": {
+            "victim_detection": registry.is_loaded("victim"),
+            "ladi_classifier":  registry.is_loaded("ladi"),
+        },
+        "gpu_available": torch.cuda.is_available(),
+        "timestamp": time.time(),
+    }
+
+
+@app.post("/detect/victims")
+async def detect_victims(
+    file: UploadFile = File(...),
+    confidence: float = 0.35,
+):
+    """
+    Detect victims and classify by triage priority.
+    Returns detections with priority scores and recommended action.
+    """
+    contents = await file.read()
+    img = read_image_from_upload(contents)
+
+    model = load_victim_model()
+    if model is None:
+        raise HTTPException(status_code=503, detail="Victim detection model unavailable")
+
+    t0 = time.time()
+    results = model.predict(source=img, conf=confidence, verbose=False)
+    elapsed = round((time.time() - t0) * 1000, 2)
+
+    raw_detections = []
+    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])
+            raw_detections.append({
+                "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},
+            })
+
+    triage = compute_triage(raw_detections)
+
+    return {
+        "detections":       triage.pop("ranked_victims", raw_detections),
+        "triage_summary":   triage,
+        "inference_time_ms": elapsed,
+    }
+
+
+@app.post("/detect/vehicles")
+async def detect_vehicles(file: UploadFile = File(...)):
+    """Detect emergency vehicles using Roboflow model."""
+    contents = await file.read()
+    img = read_image_from_upload(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)
+
+    return {
+        "detections":           detections,
+        "emergency_vehicles": {
+            "ambulance_detected":  has_ambulance,
+            "fire_truck_detected": has_fire_truck,
+            "rescue_arrived":      has_ambulance or has_fire_truck,
+        },
+        "inference_time_ms": elapsed,
+    }
+
+
+@app.post("/classify")
+async def classify_scene(file: UploadFile = File(...), top_k: int = 5):
+    """Classify disaster scene using LADI-v2."""
+    contents = await file.read()
+    img_pil = Image.open(io.BytesIO(contents)).convert("RGB")
+
+    ladi = load_ladi_model()
+    if ladi is None:
+        raise HTTPException(status_code=503, detail="LADI-v2 model unavailable")
+
+    model     = ladi["model"]
+    processor = ladi["processor"]
+
+    t0 = time.time()
+    inputs = processor(images=img_pil, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
+    elapsed = round((time.time() - t0) * 1000, 2)
+
+    id2label = model.config.id2label
+    all_scores = sorted(
+        [
+            {
+                "class": id2label[i].lower().replace(" ", "_"),
+                "confidence": round(float(probs[i]), 4),
+            }
+            for i in range(len(probs))
+        ],
+        key=lambda x: x["confidence"],
+        reverse=True,
+    )
+
+    return {
+        "top_predictions":   all_scores[:top_k],
+        "all_scores":        all_scores,
+        "inference_time_ms": elapsed,
+    }
+
+
+@app.post("/analyze/full")
+async def full_analysis(
+    file: UploadFile = File(...),
+    run_victims:  bool = True,
+    run_vehicles: bool = True,
+    run_classify: bool = True,
+):
+    """
+    Run all available models in parallel on one image.
+    This is what your rover Flask app should call for full scene analysis.
+
+    Returns a unified JSON with all detections + zone scoring.
+    """
+    import asyncio
+    from concurrent.futures import ThreadPoolExecutor
+
+    contents = await file.read()
+    t_total = time.time()
+
+    results = {}
+
+    with ThreadPoolExecutor(max_workers=3) as executor:
+        futures = {}
+
+        if run_victims:
+            async def _victims():
+                f = UploadFile(filename="frame.jpg", file=io.BytesIO(contents))
+                return await detect_victims(f)
+            futures["victims"] = asyncio.ensure_future(_victims())
+
+        if run_vehicles:
+            async def _vehicles():
+                f = UploadFile(filename="frame.jpg", file=io.BytesIO(contents))
+                return await detect_vehicles(f)
+            futures["vehicles"] = asyncio.ensure_future(_vehicles())
+
+        if run_classify:
+            async def _classify():
+                f = UploadFile(filename="frame.jpg", file=io.BytesIO(contents))
+                return await classify_scene(f)
+            futures["classification"] = asyncio.ensure_future(_classify())
+
+        for key, fut in futures.items():
+            try:
+                results[key] = await fut
+            except Exception as e:
+                results[key] = {"error": str(e)}
+
+    total_ms = round((time.time() - t_total) * 1000, 2)
+
+    # Determine overall zone color
+    victim_data   = results.get("victims", {})
+    triage        = victim_data.get("triage_summary", {})
+    classify_data = results.get("classification", {})
+
+    critical = triage.get("critical", 0)
+    high     = triage.get("high", 0)
+
+    top_class = ""
+    if classify_data.get("top_predictions"):
+        top_class = classify_data["top_predictions"][0].get("class", "")
+
+    if critical > 0 or "destroy" in top_class or "collapse" in top_class:
+        zone_color = "red"
+    elif high > 0 or "major_damage" in top_class:
+        zone_color = "orange"
+    elif triage.get("total", 0) > 0:
+        zone_color = "yellow"
+    else:
+        zone_color = "green"
+
+    return {
+        "zone_color":       zone_color,
+        "results":          results,
+        "total_time_ms":    total_ms,
+        "timestamp":        time.time(),
+    }
+
+
+# ════════════════════════════════
+# Entry Point (local testing)
+# ════════════════════════════════
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)  # HF Spaces default port