app.py

import json
import logging
import os
import subprocess
import tempfile
import time
import uuid
from pathlib import Path

from fastapi import FastAPI, File, Form, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("locate-anything-server")

app = FastAPI(title="LocateAnything Detection API")

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

MODEL_PATH = os.environ.get("MODEL_PATH", "/app/models/locate-anything-q4_k.gguf")
INFERENCE_MODE = os.environ.get("INFERENCE_MODE", "hybrid")
N_THREADS = os.environ.get("N_THREADS", "2")

_available = Path(MODEL_PATH).exists()


class DetectionResult(BaseModel):
    label: str
    box: list[float]  # [x1, y1, x2, y2] normalized 0-1


class DetectResponse(BaseModel):
    success: bool
    detections: list[DetectionResult]
    raw_text: str | None = None
    elapsed_s: float | None = None
    error: str | None = None


class HealthResponse(BaseModel):
    status: str
    model_loaded: bool
    model_path: str


@app.get("/health", response_model=HealthResponse)
async def health():
    return HealthResponse(
        status="ok" if _available else "no_model",
        model_loaded=_available,
        model_path=MODEL_PATH,
    )


@app.post("/detect", response_model=DetectResponse)
async def detect(
    image: UploadFile = File(...),
    prompt: str = Form("Locate all objects in this image."),
):
    if not _available:
        raise HTTPException(
            status_code=503,
            detail=f"Model not found at {MODEL_PATH}. Has it been downloaded?",
        )

    ext = Path(image.filename or "image.jpg").suffix or ".jpg"
    tmp_path = None
    try:
        tmp_dir = Path(tempfile.gettempdir()) / "locate-anything"
        tmp_dir.mkdir(parents=True, exist_ok=True)

        tmp_path = tmp_dir / f"{uuid.uuid4().hex}{ext}"
        with open(tmp_path, "wb") as f:
            f.write(await image.read())

        t0 = time.time()
        result = subprocess.run(
            [
                "locate-anything-cli", "detect",
                "--model", MODEL_PATH,
                "--input", str(tmp_path),
                "--prompt", prompt,
                "--mode", INFERENCE_MODE,
                "--threads", N_THREADS,
            ],
            capture_output=True,
            text=True,
            timeout=300,
        )
        elapsed = time.time() - t0

        if result.returncode != 0:
            logger.error(f"CLI stderr: {result.stderr[:500]}")
            return DetectResponse(
                success=False,
                detections=[],
                raw_text=result.stderr[:500],
                elapsed_s=elapsed,
                error=f"CLI exited with code {result.returncode}",
            )

        output = result.stdout.strip()
        data = json.loads(output) if output else {}

        raw_detections = data.get("detections", [])
        detections = [
            DetectionResult(
                label=d.get("label", "object"),
                box=_normalize_box(d.get("box", [0, 0, 0, 0])),
            )
            for d in raw_detections
        ]

        logger.info(f"Detected {len(detections)} objects in {elapsed:.1f}s")
        return DetectResponse(
            success=True,
            detections=detections,
            raw_text=output[:500],
            elapsed_s=elapsed,
        )

    except subprocess.TimeoutExpired:
        logger.error("Inference timed out after 120s")
        return DetectResponse(
            success=False, detections=[], error="Inference timed out"
        )
    except Exception as e:
        logger.error(f"Detection error: {e}")
        return DetectResponse(
            success=False, detections=[], error=str(e)
        )
    finally:
        if tmp_path and tmp_path.exists():
            tmp_path.unlink()


def _normalize_box(box: list[float]) -> list[float]:
    if len(box) != 4:
        return [0, 0, 0, 0]
    x1, y1, x2, y2 = box
    if max(x1, y1, x2, y2) > 1.0:
        x1 /= 1000.0
        y1 /= 1000.0
        x2 /= 1000.0
        y2 /= 1000.0
    return [x1, y1, x2, y2]