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
