app.py

"""
AegisRoad v3.0 — YOLOv11 Inference Server
Hugging Face Spaces Deployment
"""

import os
import io
import base64
import time
import random
from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw
import cv2
import uvicorn
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles

# ── App setup ──────────────────────────────────────────────────────────────
app = FastAPI(
    title="AegisRoad Inference API",
    description="YOLOv11 road damage detection endpoint",
    version="3.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Mount static files BEFORE defining routes
if Path("static").exists():
    app.mount("/static", StaticFiles(directory="static"), name="static")

# ── Constants ──────────────────────────────────────────────────────────────
MODEL_PATH = Path("models/best.pt")

# Fallback class map if model doesn't embed names
FALLBACK_CLASS_NAMES = {0: "D00", 1: "D10", 2: "D20", 3: "D40"}

CLASS_LABELS = {
    "D00": "Longitudinal Crack",
    "D10": "Transverse Crack",
    "D20": "Alligator Cracking",
    "D40": "Pothole",
}
SEVERITY_MAP = {
    "D00": "low",
    "D10": "medium",
    "D20": "high",
    "D40": "critical",
}
CLASS_COLORS_RGB = {
    "D00": (255, 220, 0),
    "D10": (255, 140, 0),
    "D20": (255, 60, 0),
    "D40": (220, 0, 0),
}

# ── Model loading ──────────────────────────────────────────────────────────
model = None
model_class_names = FALLBACK_CLASS_NAMES.copy()

def load_model():
    global model, model_class_names
    if MODEL_PATH.exists():
        try:
            from ultralytics import YOLO
            model = YOLO(str(MODEL_PATH))
            # ✅ Read class names FROM the model itself — fixes "unknown" class bug
            raw_names = model.names  # e.g. {0: 'D00', 1: 'D10', ...}
            model_class_names = {k: str(v).upper() for k, v in raw_names.items()}
            print(f"✅ Model loaded from {MODEL_PATH}")
            print(f"✅ Model classes: {model_class_names}")
        except Exception as e:
            print(f"⚠️  Could not load model: {e}. Running in demo mode.")
            model = None
    else:
        print(f"⚠️  {MODEL_PATH} not found. Running in demo mode.")
        model = None

@app.on_event("startup")
async def startup_event():
    load_model()

# ── Helper: draw bounding boxes ────────────────────────────────────────────
def draw_boxes(image_pil: Image.Image, detections: list) -> Image.Image:
    draw = ImageDraw.Draw(image_pil)
    for det in detections:
        cls_name = det["class"]
        conf     = det["confidence"]
        x1, y1, x2, y2 = det["bbox"]
        color     = CLASS_COLORS_RGB.get(cls_name, (255, 255, 255))
        color_hex = "#{:02x}{:02x}{:02x}".format(*color)
        # Box
        draw.rectangle([x1, y1, x2, y2], outline=color_hex, width=3)
        # Label background + text
        label  = f"{cls_name} {conf:.0%}"
        text_y = max(y1 - 22, 0)
        draw.rectangle(
            [x1, text_y, x1 + len(label) * 8 + 6, text_y + 20],
            fill=(0, 0, 0)
        )
        draw.text((x1 + 3, text_y + 2), label, fill=color_hex)
    return image_pil

# ── Helper: PIL → base64 ───────────────────────────────────────────────────
def pil_to_b64(img: Image.Image, fmt: str = "JPEG") -> str:
    buf = io.BytesIO()
    img.save(buf, format=fmt, quality=88)
    return base64.b64encode(buf.getvalue()).decode()

# ── Demo mode fallback ─────────────────────────────────────────────────────
def demo_detections(w: int, h: int) -> list:
    candidates = [
        {"class": "D40", "confidence": round(random.uniform(0.72, 0.94), 2),
         "bbox": [int(w*0.15), int(h*0.35), int(w*0.42), int(h*0.65)]},
        {"class": "D00", "confidence": round(random.uniform(0.55, 0.78), 2),
         "bbox": [int(w*0.50), int(h*0.20), int(w*0.85), int(h*0.45)]},
        {"class": "D20", "confidence": round(random.uniform(0.60, 0.80), 2),
         "bbox": [int(w*0.05), int(h*0.60), int(w*0.30), int(h*0.90)]},
    ]
    return random.sample(candidates, random.randint(1, 3))

# ── Routes ─────────────────────────────────────────────────────────────────

# ✅ FIX: serve the HTML frontend at root instead of JSON
@app.get("/")
async def root():
    return FileResponse("static/index.html")

@app.get("/health")
async def health():
    return {
        "status":       "ok",
        "model_loaded": model is not None,
        "mode":         "live" if model is not None else "demo",
    }

# ✅ NEW: debug endpoint — tells you exactly what class names your model uses
@app.get("/debug")
async def debug():
    if model is None:
        return {"model_loaded": False, "mode": "demo"}
    return {
        "model_loaded":    True,
        "mode":            "live",
        "model_path":      str(MODEL_PATH),
        "model_classes":   model_class_names,
        "num_classes":     len(model_class_names),
        "model_task":      model.task,
    }

@app.get("/classes")
async def get_classes():
    return {
        "classes": [
            {
                "id":       k,
                "code":     v,
                "label":    CLASS_LABELS.get(v, v),
                "severity": SEVERITY_MAP.get(v, "unknown"),
            }
            for k, v in model_class_names.items()
        ]
    }

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File must be an image.")

    start    = time.time()
    contents = await file.read()

    try:
        image = Image.open(io.BytesIO(contents)).convert("RGB")
    except Exception:
        raise HTTPException(status_code=400, detail="Could not read image file.")

    w, h = image.size

    # ── Live inference ────────────────────────────────────────────────────
    if model is not None:
        # ✅ Tuned thresholds: conf=0.30 reduces noise, iou=0.4 reduces duplicate boxes
        results    = model(image, conf=0.30, iou=0.4)[0]
        detections = []
        for box in results.boxes:
            cls_id = int(box.cls[0])
            # ✅ Use model's own embedded class names — fixes "unknown" bug
            cls_name = model_class_names.get(cls_id, f"cls_{cls_id}")
            conf     = float(box.conf[0])
            x1, y1, x2, y2 = [round(v) for v in box.xyxy[0].tolist()]
            detections.append({
                "class":      cls_name,
                "label":      CLASS_LABELS.get(cls_name, cls_name),
                "confidence": round(conf, 4),
                "severity":   SEVERITY_MAP.get(cls_name, "low"),
                "bbox":       [x1, y1, x2, y2],
            })
        demo = False

    # ── Demo fallback ─────────────────────────────────────────────────────
    else:
        detections = demo_detections(w, h)
        for d in detections:
            d["label"]    = CLASS_LABELS.get(d["class"], d["class"])
            d["severity"] = SEVERITY_MAP.get(d["class"], "unknown")
        demo = True

    # ── Annotate image ────────────────────────────────────────────────────
    annotated     = draw_boxes(image.copy(), detections)
    annotated_b64 = pil_to_b64(annotated)
    original_b64  = pil_to_b64(image)

    elapsed = round(time.time() - start, 3)

    # ── Summary stats ─────────────────────────────────────────────────────
    severity_counts = {"low": 0, "medium": 0, "high": 0, "critical": 0}
    for d in detections:
        sev = d["severity"]
        if sev in severity_counts:
            severity_counts[sev] += 1

    road_score = max(0, 100 - (
        severity_counts["critical"] * 30 +
        severity_counts["high"]     * 20 +
        severity_counts["medium"]   * 10 +
        severity_counts["low"]      *  5
    ))

    return JSONResponse({
        "success":           True,
        "demo_mode":         demo,
        "inference_ms":      int(elapsed * 1000),
        "image_size":        {"width": w, "height": h},
        "detections":        detections,
        "detection_count":   len(detections),
        "severity_summary":  severity_counts,
        "road_health_score": road_score,
        "annotated_image":   annotated_b64,
        "original_image":    original_b64,
    })