Spatial heuristics are broken when when run inside a notebook

[jleibs]

This is impacting at least 0.14.1

To repro, run the following in a notebook:

import rerun as rr
import numpy as np

rr.init("rerun_example_images")

image1 = np.zeros((200, 300, 3), dtype=np.uint8)
image1[:, :, 0] = 255

image2 = np.zeros((200, 300, 3), dtype=np.uint8)
image2[:, :, 1] = 255

image3 = np.zeros((400, 500, 3), dtype=np.uint8)
image3[:, :, 2] = 255

rec = rr.memory_recording()

rr.log("vrs/a", rr.Image(image1))
rr.log("vrs/b", rr.Image(image2))
rr.log("vrs/c", rr.Image(image3))

rec

See that the heuristics incorrectly stack the image:

Whereas the expected behavior is to have them split

Theory

Based on initial investigation my best through is the SpatialTopology store subscriber is getting confused when used with stream_rrd_from_event_listener, which is only used for injecting data into the viewer from notebooks.

I strongly suspect it's missing messages. Possibly related to store-subscriber initialization timing?

[jleibs]

Confirming my suspicions... I monkey-patched the html-renderer to inject a dummy RRD prior to the "real" RRD. This appears to be enough to "prime the pump" so the the StoreSubscriber does the right thing when it is subsequently injected.

dummy = rr.new_recording(application_id="dummy").memory_recording()

def patched_as_html(
    self,
    *,
    width: int = rr.recording.DEFAULT_WIDTH,
    height: int = rr.recording.DEFAULT_HEIGHT,
    app_url: str | None = None,
    timeout_ms: int = rr.recording.DEFAULT_TIMEOUT,
    other: rr.recording.MemoryRecording | None = None,
) -> str:
    from rerun import bindings
    import base64
    import logging
    import random
    import string

    if app_url is None:
        app_url = bindings.get_app_url()

    # Use a random presentation ID to avoid collisions when multiple recordings are shown in the same notebook.
    presentation_id = "".join(random.choice(string.ascii_letters) for i in range(6))

    base64_dummy = base64.b64encode(dummy.storage.concat_as_bytes(None)).decode("utf-8")
    
    if other:
        other = other.storage
    base64_data = base64.b64encode(self.storage.concat_as_bytes(other)).decode("utf-8")

    html_template = f"""
    <div id="{presentation_id}_rrd" style="display: none;" data-rrd="{base64_data}"></div>
    <div id="{presentation_id}_dummy" style="display: none;" data-dummy="{base64_dummy}"></div>
    <div id="{presentation_id}_error" style="display: none;"><p>Timed out waiting for {app_url} to load.</p>
    <p>Consider using <code>rr.start_web_viewer_server()</code></p></div>
    <script>
        {presentation_id}_timeout = setTimeout(() => {{
            document.getElementById("{presentation_id}_error").style.display = 'block';
        }}, {timeout_ms});

        window.addEventListener("message", function(abuffs) {{
            return async function {presentation_id}_onIframeReady(event) {{
                var iframe = document.getElementById("{presentation_id}_iframe");
                if (event.source === iframe.contentWindow) {{
                    clearTimeout({presentation_id}_timeout);
                    document.getElementById("{presentation_id}_error").style.display = 'none';
                    iframe.style.display = 'inline';
                    window.removeEventListener("message", {presentation_id}_onIframeReady);

                    let buffs = await abuffs;
                    iframe.contentWindow.postMessage(buffs[0], "*");
                    await new Promise(r => setTimeout(r, 200));
                    iframe.contentWindow.postMessage(buffs[1], "*");
                }}
            }}
        }}(async function() {{
            await new Promise(r => setTimeout(r, 0));
            
            var dummy_div = document.getElementById("{presentation_id}_dummy");
            var dummy_base64Data = dummy_div.dataset.dummy;
            var dummy_intermediate = atob(dummy_base64Data);
            var dummy_buff = new Uint8Array(dummy_intermediate.length);
            for (var i = 0; i < dummy_intermediate.length; i++) {{
                dummy_buff[i] = dummy_intermediate.charCodeAt(i);
            }}
            
            var div = document.getElementById("{presentation_id}_rrd");
            var base64Data = div.dataset.rrd;
            var intermediate = atob(base64Data);
            var buff = new Uint8Array(intermediate.length);
            for (var i = 0; i < intermediate.length; i++) {{
                buff[i] = intermediate.charCodeAt(i);
            }}
            return [dummy_buff, buff];
        }}()));
    </script>
    <iframe id="{presentation_id}_iframe" width="{width}" height="{height}"
        src="{app_url}?url=web_event://&persist=0&notebook=1"
        frameborder="0" style="display: none;" allowfullscreen=""></iframe>
    """

    return html_template

rr.MemoryRecording.as_html = patched_as_html