def run_distributed_export(self, block_roi: Slice5D): from lazyflow.distributed.TaskOrchestrator import TaskOrchestrator orchestrator = TaskOrchestrator() n5_file_path = Path(self.OutputFilenameFormat.value).with_suffix(".n5") output_meta = self.ImageToExport.meta if orchestrator.rank == 0: output_shape = output_meta.getShape5D() block_shape = block_roi.clamped(output_shape.to_slice_5d()).shape with z5py.File(n5_file_path, "w") as f: ds = f.create_dataset( self.OutputInternalPath.value, shape=output_meta.shape, chunks=block_shape.to_tuple(output_meta.getAxisKeys()), dtype=output_meta.dtype.__name__, ) ds.attrs["axes"] = list(reversed(output_meta.getAxisKeys())) ds[...] = 1 # FIXME: for some reason setting to 0 does nothing cutout = self.get_roi() orchestrator.orchestrate(cutout.split(block_shape=block_shape)) else: def process_tile(tile: Slice5D, rank: int): self.set_roi(tile) slices = tile.to_slices(output_meta.getAxisKeys()) with z5py.File(n5_file_path, "r+") as n5_file: dataset = n5_file[self.OutputInternalPath.value] dataset[slices] = self.ImageToExport.value orchestrator.start_as_worker(process_tile)
def retrieve(self, roi: Slice5D, address_mode: AddressMode = AddressMode.BLACK) -> Array5D: # FIXME: Remove address_mode or implement all variations and make feature extractors use the correct one out = self._allocate(roi.defined_with(self.shape).translated(-self.location), fill_value=0) local_data_roi = roi.clamped(self.roi).translated(-self.location) for tile in local_data_roi.get_tiles(self.tile_shape): tile_within_bounds = tile.clamped(self.shape) tile_data = self.get_tile(tile_within_bounds) out.set(tile_data, autocrop=True) out.setflags(write=False) return out.translated(self.location)
def _get_tile(self, tile: Slice5D) -> Array5D: first_layer_idx = bisect.bisect_left(self.layer_offsets, tile.start[self.stack_axis]) out = self._allocate(roi=tile, fill_value=0) for layer, layer_offset in zip(self.layers[first_layer_idx:], self.layer_offsets[first_layer_idx:]): if layer_offset > tile.stop[self.stack_axis]: break layer_tile = tile.clamped(layer.roi) layer_data = layer.retrieve(layer_tile) out.set(layer_data, autocrop=True) return out