main.py

import io, os, sys
from typing import List, Dict, Any
from fastapi import FastAPI, UploadFile, File, Query
from fastapi.responses import JSONResponse
from ultralytics import YOLO
from PIL import Image

# ===== 모델 로드 =====
MODEL_NAME = os.getenv("YOLO_MODEL", "/app/models/yolo11n.pt")

if not os.path.isfile(MODEL_NAME):
    # 디버그 출력으로 현재 상태 확인
    print(">>> YOLO_MODEL =", MODEL_NAME, file=sys.stderr)
    try:
        print(">>> /app/models =", os.listdir("/app/models"), file=sys.stderr)
    except Exception as e:
        print(">>> cannot list /app/models:", e, file=sys.stderr)
    raise FileNotFoundError(f"weights not found: {MODEL_NAME} (no auto-download)")

model = YOLO(MODEL_NAME)  # 여기까지 오면 파일이 확실히 있음

app = FastAPI(title="YOLO FastAPI", version="1.0.0")


@app.get("/health")
def health():
    return {"ok": True, "model": MODEL_NAME}


@app.post("/predict")
async def predict(
    file: UploadFile = File(...),
    conf: float = Query(0.25, ge=0.0, le=1.0, description="confidence threshold"),
    iou: float = Query(0.7, ge=0.0, le=1.0, description="IoU threshold (NMS)"),
):
    """
    이미지 1장을 받아 YOLO 검출 결과를 JSON으로 반환합니다.
    """
    # 1) 이미지 로드
    image_bytes = await file.read()
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")

    # 2) YOLO 추론
    results = model.predict(
        source=image,
        conf=conf,
        iou=iou,
        device="cpu",
        imgsz=640,
        verbose=False,
    )

    r = results[0]
    names = r.names  # 클래스 id -> 이름 매핑
    out: List[Dict[str, Any]] = []

    if r.boxes is not None and len(r.boxes) > 0:
        # xyxy, conf, cls를 numpy로
        xyxy = r.boxes.xyxy.cpu().numpy()
        confs = r.boxes.conf.cpu().numpy()
        clses = r.boxes.cls.cpu().numpy().astype(int)

        for i in range(len(clses)):
            out.append(
                {
                    "class_id": int(clses[i]),
                    "class_name": names.get(int(clses[i]), str(int(clses[i]))),
                    "confidence": float(round(confs[i], 4)),
                    "box": {
                        "x1": float(xyxy[i][0]),
                        "y1": float(xyxy[i][1]),
                        "x2": float(xyxy[i][2]),
                        "y2": float(xyxy[i][3]),
                    },
                }
            )

    return JSONResponse(
        {
            "model": MODEL_NAME,
            "num_detections": len(out),
            "detections": out,
        }
    )