From c5aa5a267fc77f06629b3c2850059110806a96df Mon Sep 17 00:00:00 2001
From: xuzhen <zhenx@zju.edu.cn>
Date: Sat, 6 Apr 2024 19:20:48 +0800
Subject: [PATCH] evc: improving ws viewer performance

---
 easyvolcap/runners/websocket_server.py |  3 ++-
 easyvolcap/utils/console_utils.py      | 24 ++++++++++++++++++++++--
 scripts/websocket/websocket_client.py  | 15 ++++++++-------
 3 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/easyvolcap/runners/websocket_server.py b/easyvolcap/runners/websocket_server.py
index 0664d22..d051903 100644
--- a/easyvolcap/runners/websocket_server.py
+++ b/easyvolcap/runners/websocket_server.py
@@ -45,7 +45,7 @@ def __init__(self,
                  # Camera related config
                  camera_cfg: dotdict = dotdict(),
                  jpeg_quality: int = 75,
-                 window_size: List[int] = [1080, 1920],
+                 window_size: List[int] = [768, 1366],
 
                  **kwargs,
                  ):
@@ -138,6 +138,7 @@ async def server_loop(self, websocket: websockets.WebSocket, path: str):
             if len(response):
                 camera = Camera()
                 camera.from_string(zlib.decompress(response).decode('ascii'))
+                # camera.from_string(response)
                 self.camera = camera
 
             curr_time = time.perf_counter()
diff --git a/easyvolcap/utils/console_utils.py b/easyvolcap/utils/console_utils.py
index f90a247..9bedb85 100644
--- a/easyvolcap/utils/console_utils.py
+++ b/easyvolcap/utils/console_utils.py
@@ -630,6 +630,10 @@ def __init__(self,
         if self.record_to_file:
             self.timing_record = dotdict()
 
+        self.event_acc = dotdict()
+        self.event_last = dotdict()
+        self.event_denom = dotdict()
+
     def __enter__(self):
         self.start()
 
@@ -656,10 +660,25 @@ def stop(self, print=True, back=2):
         if print: log(f"{(end - start) * 1000:8.3f} ms", self.name, back=back)  # 3 decimals, 3 digits
         return end - start  # return the difference
 
-    def record(self, event: str = ''):
+    def record(self, event: str = '', log_interval: float = -1):
         if self.disabled: return 0
         self.name = event
-        diff = self.stop(print=bool(event), back=3)
+
+        diff = self.stop(print=bool(event) and log_interval <= 0, back=3)
+        curr = time.perf_counter()
+        acc = self.event_acc.get(event, 0)
+        last = self.event_last.get(event, 0)
+        denom = self.event_denom.get(event, 0)
+
+        if (curr - last) > log_interval:  # if this is true, will never have printed
+            log(f"{(acc + diff) / (denom + 1) * 1000:8.3f} ms", event, back=2)
+            self.event_acc[event] = 0
+            self.event_denom[event] = 0
+            self.event_last[event] = curr
+        else:
+            self.event_acc[event] = acc + diff
+            self.event_denom[event] = denom + 1
+
         if self.record_to_file and event:
             if event not in self.timing_record:
                 self.timing_record[event] = []
@@ -667,6 +686,7 @@ def record(self, event: str = ''):
 
             with open(join(self.record_dir, f'{self.exp_name}.json'), 'w') as f:
                 json.dump(self.timing_record, f, indent=4)
+
         self.start()
         return diff
 
diff --git a/scripts/websocket/websocket_client.py b/scripts/websocket/websocket_client.py
index e3c8f16..75680a4 100644
--- a/scripts/websocket/websocket_client.py
+++ b/scripts/websocket/websocket_client.py
@@ -158,17 +158,18 @@ async def websocket_client():
     async with websockets.connect(uri) as websocket:
 
         while True:
-            timer.record('other')
+            timer.record('other', log_interval=2.0)
 
+            # camera_data = viewer.camera.to_string()
             camera_data = zlib.compress(viewer.camera.to_string().encode('ascii'))
-            timer.record('compress')
+            timer.record('compress', log_interval=2.0)
 
             await websocket.send(camera_data)
-            timer.record('send')
-
+            timer.record('send', log_interval=2.0)
 
             buffer = await websocket.recv()
-            timer.record('receive')
+            timer.record('receive', log_interval=2.0)
+
             try:
                 buffer = decode_jpeg(torch.from_numpy(np.frombuffer(buffer, np.uint8)), device='cuda')  # 10ms for 1080p...
             except RuntimeError as e:
@@ -177,10 +178,10 @@ async def websocket_client():
             buffer = torch.cat([buffer, torch.ones_like(buffer[..., :1])], dim=-1)
 
             with lock:
-                image = buffer # might be a cuda tensor or cpu tensor
+                image = buffer  # might be a cuda tensor or cpu tensor
 
             event.set()  # explicit synchronization
-            timer.record('decode')
+            timer.record('decode', log_interval=2.0)
 
 uri = "ws://10.76.5.252:1024"
 image = None