Spaces:
Sleeping
Sleeping
File size: 2,593 Bytes
60b4784 cada0f3 60b4784 cada0f3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | 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,
}
)
|