import io
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from PIL import Image
import torch
from ultralytics import YOLO

app = FastAPI()
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
DEV = 0 if torch.cuda.is_available() else "cpu"
model = YOLO("yolo11x.pt")
KEEP = {"boat": "SHIP", "person": "HUMAN"}
try:
    model.predict(Image.new("RGB", (640, 384)), device=DEV, verbose=False)  # warmup
except Exception:
    pass

@app.get("/")
def root():
    return {"ok": True, "model": "yolo11x", "imgsz": 1536, "conf": 0.15, "iou": 0.5, "device": str(DEV), "cuda": torch.cuda.is_available()}

@app.post("/detect")
async def detect(req: Request):
    data = await req.body()
    img = Image.open(io.BytesIO(data)).convert("RGB")
    W, H = img.size
    r = model.predict(img, conf=0.15, iou=0.5, imgsz=1536, max_det=100, device=DEV, verbose=False)[0]
    dets = []
    for b in r.boxes:
        c = model.names[int(b.cls)]
        if c not in KEEP:
            continue
        x1, y1, x2, y2 = [float(v) for v in b.xyxy[0]]
        dets.append({"label": KEEP[c], "score": round(float(b.conf), 3),
                     "box": {"xmin": x1, "ymin": y1, "xmax": x2, "ymax": y2}})
    return JSONResponse({"w": W, "h": H, "dets": dets})