36 lines
988 B
Python
36 lines
988 B
Python
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
|