main.py

from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from ultralytics import YOLO

app = FastAPI()

origins = [
    "https://yolo.n-e.kr",          # 실제 프론트 도메인
    "https://www.yolo.n-e.kr"       # www 포함한 경우
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,         
    allow_credentials=True,
    allow_methods=["*"],            # 모든 HTTP 메서드 허용
    allow_headers=["*"],            # 모든 헤더 허용
)

# 모델 로딩
model = YOLO("best.pt")  # 사전 학습된 YOLOv8 모델

@app.post("/detect")
async def detect_image(file: UploadFile = File(...)):
    contents = await file.read()

    # 이미지 임시 저장
    with open("temp.jpg", "wb") as f:
        f.write(contents)

    results = model("temp.jpg")
    boxes = results[0].boxes

    labels = []
    for box in boxes:
        conf = float(box.conf[0])
        if conf >= 0.6:
            cls_id = int(box.cls[0])
            label = model.names[cls_id]
            labels.append(label)

    return JSONResponse(content={"ingredients": labels})