from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from ultralytics import YOLO
from huggingface_hub import hf_hub_download

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=["*"],
)

# ===========================================
#       DOWNLOAD CRACK MODEL FROM HF
# ===========================================

print("🔵 Loading local crack model...")
model = YOLO("best.pt")
print("✅ Crack Model Loaded Successfully")

# ===========================================
#               PREDICTION 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, conf=0.40)
        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)