import os
import gdown
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from ultralytics import YOLO
import cv2
import numpy as np
import uvicorn

app = FastAPI()

# الموديل بتاعك
FILE_ID = '1a5-mI8D6oT2K2r0Y-j0U-pE4k3t2r9mC' 
drive_url = f'https://drive.google.com/uc?id={FILE_ID}'
model_path = 'best.pt'

if not os.path.exists(model_path):
    gdown.download(drive_url, model_path, quiet=False)

model = YOLO(model_path)

@app.get("/")
def home():
    return {"status": "Online", "model": "YOLOv8 Loaded"}

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        nparr = np.frombuffer(contents, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        results = model(img)
        predictions = []
        for result in results:
            for box in result.boxes:
                predictions.append({
                    "class": model.names[int(box.cls)],
                    "confidence": float(box.conf),
                    "bbox": [float(x) for x in box.xyxy[0]]
                })
        return JSONResponse(content={"predictions": predictions})
    except Exception as e:
        return JSONResponse(content={"error": str(e)}, status_code=500)

if __name__ == "__main__":
    # لازم بورت 7860 عشان Hugging Face
    uvicorn.run(app, host="0.0.0.0", port=7860)