Update app.py

app.py

@@ -1,9 +1,13 @@
 import cv2
 import torch
-from transformers import DPTForDepthEstimation, DPTImageProcessor
 import numpy as np
+from transformers import DPTForDepthEstimation, DPTImageProcessor
 import time
 import warnings
+import asyncio
+import json
+import websockets
+
 warnings.filterwarnings("ignore", message="It looks like you are trying to rescale already rescaled images.")
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -25,72 +29,83 @@ def manual_normalize(depth_map):
         return np.zeros_like(depth_map, dtype=np.uint8)
 
 frame_skip = 4
-frame_count = 0
 color_map = cv2.applyColorMap(np.arange(256, dtype=np.uint8), cv2.COLORMAP_INFERNO)
 
-prev_frame_time = 0
+connected = set()
 
-while True:
-    ret, frame = cap.read()
-    if not ret:
-        break
+async def broadcast(message):
+    for websocket in connected:
+        try:
+            await websocket.send(message)
+        except websockets.exceptions.ConnectionClosed:
+            connected.remove(websocket)
 
-    frame_count += 1
-    if frame_count % frame_skip != 0:
-        continue
+async def handler(websocket, path):
+    connected.add(websocket)
+    try:
+        await websocket.wait_closed()
+    finally:
+        connected.remove(websocket)
 
-    rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-    resized_frame = resize_image(rgb_frame)
+async def process_frames():
+    frame_count = 0
+    prev_frame_time = 0
 
-    inputs = processor(images=resized_frame, return_tensors="pt").to(device)
-    inputs = {k: v.to(torch.float16) for k, v in inputs.items()}
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
 
-    with torch.no_grad():
-        outputs = model(**inputs)
-        predicted_depth = outputs.predicted_depth
+        frame_count += 1
+        if frame_count % frame_skip != 0:
+            continue
 
-    depth_map = predicted_depth.squeeze().cpu().numpy()
+        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        resized_frame = resize_image(rgb_frame)
 
-    # Check Input Data
-    print(f"depth_map shape: {depth_map.shape}")
-    print(f"depth_map min: {np.min(depth_map)}, max: {np.max(depth_map)}")
-    print(f"depth_map dtype: {depth_map.dtype}")
+        inputs = processor(images=resized_frame, return_tensors="pt").to(device)
+        inputs = {k: v.to(torch.float16) for k, v in inputs.items()}
 
-    # Handle invalid values
-    depth_map = np.nan_to_num(depth_map, nan=0.0, posinf=0.0, neginf=0.0)
+        with torch.no_grad():
+            outputs = model(**inputs)
+            predicted_depth = outputs.predicted_depth
 
-    # Ensure depth_map is in float32 format
-    depth_map = depth_map.astype(np.float32)
+        depth_map = predicted_depth.squeeze().cpu().numpy()
 
-    # Check for zero-sized arrays
-    if depth_map.size == 0:
-        print("Error: depth_map is empty")
-        depth_map = np.zeros((256, 256), dtype=np.uint8)
-    else:
-        # Handle empty or constant arrays
-        if np.any(depth_map) and np.min(depth_map) != np.max(depth_map):
-            depth_map = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
+        depth_map = np.nan_to_num(depth_map, nan=0.0, posinf=0.0, neginf=0.0)
+        depth_map = depth_map.astype(np.float32)
+
+        if depth_map.size == 0:
+            depth_map = np.zeros((256, 256), dtype=np.uint8)
         else:
-            depth_map = np.zeros_like(depth_map, dtype=np.uint8)
+            if np.any(depth_map) and np.min(depth_map) != np.max(depth_map):
+                depth_map = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
+            else:
+                depth_map = np.zeros_like(depth_map, dtype=np.uint8)
 
-    # Use manual normalization as a fallback
-    if np.all(depth_map == 0):
-        depth_map = manual_normalize(depth_map)
+        if np.all(depth_map == 0):
+            depth_map = manual_normalize(depth_map)
 
-    depth_map_colored = cv2.applyColorMap(depth_map, color_map)
-    depth_map_colored = cv2.resize(depth_map_colored, (frame.shape[1], frame.shape[0]))
+        data = {
+            'depthMap': depth_map.tolist(),
+            'rgbFrame': rgb_frame.tolist()
+        }
+        
+        await broadcast(json.dumps(data))
 
-    combined = np.hstack((frame, depth_map_colored))
+        new_frame_time = time.time()
+        fps = 1 / (new_frame_time - prev_frame_time)
+        prev_frame_time = new_frame_time
 
-    new_frame_time = time.time()
-    fps = 1 / (new_frame_time - prev_frame_time)
-    prev_frame_time = new_frame_time
-    cv2.putText(combined, f"FPS: {int(fps)}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+        if cv2.waitKey(1) & 0xFF == ord('q'):
+            break
 
-    cv2.imshow('Webcam and Depth Map', combined)
+    cap.release()
+    cv2.destroyAllWindows()
 
-    if cv2.waitKey(1) & 0xFF == ord('q'):
-        break
+async def main():
+    server = await websockets.serve(handler, "localhost", 8765)
+    await asyncio.gather(server.wait_closed(), process_frames())
 
-cap.release()
-cv2.destroyAllWindows()
+if __name__ == "__main__":
+    asyncio.run(main())