feat: virtual data from precomputed to zarr

python/examples/zarr_gateway.py

@@ -0,0 +1,225 @@
+"""Serves a Neuroglancer image layer as OME-Zarr 0.4, read through the viewer.
+
+Any OME-Zarr reader, such as vizarr, can then open a format that only
+Neuroglancer reads, such as JPEG-encoded precomputed:
+
+    python zarr_gateway.py --port 9000
+
+Open the printed Neuroglancer link and keep its tab open, then open the
+printed vizarr link. Chrome asks vizarr for permission to reach localhost;
+allow it.
+"""
+
+import argparse
+import collections.abc
+import http.server
+import json
+import select
+import socket
+import urllib.parse
+
+import neuroglancer
+import neuroglancer.cli
+import numpy as np
+import tensorstore as ts
+
+ZARR_CHUNK_SHAPE = (1, 256, 256)
+
+
+class ZarrGateway:
+    def __init__(self, viewer: neuroglancer.Viewer, layer: str):
+        info = viewer.volume_info(layer).result()
+        names = list(reversed(info.dimensions.names))
+        if names != ["z", "y", "x"] or set(info.dimensions.units) != {"m"}:
+            raise ValueError(
+                f"Expected spatial dimensions x, y, z in meters, "
+                f"not {info.dimensions.names} in {info.dimensions.units}"
+            )
+        self.layer = layer
+        cache = ts.Context({"cache_pool": {"total_bytes_limit": 2**30}})
+        self.levels: list[ts.TensorStore] = [
+            viewer.volume(
+                layer,
+                dimensions=neuroglancer.CoordinateSpace(
+                    names=info.dimensions.names,
+                    units=info.dimensions.units,
+                    scales=scales,
+                ),
+                context=cache,
+            )
+            .result()
+            .T
+            for scales in info.available_scales
+        ]
+
+    def read(
+        self, path: str, client_disconnected: collections.abc.Callable[[], bool]
+    ) -> bytes | None:
+        parts = path.strip("/").split("/")
+        if parts == [".zgroup"]:
+            return json.dumps({"zarr_format": 2}).encode()
+        if parts == [".zattrs"]:
+            return json.dumps(self._multiscales()).encode()
+        if not parts[0].isdigit() or int(parts[0]) >= len(self.levels):
+            return None
+        level = self.levels[int(parts[0])]
+        if parts[1:] == [".zarray"]:
+            return json.dumps(self._array_metadata(level)).encode()
+        if len(parts) == 4 and all(part.isdigit() for part in parts[1:]):
+            return self._chunk(
+                level, [int(part) for part in parts[1:]], client_disconnected
+            )
+        return None
+
+    def serve(self, port: int) -> None:
+        gateway = self
+
+        class GatewayRequestHandler(http.server.BaseHTTPRequestHandler):
+            def do_GET(self):
+                body = gateway.read(
+                    urllib.parse.urlsplit(self.path).path, self._client_disconnected
+                )
+                if self._client_disconnected():
+                    return
+                self.send_response(404 if body is None else 200)
+                self._send_cors_headers()
+                if body is not None:
+                    self.send_header("Content-Type", "application/octet-stream")
+                    self.send_header("Content-Length", str(len(body)))
+                self.end_headers()
+                if body is not None:
+                    self.wfile.write(body)
+
+            def do_OPTIONS(self):
+                self.send_response(204)
+                self._send_cors_headers()
+                self.end_headers()
+
+            def _client_disconnected(self) -> bool:
+                readable, _, _ = select.select([self.connection], [], [], 0)
+                try:
+                    return bool(readable) and not self.connection.recv(
+                        1, socket.MSG_PEEK
+                    )
+                except ConnectionResetError:
+                    return True
+
+            def _send_cors_headers(self):
+                self.send_header("Access-Control-Allow-Origin", "*")
+                self.send_header("Access-Control-Allow-Headers", "*")
+
+        http.server.ThreadingHTTPServer(
+            ("", port), GatewayRequestHandler
+        ).serve_forever()
+
+    def _multiscales(self):
+        base_scale = _scale_in_nanometers(self.levels[0])
+        datasets = []
+        for level_index, level in enumerate(self.levels):
+            scale = _scale_in_nanometers(level)
+            # OME-Zarr places the voxel centre at its index, not the corner.
+            translation = [
+                start * size + (size - base_size) / 2
+                for start, size, base_size in zip(
+                    level.domain.inclusive_min, scale, base_scale
+                )
+            ]
+            datasets.append(
+                {
+                    "path": str(level_index),
+                    "coordinateTransformations": [
+                        {"type": "scale", "scale": scale},
+                        {"type": "translation", "translation": translation},
+                    ],
+                }
+            )
+        return {
+            "multiscales": [
+                {
+                    "version": "0.4",
+                    "name": self.layer,
+                    "axes": [
+                        {"name": name, "type": "space", "unit": "nanometer"}
+                        for name in ["z", "y", "x"]
+                    ],
+                    "datasets": datasets,
+                }
+            ]
+        }
+
+    def _array_metadata(self, level: ts.TensorStore):
+        return {
+            "zarr_format": 2,
+            "shape": list(level.shape),
+            "chunks": list(ZARR_CHUNK_SHAPE),
+            "dtype": level.dtype.numpy_dtype.str,
+            "compressor": None,
+            "filters": None,
+            "fill_value": 0,
+            "order": "C",
+            "dimension_separator": "/",
+        }
+
+    def _chunk(
+        self,
+        level: ts.TensorStore,
+        chunk_position: list[int],
+        client_disconnected: collections.abc.Callable[[], bool],
+    ) -> bytes | None:
+        level = level.translate_to[0]
+        origin = [
+            position * size for position, size in zip(chunk_position, ZARR_CHUNK_SHAPE)
+        ]
+        if any(start >= extent for start, extent in zip(origin, level.shape)):
+            return None
+        box = tuple(
+            slice(start, min(start + size, extent))
+            for start, size, extent in zip(origin, ZARR_CHUNK_SHAPE, level.shape)
+        )
+        read = level[box].read()
+        while True:
+            try:
+                voxels = read.result(timeout=0.2)
+                break
+            except TimeoutError:
+                if client_disconnected():
+                    read.cancel()
+                    return None
+        chunk = np.zeros(ZARR_CHUNK_SHAPE, dtype=level.dtype.numpy_dtype)
+        chunk[tuple(slice(0, extent) for extent in voxels.shape)] = voxels
+        return chunk.tobytes()
+
+
+def _scale_in_nanometers(level: ts.TensorStore) -> list[float]:
+    return [unit.multiplier * 1e9 for unit in level.dimension_units]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    neuroglancer.cli.add_server_arguments(parser)
+    parser.add_argument(
+        "--source",
+        default="precomputed://gs://neuroglancer-public-data/flyem_fib-25/image",
+    )
+    parser.add_argument("--port", type=int, default=9000)
+    args = parser.parse_args()
+    neuroglancer.cli.handle_server_arguments(args)
+
+    viewer = neuroglancer.Viewer()
+    with viewer.txn() as state:
+        state.layers.append(
+            name="image",
+            layer=neuroglancer.ImageLayer(source=args.source),
+            visible=False,
+        )
+    print(f"Open Neuroglancer and keep its tab open: {viewer}")
+    gateway = ZarrGateway(viewer, "image")
+    print(
+        "Open vizarr: "
+        f"https://hms-dbmi.github.io/vizarr/?source=http://localhost:{args.port}/"
+    )
+    gateway.serve(args.port)
+
+
+if __name__ == "__main__":
+    main()

python/neuroglancer/viewer_base.py

@@ -282,12 +282,20 @@ def volume(
         layer: str,
         *,
         dimensions: coordinate_space.CoordinateSpace | None = None,
+        context: "ts.Context | None" = None,
     ) -> "Future[ts.TensorStore]":
         future: Future[ts.TensorStore] = Future()
 
         def info_done(info_future):
             try:
                 info = info_future.result()
+                chunk_request_info = viewer_config_state.VolumeInfo(
+                    {
+                        key: value
+                        for key, value in info.to_json().items()
+                        if key != "availableScales"
+                    }
+                )
                 dimension_units = [
                     f"{scale} {unit}"
                     for scale, unit in zip(
@@ -317,9 +325,9 @@ def chunk_done(chunk_future):
                         except Exception as e:
                             read_promise.set_exception(e)
 
-                    self._volume_chunk(layer, info, chunk_pos).add_done_callback(
-                        chunk_done
-                    )
+                    self._volume_chunk(
+                        layer, chunk_request_info, chunk_pos
+                    ).add_done_callback(chunk_done)
                     return read_future
 
                 future.set_result(
@@ -334,6 +342,7 @@ def chunk_done(chunk_future):
                         ),
                         dimension_units=dimension_units,
                         chunk_layout=ts.ChunkLayout(read_chunk_shape=info.chunk_shape),
+                        context=context,
                     )
                 )
             except Exception as e:

python/neuroglancer/viewer_config_state.py

@@ -270,6 +270,9 @@ class VolumeInfo(JsonObjectWrapper):
     grid_origin = gridOrigin = wrapped_property("gridOrigin", array_wrapper(np.int64))
     lower_bound = lowerBound = wrapped_property("lowerBound", array_wrapper(np.int64))
     upper_bound = upperBound = wrapped_property("upperBound", array_wrapper(np.int64))
+    available_scales = availableScales = wrapped_property(
+        "availableScales", typed_list(array_wrapper(np.float64))
+    )
 
     @property
     def rank(self):

python/tests/precomputed_test.py

@@ -91,3 +91,96 @@ def test_precomputed(tempdir_server: tuple[pathlib.Path, str], webdriver, spec):
         np.testing.assert_allclose(store.read().result(), b, atol=4, rtol=0)
     else:
         np.testing.assert_equal(store.read().result(), b)
+
+
+def test_reads_every_resolution_that_volume_info_lists(
+    tempdir_server: tuple[pathlib.Path, str], webdriver
+):
+    import tensorstore as ts
+
+    tmp_path, server_url = tempdir_server
+    rng = np.random.default_rng(0)
+    full = rng.integers(0, 256, size=(41, 30, 20, 1), dtype=np.uint8)
+    half = rng.integers(0, 256, size=(20, 15, 10, 1), dtype=np.uint8)
+    for resolution, data in [([1, 1, 1], full), ([2, 2, 2], half)]:
+        store = ts.open(
+            {
+                "driver": "neuroglancer_precomputed",
+                "kvstore": {"driver": "file", "path": str(tmp_path)},
+                "multiscale_metadata": {"data_type": "uint8", "num_channels": 1},
+                "scale_metadata": {
+                    "resolution": resolution,
+                    "encoding": "raw",
+                    "chunk_size": [8, 8, 8],
+                },
+                "schema": {"domain": {"shape": list(data.shape)}},
+            },
+            create=True,
+            open=True,
+        ).result()
+        store[...] = data
+
+    with webdriver.viewer.txn() as s:
+        s.layers.append(
+            name="a",
+            layer=neuroglancer.ImageLayer(source=f"precomputed://{server_url}"),
+        )
+
+    info = webdriver.viewer.volume_info("a").result()
+    assert [list(scales) for scales in info.available_scales] == [
+        [1e-9, 1e-9, 1e-9],
+        [2e-9, 2e-9, 2e-9],
+    ]
+    for scales, data in zip(info.available_scales, [full, half]):
+        dimensions = neuroglancer.CoordinateSpace(
+            names=info.dimensions.names,
+            units=info.dimensions.units,
+            scales=scales,
+        )
+        volume = webdriver.viewer.volume("a", dimensions=dimensions).result()
+        np.testing.assert_equal(volume.read().result(), data[..., 0])
+
+
+def test_a_volume_with_a_cache_reads_each_chunk_from_the_browser_once(
+    tempdir_server: tuple[pathlib.Path, str], webdriver, monkeypatch
+):
+    import tensorstore as ts
+
+    tmp_path, server_url = tempdir_server
+    data = np.random.default_rng(0).integers(0, 256, size=(8, 8, 8, 1), dtype=np.uint8)
+    store = ts.open(
+        {
+            "driver": "neuroglancer_precomputed",
+            "kvstore": {"driver": "file", "path": str(tmp_path)},
+            "multiscale_metadata": {"data_type": "uint8", "num_channels": 1},
+            "scale_metadata": {
+                "resolution": [1, 1, 1],
+                "encoding": "raw",
+                "chunk_size": [8, 8, 8],
+            },
+            "schema": {"domain": {"shape": [8, 8, 8, 1]}},
+        },
+        create=True,
+    ).result()
+    store[...] = data
+
+    with webdriver.viewer.txn() as s:
+        s.layers.append(
+            name="a",
+            layer=neuroglancer.ImageLayer(source=f"precomputed://{server_url}"),
+        )
+
+    chunk_requests = []
+    request_chunk = webdriver.viewer._volume_chunk
+
+    def count_chunk_request(layer, info, chunk_grid_position):
+        chunk_requests.append(list(chunk_grid_position))
+        return request_chunk(layer, info, chunk_grid_position)
+
+    monkeypatch.setattr(webdriver.viewer, "_volume_chunk", count_chunk_request)
+    volume = webdriver.viewer.volume(
+        "a", context=ts.Context({"cache_pool": {"total_bytes_limit": 10_000_000}})
+    ).result()
+    for z in range(8):
+        np.testing.assert_equal(volume[:, :, z].read().result(), data[:, :, z, 0])
+    assert chunk_requests == [[0, 0, 0]]