Trial for sockets

app.py

@@ -18,6 +18,10 @@ from ECG import classify_new_ecg
 from ultralytics import YOLO
 import tensorflow as tf
 from fastapi import HTTPException
+from fastapi import WebSocket, WebSocketDisconnect
+import base64
+import cv2
+
 
 
 app = FastAPI()
@@ -245,4 +249,46 @@ async def process_image(file: UploadFile = File(...)):
         "message": "Image processed successfully",
         "KeypointsXY": keypoints.tolist(),
         "confidences": confidences.tolist()
-    })
+    })
+
+
+# WebSocket endpoint to handle image processing
+@app.websocket("/ws/image")
+async def websocket_endpoint(websocket: WebSocket):
+    model = YOLO("yolo11n-pose_float16.tflite")
+    print("Model loaded successfully")
+    await websocket.accept()
+    try:
+        while True:
+            data = await websocket.receive_text()  # Receive base64-encoded image data
+            image_data = base64.b64decode(data)  # Decode the image data
+
+            # Convert image bytes to numpy array and decode with OpenCV
+            np_arr = np.frombuffer(image_data, np.uint8)
+            frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
+            
+            # Run YOLO pose estimation
+            if model is not None:
+                try:
+                    results = model.predict(frame, save=False, conf=0.3)
+                    
+                    if results and results[0].keypoints:
+                        keypoints = results[0].keypoints.xy.cpu().numpy().tolist()  # Extract keypoints
+                        confidences = results[0].boxes.conf.cpu().numpy().tolist() if results[0].boxes else []
+
+                        # Send the results back to the client
+                        response = {
+                            "message": "Pose detected",
+                            "KeypointsXY": keypoints[:5],  # Limit to first 5 keypoints for brevity
+                            "Confidences": confidences[:5],  # Limit to first 5 confidences
+                        }
+                        await websocket.send_text(str(response))
+                    else:
+                        await websocket.send_text("❌ No keypoints detected.")
+                except Exception as e:
+                    await websocket.send_text(f"⚠️ Error processing image: {str(e)}")
+            else:
+                await websocket.send_text("⚠️ Model not loaded.")
+                
+    except WebSocketDisconnect:
+        print("🔌 Client disconnected")

requirements.txt

@@ -100,4 +100,5 @@ uvicorn==0.34.2
 Werkzeug==3.1.3
 wfdb==4.3.0
 wrapt==1.17.2
-yarl==1.20.0
+yarl==1.20.0
+websockets