Function process_yolov5_ws

ml-yolo/backend/main.py:83–111

(websocket: WebSocket, client_id: int)

@app.websocket("/yolo_ws/{client_id}")
async def process_yolov5_ws(websocket: WebSocket, client_id: int):
    await conn_mgr.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()

            # Convert to PIL image
            image = data[data.find(",") + 1:]
            dec = base64.b64decode(image + "===")
            image = Image.open(BytesIO(dec)).convert("RGB")

            # Process the image
            name = f"/data/{str(uuid.uuid4())}.png"
            image.filename = name
            classes, converted_img = yolov5(image)

            result = {
                "prediction": json.dumps(classes),
                "output": base64_encode_img(converted_img),
            }
            # logging.info("-----", json.dumps(result))

            # Send back the result
            await conn_mgr.send_message(json.dumps(result), websocket)

            # await conn_mgr.broadcast(f"Client #{client_id} says: {data}")
    except WebSocketDisconnect:
        conn_mgr.disconnect(websocket)
        await conn_mgr.broadcast(f"Client #{client_id} left the chat")


ws_client_html = """