feat: avoid base64-encoding image data (#76)

* feat: avoid base64-encoding image data

* fix: use correct mimetype

* fix: remove remaining 'b64_encoding_sum' stats

* chore: remove unused imports

* ci: bump node version

* Tiny refactor to put objectURL calls together.

* Can do both

* Turn websocket on by default

* Fix test and add test for websocket use

* flake

---------

Co-authored-by: Almar Klein <almar@almarklein.org>

Author: manzt

js/lib/widget.js

@@ -84,9 +84,13 @@ export class RemoteFrameBufferModel extends DOMWidgetModel {
         this._request_animation_frame();
     }
 
+    /**
+     * @param {Object} msg
+     * @param {DataView[]} buffers
+     */
     on_msg(msg, buffers) {
         if (msg.type === 'framebufferdata') {
-            this.frames.push(msg);
+            this.frames.push({ ...msg, buffers: buffers });
         }
     }
 
@@ -137,9 +141,18 @@ export class RemoteFrameBufferModel extends DOMWidgetModel {
         }
         // Pick the oldest frame from the stack
         let frame = this.frames.shift();
+        let new_src;
+        if (frame.buffers.length > 0) {
+            let blob = new Blob([frame.buffers[0].buffer], { type: frame.mimetype });
+            new_src = URL.createObjectURL(blob);
+        } else {
+            new_src = frame.data_b64;
+        }
+        let old_src = this.img_elements?.[0]?.src;
+        if (old_src.startsWith('blob:')) { URL.revokeObjectURL(old_src); }
         // Update the image sources
         for (let img of this.img_elements) {
-            img.src = frame.src;
+            img.src = new_src;
         }
         // Let the server know we processed the image (even if it's not shown yet)
         this.last_frame = frame;

jupyter_rfb/_utils.py

@@ -17,24 +17,24 @@ def array2compressed(array, quality=90):
     """Convert the given image (a numpy array) as a compressed array.
 
     If the quality is 100, a PNG is returned. Otherwise, JPEG is
-    preferred and PNG is used as a fallback. Returns (preamble, bytes).
+    preferred and PNG is used as a fallback. Returns (mimetype, bytes).
     """
 
     # Drop alpha channel if there is one
     if len(array.shape) == 3 and array.shape[2] == 4:
         array = array[:, :, :3]
 
     if quality >= 100:
-        preamble = "data:image/png;base64,"
+        mimetype = "image/png"
         result = array2png(array)
     else:
-        preamble = "data:image/jpeg;base64,"
+        mimetype = "image/jpeg"
         result = array2jpg(array, quality)
         if result is None:
-            preamble = "data:image/png;base64,"
+            mimetype = "image/png"
             result = array2png(array)
 
-    return preamble, result
+    return mimetype, result
 
 
 class RFBOutputContext(ipywidgets.Output):

jupyter_rfb/widget.py

@@ -91,6 +91,7 @@ def __init__(self, *args, **kwargs):
         self._rfb_last_resize_event = None
         self._rfb_warned_png = False
         self._rfb_lossless_draw_info = None
+        self._use_websocket = True  # Could be a prop, private for now
         # Init stats
         self.reset_stats()
         # Setup events
@@ -285,12 +286,16 @@ def _rfb_send_frame(self, array, is_lossless_redraw=False):
 
         # Turn array into a based64-encoded JPEG or PNG
         t1 = time.perf_counter()
-        preamble, data = array2compressed(array, quality)
+        mimetype, data = array2compressed(array, quality)
+        if self._use_websocket:
+            datas = [data]
+            data_b64 = None
+        else:
+            datas = []
+            data_b64 = f"data:{mimetype};base64," + encodebytes(data).decode()
         t2 = time.perf_counter()
-        src = preamble + encodebytes(data).decode()
-        t3 = time.perf_counter()
 
-        if "jpeg" in preamble:
+        if "jpeg" in mimetype:
             self._rfb_schedule_lossless_draw(array)
         else:
             self._rfb_cancel_lossless_draw()
@@ -308,7 +313,6 @@ def _rfb_send_frame(self, array, is_lossless_redraw=False):
         else:
             # Stats
             self._rfb_stats["img_encoding_sum"] += t2 - t1
-            self._rfb_stats["b64_encoding_sum"] += t3 - t2
             self._rfb_stats["sent_frames"] += 1
             if self._rfb_stats["start_time"] <= 0:  # Start measuring
                 self._rfb_stats["start_time"] = timestamp
@@ -317,11 +321,12 @@ def _rfb_send_frame(self, array, is_lossless_redraw=False):
         # Compose message and send
         msg = dict(
             type="framebufferdata",
-            src=src,
+            mimetype=mimetype,
+            data_b64=data_b64,
             index=self._rfb_frame_index,
             timestamp=timestamp,
         )
-        self.send(msg)
+        self.send(msg, datas)
 
     # ----- related to stats
 
@@ -336,7 +341,6 @@ def reset_stats(self):
             "roundtrip_sum": 0,
             "delivery_sum": 0,
             "img_encoding_sum": 0,
-            "b64_encoding_sum": 0,
         }
 
     def get_stats(self):
@@ -364,7 +368,6 @@ def get_stats(self):
             "roundtrip": d["roundtrip_sum"] / roundtrip_count_div,
             "delivery": d["delivery_sum"] / roundtrip_count_div,
             "img_encoding": d["img_encoding_sum"] / sent_frames_div,
-            "b64_encoding": d["b64_encoding_sum"] / sent_frames_div,
             "fps": d["confirmed_frames"] / fps_div,
         }
 

tests/test_widget.py

@@ -25,8 +25,10 @@ def __init__(self):
         self.has_visible_views = True
         self.msgs = []
 
-    def send(self, msg):
+    def send(self, msg, buffers):
         """Overload the send method so we can check what was sent."""
+        msg = msg.copy()
+        msg["buffers"] = buffers
         self.msgs.append(msg)
 
     def get_frame(self):
@@ -47,6 +49,14 @@ def trigger(self, request):
             self._rfb_draw_requested = True
         self._rfb_maybe_draw()
 
+    def flush(self):
+        """Prentend to flush a frame by setting the widget's frame feedback."""
+        if not len(self.msgs):
+            return
+        self.frame_feedback["index"] = len(self.msgs)
+        self.frame_feedback["timestamp"] = self.msgs[-1]["timestamp"]
+        self.frame_feedback["localtime"] = time.time()
+
 
 def test_widget_frames_and_stats_1():
     """Test sending frames with max 1 in-flight, and how it affects stats."""
@@ -69,10 +79,7 @@ def test_widget_frames_and_stats_1():
     assert fs.get_stats()["sent_frames"] == 1
     assert fs.get_stats()["confirmed_frames"] == 0
 
-    # Flush
-    fs.frame_feedback["index"] = 1
-    fs.frame_feedback["timestamp"] = fs.msgs[-1]["timestamp"]
-    fs.frame_feedback["localtime"] = time.time()
+    fs.flush()
 
     # Trigger, the previous request is still open
     fs.trigger(False)
@@ -81,10 +88,7 @@ def test_widget_frames_and_stats_1():
     assert fs.get_stats()["sent_frames"] == 2
     assert fs.get_stats()["confirmed_frames"] == 1
 
-    # Flush
-    fs.frame_feedback["index"] = 2
-    fs.frame_feedback["timestamp"] = fs.msgs[-1]["timestamp"]
-    fs.frame_feedback["localtime"] = time.time()
+    fs.flush()
 
     fs.trigger(False)
     assert len(fs.msgs) == 2
@@ -135,10 +139,7 @@ def test_widget_frames_and_stats_3():
     assert fs.get_stats()["sent_frames"] == 3
     assert fs.get_stats()["confirmed_frames"] == 0
 
-    # Flush
-    fs.frame_feedback["index"] = 3
-    fs.frame_feedback["timestamp"] = fs.msgs[-1]["timestamp"]
-    fs.frame_feedback["localtime"] = time.time()
+    fs.flush()
 
     # Trigger with True. We request a new frame, but there was a request open
     fs.trigger(True)
@@ -147,10 +148,7 @@ def test_widget_frames_and_stats_3():
     assert fs.get_stats()["sent_frames"] == 4
     assert fs.get_stats()["confirmed_frames"] == 3
 
-    # Flush
-    fs.frame_feedback["index"] = 4
-    fs.frame_feedback["timestamp"] = fs.msgs[-1]["timestamp"]
-    fs.frame_feedback["localtime"] = time.time()
+    fs.flush()
 
     # Trigger, but nothing to send (no frame pending)
     fs.trigger(False)
@@ -173,10 +171,7 @@ def test_widget_frames_and_stats_3():
     fs.trigger(True)
     assert len(fs.msgs) == 7
 
-    # Flush
-    fs.frame_feedback["index"] = 7
-    fs.frame_feedback["timestamp"] = fs.msgs[-1]["timestamp"]
-    fs.frame_feedback["localtime"] = time.time()
+    fs.flush()
 
     # Trigger with False. no new request, but there was a request open
     fs.trigger(False)
@@ -259,10 +254,7 @@ def test_has_visible_views():
     fs.trigger(True)
     assert len(fs.msgs) == 1
 
-    # Flush
-    fs.frame_feedback["index"] = 1
-    fs.frame_feedback["timestamp"] = fs.msgs[-1]["timestamp"]
-    fs.frame_feedback["localtime"] = time.time()
+    fs.flush()
 
     fs.has_visible_views = False
     for _ in range(3):
@@ -302,3 +294,35 @@ def test_snapshot():
     s = w.snapshot()
     assert isinstance(s, Snapshot)
     assert np.all(s.data == w.get_frame())
+
+
+def test_use_websocket():
+    """Test the use of websocket and base64."""
+
+    w = MyRFB()
+
+    # The default uses a websocket
+    w.flush()
+    w.trigger(True)
+    msg = w.msgs[-1]
+    assert len(msg["buffers"]) == 1
+    assert isinstance(msg["buffers"][0], bytes)
+    assert msg["data_b64"] is None
+
+    # Websocket use can be turned off, falling back to base64 encoded images instead
+    w._use_websocket = False
+    w.flush()
+    w.trigger(True)
+    msg = w.msgs[-1]
+    assert len(msg["buffers"]) == 0
+    assert isinstance(msg["data_b64"], str)
+    assert msg["data_b64"].startswith("data:image/jpeg;base64,")
+
+    # Turn it back on
+    w._use_websocket = True
+    w.flush()
+    w.trigger(True)
+    msg = w.msgs[-1]
+    assert len(msg["buffers"]) == 1
+    assert isinstance(msg["buffers"][0], bytes)
+    assert msg["data_b64"] is None