from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
import cv2
import numpy as np
import io
from PIL import Image
from ultralytics import YOLO

app = FastAPI(title="YOLOv11 Detection API")

# Load the YOLO model
try:
    model = YOLO('yolo11n.pt')
except Exception as e:
    print(f"Error loading model: {e}")
    model = None

@app.get("/")
async def root():
    return {"message": "YOLOv11 Detection API is running. Go to /docs for API documentation."}

@app.get("/health")
async def health():
    if model is not None:
        return {"status": "healthy", "model": "yolo11n.pt"}
    else:
        raise HTTPException(status_code=503, detail="Model not loaded")

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    # Read the uploaded image
    try:
        contents = await file.read()
        image = Image.open(io.BytesIO(contents)).convert("RGB")
        img_array = np.array(image)
        # Convert RGB to BGR for OpenCV/YOLO if needed
        img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid image: {e}")

    # Run inference
    results = model(img_bgr, verbose=False)
    
    detections = []
    for box in results[0].boxes:
        class_id = int(box.cls[0])
        class_name = model.names[class_id]
        confidence = float(box.conf[0])
        x1, y1, x2, y2 = box.xyxy[0].tolist()
        
        detections.append({
            "class": class_name,
            "confidence": confidence,
            "bbox": [x1, y1, x2, y2]
        })

    return JSONResponse(content={
        "filename": file.filename,
        "detections": detections,
        "count": len(detections)
    })

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