完成训练模块的转移

utils/websocket_server.py

@@ -0,0 +1,47 @@
+from fastapi import WebSocket
+
+
+class SocketManager:
+    def __init__(self):
+        self.rooms = {}
+
+    async def add_to_room(self, room: str, websocket: WebSocket):
+        if room not in self.rooms:
+            self.rooms[room] = []
+        self.rooms[room].append(websocket)
+
+    async def remove_from_room(self, room: str, websocket: WebSocket):
+        if room in self.rooms:
+            self.rooms[room].remove(websocket)
+            if len(self.rooms[room]) == 0:
+                del self.rooms[room]
+                if room.startswith('detect_rtsp_'):
+                    print()
+
+    async def broadcast_to_room(self, room: str, message: str, exclude_websocket: WebSocket = None):
+        if room in self.rooms:
+            for ws in self.rooms[room]:
+                if ws != exclude_websocket:
+                    try:
+                        await ws.send_text(message)
+                    except:
+                        await self.remove_from_room(room, ws)
+
+    async def send_to_room(self, room: str, message: str):
+        if room in self.rooms:
+            for ws in self.rooms[room]:
+                try:
+                    await ws.send_text(message)
+                except Exception as e:
+                    print(e)
+
+    async def send_stream_to_room(self, room: str, message):
+        if room in self.rooms:
+            for ws in self.rooms[room]:
+                try:
+                    await ws.send_bytes(message)
+                except Exception as e:
+                    print(e)
+
+
+room_manager = SocketManager()