Update app.py

app.py

@@ -52,9 +52,8 @@ def process_frame(image):
     depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
     depth_map = (depth_map * 255).byte()
     
-    # Reshape depth_map to match color_map dimensions
-    depth_map = depth_map.view(-1)
-    depth_map_colored = color_map[depth_map].view(96, 128, 3)
+    # Apply color map
+    depth_map_colored = color_map[depth_map].squeeze()
     
     return cv2.cvtColor(depth_map_colored.cpu().numpy(), cv2.COLOR_BGR2RGB)