From de092054904655224dc8bc760c772da8ce030853 Mon Sep 17 00:00:00 2001
From: "pengsida.psd" <pengsida.psd@antgroup.com>
Date: Sat, 6 Apr 2024 19:01:56 +0800
Subject: [PATCH] evc: run websocket client on osx

---
 easyvolcap/runners/volumetric_video_viewer.py |  1 -
 scripts/websocket/websocket_client.py         | 26 ++++++++++++-------
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/easyvolcap/runners/volumetric_video_viewer.py b/easyvolcap/runners/volumetric_video_viewer.py
index f690a39..fc61e84 100644
--- a/easyvolcap/runners/volumetric_video_viewer.py
+++ b/easyvolcap/runners/volumetric_video_viewer.py
@@ -1322,7 +1322,6 @@ def glfw_framebuffer_size_callback(self, window, width, height):
         ori_W = self.W
         self.H = height
         self.W = width
-
         if ori_W > 0 and ori_H > 0:
             ratio = min(self.W / ori_W, self.H / ori_H)
             self.camera.K[0, 0] *= ratio   # only update one of them
diff --git a/scripts/websocket/websocket_client.py b/scripts/websocket/websocket_client.py
index 00912f1..e3c8f16 100644
--- a/scripts/websocket/websocket_client.py
+++ b/scripts/websocket/websocket_client.py
@@ -27,7 +27,7 @@
 
 class Viewer(VolumetricVideoViewer):
     def __init__(self,
-                 window_size: List[int] = [1080, 1920],  # height, width
+                 window_size: List[int] = [768, 1366],  # height, width
                  window_title: str = f'EasyVolcap WebSocket Client',  # MARK: global config
 
                  font_size: int = 18,
@@ -122,8 +122,12 @@ def render(self):
         global image
         event.wait()
         event.clear()
-        if image.shape[0] == self.H and image.shape[1] == self.W:
-            self.quad.copy_to_texture(image)
+
+        with lock:
+            buffer = image
+
+        if buffer.shape[0] == self.H and buffer.shape[1] == self.W:
+            self.quad.copy_to_texture(buffer)
             self.quad.draw()
         return None, None
 
@@ -162,22 +166,26 @@ async def websocket_client():
             await websocket.send(camera_data)
             timer.record('send')
 
-            # await websocket.send('')
-            # timer.record('send')
 
             buffer = await websocket.recv()
             timer.record('receive')
+            try:
+                buffer = decode_jpeg(torch.from_numpy(np.frombuffer(buffer, np.uint8)), device='cuda')  # 10ms for 1080p...
+            except RuntimeError as e:
+                buffer = decode_jpeg(torch.from_numpy(np.frombuffer(buffer, np.uint8)), device='cpu')  # 10ms for 1080p...
+            buffer = buffer.permute(1, 2, 0)
+            buffer = torch.cat([buffer, torch.ones_like(buffer[..., :1])], dim=-1)
+
+            with lock:
+                image = buffer # might be a cuda tensor or cpu tensor
 
-            image = decode_jpeg(torch.from_numpy(np.frombuffer(buffer, np.uint8)), device='cuda')  # 10ms for 1080p...
-            image = image.permute(1, 2, 0)
-            image = torch.cat([image, torch.ones_like(image[..., :1])], dim=-1)
-            # image = torch.from_numpy(np.frombuffer(buffer, np.uint8)).to('cuda', non_blocking=True)
             event.set()  # explicit synchronization
             timer.record('decode')
 
 uri = "ws://10.76.5.252:1024"
 image = None
 viewer = Viewer()
+lock = threading.Lock()
 event = threading.Event()
 timer.disabled = False