Added ConnectionManager for better connection handling

app.py

@@ -4,6 +4,10 @@ from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.logger import logger
 
+# Connection Manager
+from connectionManager import ConnectionManager
+
+# Model
 from model import Model
 import base64
 from io import BytesIO
@@ -44,10 +48,12 @@ async def startup_event():
 
     # Initialize the pytorch model
     model = Model()
+    connectionManager = ConnectionManager()
 
     # add model and other preprocess tools too app state
     app.package = {
-        "model": model
+        "model": model,
+        "connectionManager": connectionManager
     }
 
 @app.get("/ping")
@@ -102,28 +108,31 @@ async def do_predict(body: ImageBody):
     img_str = 'data:image/jpeg;base64,' + base64.b64encode(buffered.getvalue()).decode("utf-8")
     
     return {"ok": True, "status": "FINISHED", "result": img_str}
-
+    
 
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
-    await websocket.accept()
+    connectionManager = app.package.get('connectionManager')
+    await connectionManager.connect(websocket)
+    await connectionManager.send_json({"ok": True, "status": "CONNECTED"}, websocket)
     while True:
         try: 
-          data = await websocket.receive_json()
+          data = await connectionManager.receive_json(websocket)
+          if (data is None):
+              # Wait for data
+              if not connectionManager.isConnected(websocket):
+                break
+              if connectionManager.shouldDisconnect(websocket):
+                await websocket.close()
+                connectionManager.disconnect(websocket)
+                break
+              continue
           image: str = data.get('image')
           threshold: float = data.get('threshold') or 0.15
           num_objects: int = data.get('num_objects') or 1
 
-          await websocket.send_json({"ok": True, "status": "STARTED"})
+          await connectionManager.send_json({"ok": True, "status": "STARTED"}, websocket)
 
-          if image == None:
-              await websocket.send_json({
-                  "ok": False,
-                  "status": "ERROR",
-                  "message": "No image provided"
-              })
-              break
-          
           # Run the algorithm
           result = predict(app.package, image, threshold, num_objects)
 
@@ -132,12 +141,15 @@ async def websocket_endpoint(websocket: WebSocket):
           result.save(buffered, format="JPEG")
           img_str = 'data:image/jpeg;base64,' + base64.b64encode(buffered.getvalue()).decode("utf-8")
 
-          await websocket.send_json({"ok": True, "status": "FINISHED", "result": img_str})
+          await connectionManager.send_json({"ok": True, "status": "FINISHED", "result": img_str}, websocket)
 
           await websocket.close()
+          connectionManager.disconnect(websocket)
         except WebSocketDisconnect:
+          print('websocket disconnect')
+          connectionManager.disconnect(websocket)
           break 
 
 if __name__ == '__main__':
     # server api
-    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)
+    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)

connectionManager.py

@@ -0,0 +1,55 @@
+from fastapi import WebSocket
+
+from datetime import datetime
+import json
+from typing import Dict, List
+
+class Connection: 
+    websocket: WebSocket
+    connection_time: datetime
+
+    def __init__(self, websocket: WebSocket, connection_time: datetime):
+        self.websocket = websocket
+        self.connection_time = connection_time
+
+class ConnectionManager:
+    timeout = 60 * 5 # 5 minutes
+
+    def __init__(self):
+        self.active_connections: List[Connection] = []
+
+    async def connect(self, websocket: WebSocket):
+        await websocket.accept()
+        # Add connection time and websocket to active connections
+        self.active_connections.append(Connection(websocket=websocket, connection_time=datetime.now()))
+
+    def isConnected(self, websocket: WebSocket):
+        for connection in self.active_connections:
+            if connection.websocket == websocket:
+                return True
+        return False
+  
+    def shouldDisconnect(self, websocket: WebSocket):
+        for connection in self.active_connections:
+            if connection.websocket == websocket:
+                if (datetime.now() - connection.connection_time).total_seconds() > self.timeout:
+                    return True
+        return False
+
+    async def receive_json(self, websocket: WebSocket):
+        if not self.isConnected(websocket):
+          return None
+        data = await websocket.receive_json()
+        return data
+
+    def disconnect(self, websocket: WebSocket):
+        for connection in self.active_connections:
+            if connection.websocket == websocket:
+                self.active_connections.remove(connection)
+                return True
+        return False
+
+    async def send_json(self, json, websocket: WebSocket):
+        # Only send the message if the connection is still active
+        if self.isConnected(websocket):
+          await websocket.send_json(json)

requirements.txt

@@ -8,7 +8,7 @@ torchvision==0.14.1
 tqdm==4.64.1
 
 pycocotools==2.0.6
-fastapi==0.94.0
+fastapi==0.94.1
 websockets==10.4
 pydantic==1.8.2
 uvicorn[standard]==0.21.0