from fastapi import FastAPI, WebSocket
from ultralytics import YOLO
import torch
import cv2
import base64
import numpy as np

app = FastAPI()

model = YOLO("yolov8n.pt")
if torch.cuda.is_available():
    model.to('cuda')

@app.websocket("/detect")
async def detect(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            # Decode base64 image bytes
            img_data = base64.b64decode(data)
            np_arr = np.frombuffer(img_data, np.uint8)
            frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

            results = model(frame, stream=False)
            boxes = []
            for r in results:
                for box in r.boxes:
                    boxes.append({
                        "class": int(box.cls[0]),
                        "confidence": float(box.conf[0]),
                    })

            await websocket.send_json({"detections": boxes})
    except:
        pass