main.py

from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from ultralytics import YOLO
import shutil
import uuid
import os

app = FastAPI()

# Load YOLO model
model = YOLO("yolo11s.pt")

# Writable directory on Hugging Face
UPLOAD_DIR = "/tmp/uploads"
os.makedirs(UPLOAD_DIR, exist_ok=True)


@app.get("/")
def home():
    return {"status": "running", "docs": "/docs"}


@app.post("/detect")
async def detect(file: UploadFile = File(...)):
    """
    Run YOLO detection.
    Returns:
    - bounding boxes
    - class labels
    - confidence scores
    """
    # Save uploaded file
    file_id = str(uuid.uuid4())
    file_path = os.path.join(UPLOAD_DIR, f"{file_id}_{file.filename}")
    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    # Run YOLO detection
    results = model(file_path)

    # Extract predictions
    predictions = []
    for box in results[0].boxes:
        predictions.append({
            "class": results[0].names[int(box.cls)],
            "confidence": float(box.conf),
            "bbox": [float(x) for x in box.xyxy[0].tolist()]  # [x1, y1, x2, y2]
        })

    return JSONResponse(content={"predictions": predictions})