app.py

"""
app.py — FastAPI wrapper for the traffic violation detection pipeline.
Deployed on HuggingFace Spaces (CPU Docker) at port 7860.

Endpoints:
  GET  /health   — liveness check
  POST /predict  — upload image, get JSON violations + all_vehicles + summary
"""

import os
from contextlib import asynccontextmanager
from pathlib import Path

import cv2
import numpy as np
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.responses import JSONResponse
from ultralytics import YOLO
import uvicorn

# ── ensure all model paths resolve from the app's working directory ───────────
os.chdir(Path(__file__).parent)

from enhanced_pipeline import (
    COCO_MODEL, S1_LOCAL, S3_LOCAL, S4_LOCAL,
    load_depth_model, run_pipeline,
)

# ── global model store (loaded once at startup) ───────────────────────────────
_models: dict = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load all models on startup; release on shutdown."""
    print("[*] Loading depth model (DepthAnythingV2)...")
    load_depth_model()

    print("[*] Loading YOLO models...")
    _models["coco"]   = YOLO(COCO_MODEL)
    _models["s1"]     = YOLO(str(S1_LOCAL))
    _models["helmet"] = YOLO(str(S3_LOCAL))
    _models["plate"]  = YOLO(str(S4_LOCAL))
    print("[✓] All models ready. API is live.\n")

    yield  # app runs here

    _models.clear()
    print("[*] Models released.")


# ── app ───────────────────────────────────────────────────────────────────────
app = FastAPI(
    title="Traffic Violation Detection API",
    description=(
        "5-stage cascaded pipeline: person/bike detection → depth-filtered "
        "spatial association → helmet classification → license plate OCR. "
        "Returns per-vehicle violation data in JSON."
    ),
    version="1.0.0",
    lifespan=lifespan,
)


@app.get("/health")
def health():
    """Liveness check — returns ok once models are loaded."""
    if len(_models) < 4:
        raise HTTPException(503, "Models not ready yet")
    return {"status": "ok", "models_loaded": list(_models.keys())}


@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    """
    Upload an image (jpg/png) and receive a full violation report.

    Response shape:
    {
      "violations": [
        { "num_riders": int, "helmet_violations": int, "license_plate": str|null }
      ],
      "all_vehicles": [
        {
          "bike_bbox":        [x1,y1,x2,y2],
          "num_riders":       int,
          "with_helmet":      int,
          "without_helmet":   int,
          "triple_riding":    bool,
          "helmet_violation": bool,
          "license_plate":    str|null,
          "plate_bbox":       [x1,y1,x2,y2]|null,
          "riders": [
            { "person_bbox":[x1,y1,x2,y2], "head_bbox":[x1,y1,x2,y2], "helmet": str }
          ],
          "is_violation": bool
        }
      ],
      "summary": { "total_bikes": int, "total_violations": int }
    }
    """
    # Validate content type
    if not file.content_type.startswith("image/"):
        raise HTTPException(400, f"Expected an image file, got: {file.content_type}")

    # Decode uploaded bytes → numpy BGR image
    raw = await file.read()
    arr = np.frombuffer(raw, np.uint8)
    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if img is None:
        raise HTTPException(400, "Could not decode image — ensure it is a valid JPEG/PNG.")

    # Run full pipeline (CPU, no save/display)
    result = run_pipeline(
        img,
        _models["coco"],
        _models["s1"],
        _models["helmet"],
        _models["plate"],
        save=False,
        debug=False,
    )

    if "error" in result:
        raise HTTPException(500, result["error"])

    return JSONResponse(content=result)


if __name__ == "__main__":
    uvicorn.run("app:app", host="0.0.0.0", port=7860, log_level="info")