fixed websocket image piepline

api/src/server.py

@@ -8,7 +8,7 @@ import numpy as np
 from pydantic import BaseModel
 import torch
 from torchvision import transforms
-from fastapi import Depends, FastAPI, File, HTTPException, Header, UploadFile, WebSocket
+from fastapi import Depends, FastAPI, File, HTTPException, Header, UploadFile, WebSocket, WebSocketDisconnect
 from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
 import io
@@ -101,6 +101,7 @@ async def predict(file: UploadFile = File(...)):
 async def predict_face(websocket: WebSocket):
     await websocket.accept()
 
+    try:
         while True:
             jpg_bytes = await websocket.receive_bytes()
 
@@ -108,7 +109,6 @@ async def predict_face(websocket: WebSocket):
             frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
 
             frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-
             image = Image.fromarray(frame)
 
             H, W, _ = frame.shape
@@ -116,9 +116,14 @@ async def predict_face(websocket: WebSocket):
             scale_h = H / IMAGE_SIZE_YOLO
 
             image = transform_face(image).unsqueeze(0).to(device)
+
             with torch.no_grad():
                 pred = modeL_face(image)
-            bboxes, _, _ = convert_prediction(pred.squeeze(0), image.squeeze(0), threshold=0.9)
+                bboxes, _, _ = convert_prediction(
+                    pred.squeeze(0),
+                    image.squeeze(0),
+                    threshold=0.9
+                )
 
             boxes_to_send = []
             for bbox in bboxes:
@@ -128,9 +133,10 @@ async def predict_face(websocket: WebSocket):
                 ymax = int(bbox[3] * scale_h)
                 boxes_to_send.append([xmin, ymin, xmax, ymax])
 
-        await websocket.send_json({
-            "bboxes": boxes_to_send
-        })
+            await websocket.send_json({"bboxes": boxes_to_send})
+
+    except WebSocketDisconnect:
+        print("Client disconnected")
 
 
 load_dotenv("./database/.env")

website/src/Object_Detection.jsx

@@ -85,6 +85,7 @@ function Object_Detection() {
         };
 
         socket.onmessage = (event) => {
+            console.log("message")
             console.log(event.data);
         };