src/api.py

"""
FastAPI service wrapping the AiGiS pipeline for use from a website or another app.

    POST /caption  (multipart file upload: "image")
        -> runs YOLO (detector.py) + a VLM (vlm_caption.py) on the image
        -> returns JSON: detections, compliance findings, VLM caption/observation,
           and the annotated image (boxes drawn) as base64 JPEG

    Query params:
      skip_vlm      (bool)  detector only, no VLM call (no API cost)
      skip_detector (bool)  VLM only — image judged without YOLO detections
      model         (str)   registry key from vlm_caption.MODEL_REGISTRY
                            ("sonnet", "opus", "haiku", "qwen-8b", "qwen-30b",
                            "qwen-32b") or a raw model id. Defaults to "sonnet".

The YOLO weights are pulled once at startup from the Hugging Face model repo

Run:
    cd src
    uvicorn api:app --reload --host 0.0.0.0 --port 8000
"""
from __future__ import annotations

import base64
import mimetypes
import os
from contextlib import asynccontextmanager
from pathlib import Path

import cv2
import numpy as np
from dotenv import load_dotenv
from fastapi import FastAPI, File, Header, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import HfApi, hf_hub_download

import vlm_caption
from detector import SafetyDetector

REPO_ROOT = Path(__file__).resolve().parent.parent
load_dotenv(REPO_ROOT / ".env")

HF_REPO_ID = os.environ.get("AIGIS_HF_REPO_ID", "ZahraFatima/aigis-yolo")
HF_FILENAME = os.environ.get("AIGIS_HF_FILENAME", "best.pt")
DATA_YAML = os.environ.get("AIGIS_DATA_YAML", str(REPO_ROOT / "config" / "sh17.yaml"))
CONF = float(os.environ.get("AIGIS_CONF", "0.25"))

# For the /admin/restart endpoint only (this Space, not the local/AWS deployment):
# HF_TOKEN needs "write" access to THIS_SPACE_ID, and RESTART_SECRET is a shared
# password the BE team uses instead of needing their own HF account.
HF_TOKEN = os.environ.get("HF_TOKEN")
THIS_SPACE_ID = os.environ.get("SPACE_ID", "DarkthornLabs1/aigis-api")
RESTART_SECRET = os.environ.get("RESTART_SECRET")

# Which env var each VLM provider needs, so a missing key gives a clean 500
# instead of a raw SDK exception.
_PROVIDER_KEY = {"anthropic": "ANTHROPIC_API_KEY", "dashscope": "DASHSCOPE_API_KEY"}

detector: SafetyDetector | None = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global detector
    weights_path = hf_hub_download(repo_id=HF_REPO_ID, filename=HF_FILENAME)
    detector = SafetyDetector(weights_path, DATA_YAML, conf=CONF)
    yield


app = FastAPI(title="AiGiS Safety Captioning API", version="0.1.0", lifespan=lifespan)

# Permissive by default so any site/MVP can call this during development.
# For production, set AIGIS_CORS_ORIGINS to a comma-separated allowlist.
app.add_middleware(
    CORSMiddleware,
    allow_origins=os.environ.get("AIGIS_CORS_ORIGINS", "*").split(","),
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/health")
def health():
    return {"status": "ok", "model_loaded": detector is not None}


@app.post("/admin/restart")
def restart(x_restart_token: str = Header(default="")):
    if not RESTART_SECRET or x_restart_token != RESTART_SECRET:
        raise HTTPException(status_code=403, detail="Invalid or missing X-Restart-Token header")
    if not HF_TOKEN:
        raise HTTPException(status_code=500, detail="HF_TOKEN not configured on server")
    try:
        HfApi(token=HF_TOKEN).restart_space(repo_id=THIS_SPACE_ID)
    except Exception as exc:
        raise HTTPException(status_code=502, detail=f"Restart request failed: {exc}") from exc
    return {"status": "restart triggered"}


@app.post("/caption")
async def caption_image(
    image: UploadFile = File(...),
    skip_vlm: bool = False,
    skip_detector: bool = False,
    model: str = vlm_caption.DEFAULT_MODEL,
):
    if detector is None:
        raise HTTPException(status_code=503, detail="Model not loaded yet")

    raw = await image.read()
    if not raw:
        raise HTTPException(status_code=400, detail="Empty file upload")

    frame = cv2.imdecode(np.frombuffer(raw, dtype=np.uint8), cv2.IMREAD_COLOR)
    if frame is None:
        raise HTTPException(status_code=400, detail="Could not decode image")

    record: dict = {}
    det_dict: dict | None = None
    annotated = frame

    if skip_detector:
        record["detector"] = {"note": "detector skipped (skip_detector=true)"}
    else:
        det_result, annotated = detector.detect_and_annotate(frame, name=image.filename or "upload.jpg")
        det_dict = det_result.to_dict()
        record["detector"] = det_dict

    if not skip_vlm:
        _, provider, _ = vlm_caption._resolve(model)
        needed_key = _PROVIDER_KEY.get(provider)
        if needed_key and needed_key not in os.environ:
            raise HTTPException(
                status_code=500,
                detail=f"{needed_key} not set on server (required for model '{model}')",
            )
        media_type = image.content_type or mimetypes.guess_type(image.filename or "")[0] or "image/jpeg"
        record["vlm"] = vlm_caption.caption_bytes(
            raw, media_type, detection_result=det_dict, model=model, use_detector=not skip_detector
        )

    ok, buf = cv2.imencode(".jpg", annotated)
    if not ok:
        raise HTTPException(status_code=500, detail="Failed to encode annotated image")
    record["annotated_image_base64"] = base64.b64encode(buf.tobytes()).decode()

    return record


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)