from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from PIL import Image
import io
from huggingface_hub import hf_hub_download
from ultralytics import YOLO

app = FastAPI()

# Download YOLOv8 model from Hugging Face
YOLO_MODEL_PATH = hf_hub_download(repo_id="sharktide/RDiCC", filename="runs/detect/train/weights/best.pt")

# Load YOLOv8 model
yolo_model = YOLO(YOLO_MODEL_PATH)

@app.get("/working")
def working():
    return JSONResponse(content={"status": "working"})

@app.post("/detect")
async def detect(file: UploadFile = File(...)):
    # Read and convert image
    image = Image.open(io.BytesIO(await file.read())).convert("RGB")

    # Run YOLO inference
    results = yolo_model(image)

    # Parse results
    detections = []
    for r in results:
        for box in r.boxes:
            cls_id = int(box.cls)
            label = yolo_model.names[cls_id]
            conf = float(box.conf)
            xyxy = box.xyxy[0].tolist()
            detections.append({
                "label": label,
                "confidence": round(conf, 3),
                "box": [round(x, 2) for x in xyxy]
            })

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