diff --git a/gunpowder/__init__.py b/gunpowder/__init__.py index d3242f24..f51f269d 100644 --- a/gunpowder/__init__.py +++ b/gunpowder/__init__.py @@ -6,7 +6,7 @@ from .array_spec import ArraySpec from .batch import Batch from .batch_request import BatchRequest -from .build import build +from .build import build, build_neuroglancer from .coordinate import Coordinate from .graph import Graph, Node, Edge, GraphKey, GraphKeys from .graph_spec import GraphSpec diff --git a/gunpowder/build.py b/gunpowder/build.py index 1b82466c..6dbb528c 100644 --- a/gunpowder/build.py +++ b/gunpowder/build.py @@ -23,3 +23,50 @@ def __exit__(self, type, value, traceback): logger.debug("leaving context, tearing down pipeline") self.pipeline.internal_teardown() logger.debug("tear down completed") + + +import neuroglancer +from .neuroglancer.event import step_next + + +class build_neuroglancer(object): + def __init__(self, pipeline): + self.pipeline = pipeline + + def __enter__(self): + neuroglancer.set_server_bind_address("0.0.0.0") + viewer = neuroglancer.Viewer() + + viewer.actions.add("continue", step_next) + + with viewer.config_state.txn() as s: + s.input_event_bindings.data_view["keyt"] = "continue" + with viewer.txn() as s: + s.layout = neuroglancer.row_layout( + [ + neuroglancer.column_layout( + [ + neuroglancer.LayerGroupViewer(layers=[]), + neuroglancer.LayerGroupViewer(layers=[]), + ] + ), + ] + ) + + try: + self.pipeline.setup(viewer) + except: + logger.error( + "something went wrong during the setup of the pipeline, calling tear down" + ) + self.pipeline.internal_teardown() + logger.debug("tear down completed") + raise + + print(viewer) + return self.pipeline + + def __exit__(self, type, value, traceback): + logger.debug("leaving context, tearing down pipeline") + self.pipeline.internal_teardown() + logger.debug("tear down completed") diff --git a/gunpowder/neuroglancer/__init__.py b/gunpowder/neuroglancer/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/gunpowder/neuroglancer/add_layer.py b/gunpowder/neuroglancer/add_layer.py new file mode 100644 index 00000000..1a77098a --- /dev/null +++ b/gunpowder/neuroglancer/add_layer.py @@ -0,0 +1,296 @@ +from .scale_pyramid import ScalePyramid +import neuroglancer + +""" +TAKEN FROM funlib.show.neuroglancer +funlib.show.neuroglancer should probably be made installable +and used as an optional dependency +""" + + +rgb_shader_code = ''' +void main() { + emitRGB( + %f*vec3( + toNormalized(getDataValue(%i)), + toNormalized(getDataValue(%i)), + toNormalized(getDataValue(%i))) + ); +}''' + +color_shader_code = ''' +void main() { + emitRGBA( + vec4( + %f, %f, %f, + toNormalized(getDataValue())) + ); +}''' + +binary_shader_code = ''' +void main() { + emitGrayscale(255.0*toNormalized(getDataValue())); +}''' + +heatmap_shader_code = ''' +void main() { + float v = toNormalized(getDataValue(0)); + vec4 rgba = vec4(0,0,0,0); + if (v != 0.0) { + rgba = vec4(colormapJet(v), 1.0); + } + emitRGBA(rgba); +}''' + + +def parse_dims(array): + + if type(array) == list: + array = array[0] + + dims = len(array.data.shape) + spatial_dims = array.spec.roi.dims + channel_dims = dims - spatial_dims + + print("dims :", dims) + print("spatial dims:", spatial_dims) + print("channel dims:", channel_dims) + + return dims, spatial_dims, channel_dims + + +def create_coordinate_space(array, spatial_dim_names, channel_dim_names, unit): + + dims, spatial_dims, channel_dims = parse_dims(array) + assert spatial_dims > 0 + + if channel_dims > 0: + channel_names = channel_dim_names[-channel_dims:] + else: + channel_names = [] + spatial_names = spatial_dim_names[-spatial_dims:] + names = channel_names + spatial_names + units = [""] * channel_dims + [unit] * spatial_dims + scales = [1] * channel_dims + list(array.spec.voxel_size) + + print("Names :", names) + print("Units :", units) + print("Scales :", scales) + + return neuroglancer.CoordinateSpace( + names=names, + units=units, + scales=scales) + + +def create_shader_code( + shader, + channel_dims, + rgb_channels=None, + color=None, + scale_factor=1.0): + + if shader is None: + if channel_dims > 1: + shader = 'rgb' + else: + return None + + if rgb_channels is None: + rgb_channels = [0, 1, 2] + + if shader == 'rgb': + return rgb_shader_code % ( + scale_factor, + rgb_channels[0], + rgb_channels[1], + rgb_channels[2]) + + if shader == 'color': + assert color is not None, \ + "You have to pass argument 'color' to use the color shader" + return color_shader_code % ( + color[0], + color[1], + color[2], + ) + + if shader == 'binary': + return binary_shader_code + + if shader == 'heatmap': + return heatmap_shader_code + + +def add_layer( + context, + array, + name, + spatial_dim_names=None, + channel_dim_names=None, + opacity=None, + shader=None, + rgb_channels=None, + color=None, + visible=True, + value_scale_factor=1.0, + units='nm'): + + """Add a layer to a neuroglancer context. + + Args: + + context: + + The neuroglancer context to add a layer to, as obtained by + ``viewer.txn()``. + + array: + + A ``daisy``-like array, containing attributes ``roi``, + ``voxel_size``, and ``data``. If a list of arrays is given, a + ``ScalePyramid`` layer is generated. + + name: + + The name of the layer. + + spatial_dim_names: + + The names of the spatial dimensions. Defaults to ``['t', 'z', 'y', + 'x']``. The last elements of this list will be used (e.g., if your + data is 2D, the channels will be ``['y', 'x']``). + + channel_dim_names: + + The names of the non-spatial (channel) dimensions. Defaults to + ``['b^', 'c^']``. The last elements of this list will be used + (e.g., if your data is 2D but the shape of the array is 3D, the + channels will be ``['c^']``). + + opacity: + + A float to define the layer opacity between 0 and 1. + + shader: + + A string to be used as the shader. Possible values are: + + None : neuroglancer's default shader + 'rgb' : An RGB shader on dimension `'c^'`. See argument + ``rgb_channels``. + 'color' : Shows intensities as a constant color. See argument + ``color``. + 'binary' : Shows a binary image as black/white. + 'heatmap': Shows an intensity image as a jet color map. + + rgb_channels: + + Which channels to use for RGB (default is ``[0, 1, 2]``). + + color: + + A list of floats representing the RGB values for the constant color + shader. + + visible: + + A bool which defines the initial layer visibility. + + value_scale_factor: + + A float to scale array values with for visualization. + + units: + + The units used for resolution and offset. + """ + + if channel_dim_names is None: + channel_dim_names = ["b", "c^"] + if spatial_dim_names is None: + spatial_dim_names = ["t", "z", "y", "x"] + + if rgb_channels is None: + rgb_channels = [0, 1, 2] + + is_multiscale = type(array) == list + + dims, spatial_dims, channel_dims = parse_dims(array) + + if is_multiscale: + + dimensions = [] + for a in array: + dimensions.append( + create_coordinate_space( + a, + spatial_dim_names, + channel_dim_names, + units)) + + # why only one offset, shouldn't that be a list? + voxel_offset = [0] * channel_dims + \ + list(array[0].roi.offset / array[0].voxel_size) + + layer = ScalePyramid( + [ + neuroglancer.LocalVolume( + data=a.data, + voxel_offset=voxel_offset, + dimensions=array_dims + ) + for a, array_dims in zip(array, dimensions) + ] + ) + + else: + + voxel_offset = [0] * channel_dims + \ + list(array.spec.roi.offset / array.spec.voxel_size) + + dimensions = create_coordinate_space( + array, + spatial_dim_names, + channel_dim_names, + units) + + layer = neuroglancer.LocalVolume( + data=array.data, + voxel_offset=voxel_offset, + dimensions=dimensions, + ) + + shader_code = create_shader_code( + shader, + channel_dims, + rgb_channels, + color, + value_scale_factor) + + if opacity is not None: + if shader_code is None: + context.layers.append( + name=name, + layer=layer, + visible=visible, + opacity=opacity) + else: + context.layers.append( + name=name, + layer=layer, + visible=visible, + shader=shader_code, + opacity=opacity) + else: + if shader_code is None: + context.layers.append( + name=name, + layer=layer, + visible=visible) + else: + context.layers.append( + name=name, + layer=layer, + visible=visible, + shader=shader_code) diff --git a/gunpowder/neuroglancer/event.py b/gunpowder/neuroglancer/event.py new file mode 100644 index 00000000..3ad71cf4 --- /dev/null +++ b/gunpowder/neuroglancer/event.py @@ -0,0 +1,10 @@ +import threading + +BATCH_STEP = threading.Event() + +def step_next(event): + BATCH_STEP.set() + +def wait_for_step(): + BATCH_STEP.wait() + BATCH_STEP.clear() \ No newline at end of file diff --git a/gunpowder/neuroglancer/scale_pyramid.py b/gunpowder/neuroglancer/scale_pyramid.py new file mode 100644 index 00000000..133342d2 --- /dev/null +++ b/gunpowder/neuroglancer/scale_pyramid.py @@ -0,0 +1,99 @@ +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.debug("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.debug("min_voxel_size: %s", self.min_voxel_size) + logger.debug("scale keys: %s", self.volume_layers.keys()) + logger.debug(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(",")) + + return self.volume_layers[scale].get_encoded_subvolume( + data_format, start, end, scale_key=",".join(("1",) * self.dims) + ) + + 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() \ No newline at end of file diff --git a/gunpowder/neuroglancer/visualize.py b/gunpowder/neuroglancer/visualize.py new file mode 100644 index 00000000..a4e2397b --- /dev/null +++ b/gunpowder/neuroglancer/visualize.py @@ -0,0 +1,79 @@ +import neuroglancer +import gunpowder as gp + +import numpy as np + +from typing import Union + + +def visualize(viewer, batch_or_request): + going_up = isinstance(batch_or_request, gp.BatchRequest) + with viewer.txn() as s: + # reverse order for raw so we can set opacity to 1, this + # way higher res raw replaces low res when available + for name, array_or_spec in batch_or_request.items(): + if going_up: + name = f"request-{name}" + opacity = 0 + spec = array_or_spec.copy() + if spec.voxel_size is None: + spec.voxel_size = spec.roi.shape + data = np.zeros(spec.roi.shape / spec.voxel_size) + else: + name = f"batch-{name}" + opacity = 0.5 + spec = array_or_spec.spec.copy() + data = array_or_spec.data + + channel_dims = len(data.shape) - len(spec.voxel_size) + assert channel_dims <= 1 + + dims = neuroglancer.CoordinateSpace( + names=["c^", "z", "y", "x"][-len(data.shape) :], + units="nm", + scales=tuple([1] * channel_dims) + tuple(spec.voxel_size), + ) + + local_vol = neuroglancer.LocalVolume( + data=data, + voxel_offset=tuple([0] * channel_dims) + + tuple(spec.roi.begin / spec.voxel_size), + dimensions=dims, + ) + + s.layers[str(name)] = neuroglancer.ImageLayer( + source=local_vol, opacity=opacity + ) + + s.layout = neuroglancer.row_layout( + [ + neuroglancer.column_layout( + [ + neuroglancer.LayerGroupViewer( + layers=[ + str(l.name) + for l in s.layers + if not ( + l.name.startswith("batch") + or l.name.startswith("request") + ) + ] + ), + ] + ), + neuroglancer.column_layout( + [ + neuroglancer.LayerGroupViewer( + layers=[ + str(l.name) + for l in s.layers + if ( + (l.name.startswith("batch") and not going_up) + or (l.name.startswith("request") and going_up) + ) + ] + ), + ] + ), + ] + ) diff --git a/gunpowder/nodes/batch_provider.py b/gunpowder/nodes/batch_provider.py index 304e1e3a..3faf694c 100644 --- a/gunpowder/nodes/batch_provider.py +++ b/gunpowder/nodes/batch_provider.py @@ -12,6 +12,8 @@ from gunpowder.array_spec import ArraySpec from gunpowder.graph import GraphKey from gunpowder.graph_spec import GraphSpec +from gunpowder.neuroglancer.event import wait_for_step +from gunpowder.neuroglancer.visualize import visualize logger = logging.getLogger(__name__) @@ -52,6 +54,8 @@ class BatchProvider(object): instead. """ + viewer = None + def add_upstream_provider(self, provider): self.get_upstream_providers().append(provider) return provider @@ -70,7 +74,7 @@ def remove_placeholders(self): return True return self._remove_placeholders - def setup(self): + def setup(self, viewer=None): """To be implemented in subclasses. Called during initialization of the DAG. Callees can assume that all @@ -190,6 +194,9 @@ def request_batch(self, request): upstream_request = request.copy() if self.remove_placeholders: upstream_request.remove_placeholders() + + self.observe_request(request) + batch = self.provide(upstream_request) request.remove_placeholders() @@ -198,6 +205,8 @@ def request_batch(self, request): self.remove_unneeded(batch, request) + self.observe_batch(batch) + logger.debug("%s provides %s", self.name(), batch) except Exception as e: @@ -210,6 +219,21 @@ def request_batch(self, request): return batch + def setup_viewer(self, viewer): + self.viewer = viewer + + def observe_request(self, request): + if self.viewer is not None: + print("Waiting for step...") + visualize(self.viewer, request) + wait_for_step() + + def observe_batch(self, batch): + if self.viewer is not None: + print("Waiting for step...") + visualize(self.viewer, batch) + wait_for_step() + def set_seeds(self, request): seed = request.random_seed random.seed(seed) diff --git a/gunpowder/nodes/zarr_source.py b/gunpowder/nodes/zarr_source.py index b7133580..187c74dc 100644 --- a/gunpowder/nodes/zarr_source.py +++ b/gunpowder/nodes/zarr_source.py @@ -5,6 +5,7 @@ from gunpowder.roi import Roi from gunpowder.array import Array from gunpowder.array_spec import ArraySpec +from gunpowder.neuroglancer.add_layer import add_layer from .batch_provider import BatchProvider from zarr._storage.store import BaseStore @@ -124,6 +125,17 @@ def setup(self): self.provides(array_key, spec) + def setup_viewer(self, viewer): + self.viewer = viewer + with viewer.txn() as s: + with self._open_file(self.store) as data_file: + for array_key, ds_name in self.datasets.items(): + if ds_name not in data_file: + raise RuntimeError("%s not in %s" % (ds_name, self.store)) + + spec = self.__read_spec(array_key, data_file, ds_name) + add_layer(s, Array(data_file[ds_name], spec), f"{array_key}_SOURCE") + def provide(self, request): timing = Timing(self) timing.start() diff --git a/gunpowder/pipeline.py b/gunpowder/pipeline.py index cad87f1a..73d5dc2f 100644 --- a/gunpowder/pipeline.py +++ b/gunpowder/pipeline.py @@ -77,7 +77,7 @@ def copy(self): return pipeline - def setup(self): + def setup(self, viewer=None): """Connect all batch providers in the pipeline and call setup for each, from source to sink.""" @@ -94,6 +94,8 @@ def connect(node): def node_setup(node): try: node.output.setup() + if viewer is not None: + node.output.setup_viewer(viewer) except Exception as e: raise PipelineSetupError(node.output) from e