fastapi yolo detector

app.py

@@ -0,0 +1,36 @@
+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

requirements.txt

@@ -0,0 +1,6 @@
+fastapi
+uvicorn
+torch
+torchvision
+ultralytics
+opencv-python