Spaces:
Sleeping
Sleeping
| """ | |
| 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 | |
| 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=["*"], | |
| ) | |
| def health(): | |
| return {"status": "ok", "model_loaded": detector is not None} | |
| 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"} | |
| 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) | |