update http_tensorstore

cellmap_flow/cli/cli.py

@@ -2,7 +2,9 @@
 import logging
 import click
 
+from cellmap_flow.server import CellMapFlowServer
 from cellmap_flow.utils.bsub_utils import start_hosts
+from cellmap_flow.utils.data import ScriptModelConfig
 from cellmap_flow.utils.neuroglancer_utils import generate_neuroglancer_link
 
 
@@ -152,6 +154,26 @@ def bioimage(model_path, data_path, queue, charge_group):
     run(command, data_path, queue, charge_group)
 
 
+@cli.command()
+@click.option(
+    "--script_path",
+    "-s",
+    type=str,
+    help="Path to the Python script containing model specification",
+)
+@click.option("--dataset", "-d", type=str, help="Path to the dataset")
+def script_server_check(script_path, dataset):
+    model_config = ScriptModelConfig(script_path=script_path)
+    server = CellMapFlowServer(dataset, model_config)
+    chunk_x = 2
+    chunk_y = 2
+    chunk_z = 2
+
+    server._chunk_impl(None, None, chunk_x, chunk_y, chunk_z, None)
+
+    print("Server check passed")
+
+
 def run(
     command,
     dataset_path,

cellmap_flow/image_data_interface.py

@@ -31,7 +31,6 @@ def __init__(
             self.output_voxel_size = output_voxel_size
         else:
             self.output_voxel_size = self.voxel_size
-
 
     @property
     def ts(self):
@@ -51,5 +50,4 @@ def to_ndarray_ts(self, roi=None):
             self.swap_axes,
             self.custom_fill_value,
         )
-        self.ts = None
         return res

cellmap_flow/inferencer.py

@@ -28,7 +28,7 @@ def predict(read_roi, write_roi, config, **kwargs):
     device = kwargs.get("device")
     if device is None:
         raise ValueError("device must be provided in kwargs")
-    
+
     use_half_prediction = kwargs.get("use_half_prediction", False)
 
     raw_input = idi.to_ndarray_ts(read_roi)
@@ -54,7 +54,7 @@ def __init__(self, model_config: ModelConfig, use_half_prediction=False):
             self.model_config.config, "write_shape"
         ):
             self.context = (
-               Coordinate(self.model_config.config.read_shape)
+                Coordinate(self.model_config.config.read_shape)
                 - Coordinate(self.model_config.config.write_shape)
             ) / 2
 
@@ -115,7 +115,12 @@ def process_chunk_basic(self, idi, roi):
 
         input_roi = output_roi.grow(self.context, self.context)
         result = self.model_config.config.predict(
-            input_roi, output_roi, self.model_config.config, idi=idi, device=self.device, use_half_prediction=self.use_half_prediction
+            input_roi,
+            output_roi,
+            self.model_config.config,
+            idi=idi,
+            device=self.device,
+            use_half_prediction=self.use_half_prediction,
         )
         write_data = self.model_config.config.normalize_output(result)
 

cellmap_flow/norm/input_normalize.py

@@ -33,13 +33,13 @@ def __init__(self, min_value=0.0, max_value=255.0):
 
     def normalize(self, data: np.ndarray) -> np.ndarray:
         data = data.astype(np.float32)
-        data.clip(self.min_value, self.max_value)
+        data = data.clip(self.min_value, self.max_value)
         return ((data - self.min_value) / (self.max_value - self.min_value)).astype(
             np.float32
         )
 
 
-NormalizationMethods = [f.name() for f in InputNormalizer.__subclasses__()]
+NormalizationMethods = [f for f in InputNormalizer.__subclasses__()]
 
 
 def get_normalization(elms: dict) -> InputNormalizer:

cellmap_flow/server.py

@@ -30,7 +30,17 @@ def __init__(self, dataset_name: str, model_config: ModelConfig):
         """
         Initialize the server and set up routes via decorators.
         """
-        self.block_shape = [int(x) for x in model_config.config.block_shape]
+
+        # this is zyx
+        self.read_block_shape = [int(x) for x in model_config.config.block_shape]
+
+        # this needs to have z and x swapped
+        self.n5_block_shape = self.read_block_shape.copy()
+        self.n5_block_shape[0], self.n5_block_shape[2] = (
+            self.n5_block_shape[2],
+            self.n5_block_shape[0],
+        )
+
         self.input_voxel_size = Coordinate(model_config.config.input_voxel_size)
         self.output_voxel_size = Coordinate(model_config.config.output_voxel_size)
         self.output_channels = model_config.config.output_channels
@@ -41,22 +51,29 @@ def __init__(self, dataset_name: str, model_config: ModelConfig):
         self.idi_raw = ImageDataInterface(
             dataset_name, target_resolution=self.input_voxel_size
         )
+        output_shape = (
+            np.array(self.idi_raw.shape)
+            * np.array(self.input_voxel_size)
+            / np.array(self.output_voxel_size)
+        )
+
         if ".zarr" in dataset_name:
             # Convert from (z, y, x) -> (x, y, z) plus channels
             self.vol_shape = np.array(
-                [*np.array(self.idi_raw.shape)[::-1], self.output_channels]
+                [
+                    *output_shape[::-1],
+                    self.output_channels,
+                ]
             )
             self.axis = ["x", "y", "z", "c^"]
         else:
             # For non-Zarr data
-            self.vol_shape = np.array(
-                [*np.array(self.idi_raw.shape), self.output_channels]
-            )
+            self.vol_shape = np.array([*output_shape, self.output_channels])
             self.axis = ["z", "y", "x", "c^"]
 
         # Chunk encoding for N5
         self.chunk_encoder = N5ChunkWrapper(
-            np.uint8, self.block_shape, compressor=numcodecs.GZip()
+            np.uint8, self.n5_block_shape, compressor=numcodecs.GZip()
         )
 
         # Create and configure Flask
@@ -269,7 +286,7 @@ def _attributes_impl(self, dataset, scale):
                 "translate": [0.0, 0.0, 0.0, 0.0],
             },
             "compression": {"type": "gzip", "useZlib": False, "level": -1},
-            "blockSize": list(self.block_shape),
+            "blockSize": list(self.n5_block_shape),
             "dataType": "uint8",
             "dimensions": self.vol_shape.tolist(),
         }
@@ -292,8 +309,8 @@ def _input_normalize_impl(self, norm_type, min_value, max_value):
             )
 
     def _chunk_impl(self, dataset, scale, chunk_x, chunk_y, chunk_z, chunk_c):
-        corner = self.block_shape[:3] * np.array([chunk_z, chunk_y, chunk_x])
-        box = np.array([corner, self.block_shape[:3]]) * self.output_voxel_size
+        corner = self.read_block_shape[:3] * np.array([chunk_z, chunk_y, chunk_x])
+        box = np.array([corner, self.read_block_shape[:3]]) * self.output_voxel_size
         roi = Roi(box[0], box[1])
         chunk_data = self.inferencer.process_chunk(self.idi_raw, roi)
         return (

cellmap_flow/utils/neuroglancer_utils.py

@@ -1,28 +1,55 @@
 import neuroglancer
 import itertools
 import logging
+import os
 
 neuroglancer.set_server_bind_address("0.0.0.0")
 
 logger = logging.getLogger(__name__)
 
 from cellmap_flow.image_data_interface import ImageDataInterface
+from cellmap_flow.utils.scale_pyramid import ScalePyramid
+
+
 # TODO support multiresolution datasets
-def get_raw_layer(dataset_path, filetype):
-    if filetype == "zarr":
+def get_raw_layer(dataset_path, filetype, is_multiscale=False):
+    if filetype == "n5":
         axis = ["x", "y", "z"]
     else:
         axis = ["z", "y", "x"]
-    image = ImageDataInterface(dataset_path)
-    return neuroglancer.ImageLayer(
-        source=neuroglancer.LocalVolume(
-            data=image.ts,
-            dimensions=neuroglancer.CoordinateSpace(
-                names=axis,
-                units="nm",
-                scales=image.voxel_size,
-            ),
-            voxel_offset=image.offset,
+
+    layers = []
+
+    if is_multiscale:
+        scales = [
+            f for f in os.listdir(dataset_path) if f[0] == "s" and f[1:].isdigit()
+        ]
+        scales.sort(key=lambda x: int(x[1:]))
+        for scale in scales:
+            image = ImageDataInterface(f"{os.path.join(dataset_path, scale)}")
+            layers.append(
+                neuroglancer.LocalVolume(
+                    data=image.ts,
+                    dimensions=neuroglancer.CoordinateSpace(
+                        names=axis,
+                        units="nm",
+                        scales=image.voxel_size,
+                    ),
+                    voxel_offset=image.offset,
+                )
+            )
+        return ScalePyramid(layers)
+    else:
+        image = ImageDataInterface(dataset_path)
+        return neuroglancer.ImageLayer(
+            source=neuroglancer.LocalVolume(
+                data=image.ts,
+                dimensions=neuroglancer.CoordinateSpace(
+                    names=axis,
+                    units="nm",
+                    scales=image.voxel_size,
+                ),
+                voxel_offset=image.offset,
             )
         )
 
@@ -33,20 +60,24 @@ def generate_neuroglancer_link(dataset_path, inference_dict):
 
     # Add a layer to the viewer
     with viewer.txn() as s:
+        is_multi_scale = False
         # if multiscale dataset
-        # if (
-        #     dataset_path.split("/")[-1].startswith("s")
-        #     and dataset_path.split("/")[-1][1:].isdigit()
-        # ):
-        #     dataset_path = dataset_path.rsplit("/", 1)[0]
+        if (
+            dataset_path.split("/")[-1].startswith("s")
+            and dataset_path.split("/")[-1][1:].isdigit()
+        ):
+            dataset_path = dataset_path.rsplit("/", 1)[0]
+            is_multi_scale = True
+
         if ".zarr" in dataset_path:
             filetype = "zarr"
         elif ".n5" in dataset_path:
             filetype = "n5"
         else:
             filetype = "precomputed"
         if dataset_path.startswith("/"):
-            s.layers["raw"] = get_raw_layer(dataset_path, filetype)
+            layer = get_raw_layer(dataset_path, filetype, is_multi_scale)
+            s.layers.append("raw", layer)
             # if "nrs/cellmap" in dataset_path:
             #     security = "https"
             #     dataset_path = dataset_path.replace("/nrs/cellmap/", "nrs/")

cellmap_flow/utils/scale_pyramid.py

@@ -0,0 +1,113 @@
+# copied from https://github.com/funkelab/funlib.show.neuroglancer/blob/master/funlib/show/neuroglancer/scale_pyramid.py
+
+import neuroglancer
+import operator
+import logging
+
+import numpy as np
+
+
+logger = logging.getLogger(__name__)
+
+
+class ScalePyramid(neuroglancer.LocalVolume):
+    """A neuroglancer layer that provides volume data on different scales.
+    Mimics a LocalVolume.
+
+    Args:
+
+            volume_layers (``list`` of ``LocalVolume``):
+
+                One ``LocalVolume`` per provided resolution.
+    """
+
+    def __init__(self, volume_layers):
+        volume_layers = volume_layers
+
+        super(neuroglancer.LocalVolume, self).__init__()
+
+        logger.info("Creating scale pyramid...")
+
+        self.min_voxel_size = min(
+            [tuple(layer.dimensions.scales) for layer in volume_layers]
+        )
+        self.max_voxel_size = max(
+            [tuple(layer.dimensions.scales) for layer in volume_layers]
+        )
+
+        self.dims = len(volume_layers[0].dimensions.scales)
+        self.volume_layers = {
+            tuple(
+                int(x)
+                for x in map(
+                    operator.truediv, layer.dimensions.scales, self.min_voxel_size
+                )
+            ): layer
+            for layer in volume_layers
+        }
+
+        logger.info("min_voxel_size: %s", self.min_voxel_size)
+        logger.info("scale keys: %s", self.volume_layers.keys())
+        logger.info(self.info())
+
+    @property
+    def volume_type(self):
+        return self.volume_layers[(1,) * self.dims].volume_type
+
+    @property
+    def token(self):
+        return self.volume_layers[(1,) * self.dims].token
+
+    def info(self):
+        reference_layer = self.volume_layers[(1,) * self.dims]
+        # return reference_layer.info()
+
+        reference_info = reference_layer.info()
+
+        info = {
+            "dataType": reference_info["dataType"],
+            "encoding": reference_info["encoding"],
+            "generation": reference_info["generation"],
+            "coordinateSpace": reference_info["coordinateSpace"],
+            "shape": reference_info["shape"],
+            "volumeType": reference_info["volumeType"],
+            "voxelOffset": reference_info["voxelOffset"],
+            "chunkLayout": reference_info["chunkLayout"],
+            "downsamplingLayout": reference_info["downsamplingLayout"],
+            "maxDownsampling": int(
+                np.prod(np.array(self.max_voxel_size) // np.array(self.min_voxel_size))
+            ),
+            "maxDownsampledSize": reference_info["maxDownsampledSize"],
+            "maxDownsamplingScales": reference_info["maxDownsamplingScales"],
+        }
+
+        return info
+
+    def get_encoded_subvolume(self, data_format, start, end, scale_key=None):
+        if scale_key is None:
+            scale_key = ",".join(("1",) * self.dims)
+
+        scale = tuple(int(s) for s in scale_key.split(","))
+        closest_scale = None
+        min_diff = np.inf
+        for volume_scales in self.volume_layers.keys():
+            scale_diff = np.array(scale) // np.array(volume_scales)
+            if any(scale_diff < 1):
+                continue
+            scale_diff = scale_diff.max()
+            if scale_diff < min_diff:
+                min_diff = scale_diff
+                closest_scale = volume_scales
+
+        assert closest_scale is not None
+        relative_scale = np.array(scale) // np.array(closest_scale)
+
+        return self.volume_layers[closest_scale].get_encoded_subvolume(
+            data_format, start, end, scale_key=",".join(map(str, relative_scale))
+        )
+
+    def get_object_mesh(self, object_id):
+        return self.volume_layers[(1,) * self.dims].get_object_mesh(object_id)
+
+    def invalidate(self):
+        return self.volume_layers[(1,) * self.dims].invalidate()