from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from ultralytics import YOLO
import numpy as np
from PIL import Image
import io
import uvicorn

app = FastAPI()

# CORS 설정
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 모델 로드
print("🔵 Loading local crack model...")
model = YOLO("best.pt")
print("✅ Crack Model Loaded Successfully")

# ===========================================
# 1. (중요) 이 부분이 있어야 "Not Found"가 안 뜹니다!
# ===========================================
@app.get("/")
def read_root():
    return {"message": "ConcreteAI Crack Detection API is running!", "status": "OK"}

# ===========================================
# 2. 예측 API
# ===========================================
@app.post("/predict")
async def predict(img: UploadFile = File(...)):
    try:
        bytes_data = await img.read()
        image = Image.open(io.BytesIO(bytes_data)).convert("RGB")
        np_img = np.array(image)
        results = model(np_img)
        result = results[0]
        
        # crack detection: check boxes
        if result.boxes is None or len(result.boxes) == 0:
            return {
                "data": [
                    {"label": "normal", "confidence": 1.0}
                ]
            }
        
        # There are crack boxes
        conf = float(result.boxes.conf.max().item())
        return {
            "data": [
                {
                    "label": "crack",
                    "confidence": conf
                }
            ]
        }
    except Exception as e:
        print("❌ Prediction error:", e)
        return {
            "data": [{"label": "normal", "confidence": 1.0}],
            "error": str(e)
        }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)