def __init__(self, dataset: "Dataset",
properties: LayerProperties) -> None:
"""
Do not use this constructor manually. Instead use `Dataset``.add_layer()` to create a `Layer`.
"""
# It is possible that the properties on disk do not contain the number of channels.
# Therefore, the parameter is optional. However at this point, 'num_channels' was already inferred.
assert properties.num_channels is not None
self._name: str = (
properties.name
) # The name is also stored in the properties, but the name is required to get the properties.
self._dataset = dataset
self._dtype_per_channel = _element_class_to_dtype_per_channel(
properties.element_class, properties.num_channels)
self._mags: Dict[Mag, MagView] = {}
self.path.mkdir(parents=True, exist_ok=True)
for mag in properties.mags:
self._setup_mag(Mag(mag.mag))
# Only keep the properties of mags that were initialized.
# Sometimes the directory of a mag is removed from disk manually, but the properties are not updated.
self._properties.mags = [
res for res in self._properties.mags if Mag(res.mag) in self._mags
]
def test_buffered_slice_writer() -> None:
test_img = np.arange(24 * 24).reshape(24, 24).astype(np.uint16) + 1
dtype = test_img.dtype
origin = Vec3Int.zeros()
layer_name = "color"
mag = Mag(1)
dataset_dir = TESTOUTPUT_DIR / "buffered_slice_writer"
dataset_path = str(dataset_dir / layer_name / mag.to_layer_name())
rmtree(dataset_dir)
ds = Dataset(dataset_dir, voxel_size=(1, 1, 1))
mag_view = ds.add_layer("color", COLOR_CATEGORY,
dtype_per_channel=dtype).add_mag(mag)
with mag_view.get_buffered_slice_writer(absolute_offset=origin) as writer:
for i in range(13):
writer.send(test_img)
with wkw.Dataset.open(dataset_path, wkw.Header(dtype)) as data:
try:
read_data = data.read(origin, (24, 24, 13))
if read_data[read_data.nonzero()].size != 0:
raise AssertionError(
"Nothing should be written on the disk. But found data with shape: {}"
.format(read_data.shape))
except wkw.wkw.WKWException:
pass
for i in range(13, 32):
writer.send(test_img)
with wkw.Dataset.open(dataset_path, wkw.Header(dtype)) as data:
read_data = data.read(origin, (24, 24, 32))
assert np.squeeze(read_data).shape == (24, 24, 32), (
"The read data should have the shape: (24, 24, 32) "
"but has a shape of: {}".format(np.squeeze(read_data).shape))
assert read_data.size == read_data[read_data.nonzero()].size, (
"The read data contains zeros while the "
"written image has no zeros")
for i in range(32, 35):
writer.send(test_img)
with wkw.Dataset.open(dataset_path, wkw.Header(dtype)) as data:
read_data = data.read(origin, (24, 24, 35))
read_data = np.squeeze(read_data)
assert read_data.shape == (24, 24, 35), (
"The read data should have the shape: (24, 24, 35) "
"but has a shape of: {}".format(np.squeeze(read_data).shape))
assert read_data.size == read_data[read_data.nonzero()].size, (
"The read data contains zeros while the "
"written image has no zeros")
test_img_3d = np.zeros((test_img.shape[0], test_img.shape[1], 35))
for i in np.arange(35):
test_img_3d[:, :, i] = test_img
# check if the data are correct
assert np.array_equal(
test_img_3d, read_data), ("The data from the disk is not the same "
"as the data that should be written.")
def _extract_num_channels(
num_channels_in_properties: Optional[int],
path: Path,
layer: str,
mag: Optional[Union[int, Mag]],
) -> int:
# if a wk dataset is not created with this API, then it most likely doesn't have the attribute 'numChannels' in the
# datasource-properties.json. In this case we need to extract the number of channels from the 'header.wkw'.
if num_channels_in_properties is not None:
return num_channels_in_properties
if mag is None:
# Unable to extract the 'num_channels' from the 'header.wkw' if the dataset has no magnifications.
# This should never be the case because wkw-datasets that are created without this API always have a magnification.
raise RuntimeError(
"Cannot extract the number of channels of a dataset without a properties file and without any magnifications"
)
mag = Mag(mag)
array_file_path = path / layer / mag.to_layer_name()
try:
array = BaseArray.open(array_file_path)
except ArrayException as e:
raise Exception(
f"The dataset you are trying to open does not have the attribute 'numChannels' for layer {layer}. "
f"However, this attribute is necessary. To mitigate this problem, it was tried to locate "
f"the file {array_file_path} to extract the num_channels from there. "
f"Since this file does not exist, the attempt to open the dataset failed. "
f"Please add the attribute manually to solve the problem. "
f"If the layer does not contain any data, you can also delete the layer and add it again.",
) from e
return array.info.num_channels
def add_mag_for_existing_files(
self,
mag: Union[int, str, list, tuple, np.ndarray, Mag],
) -> MagView:
"""
Creates a new mag based on already existing files.
Raises an IndexError if the specified `mag` does not exists.
"""
self.dataset._ensure_writable()
mag = Mag(mag)
assert (
mag not in self.mags
), f"Cannot add mag {mag} as it already exists for layer {self}"
self._setup_mag(mag)
mag_view = self._mags[mag]
mag_array_info = mag_view.info
self._properties.mags.append(
MagViewProperties(
mag=mag,
cube_length=(mag_array_info.shard_size.x
if mag_array_info.data_format == DataFormat.WKW
else None),
))
self.dataset._export_as_json()
return mag_view
def calculate_default_coarsest_mag(dataset_size: Vec3IntLike) -> Mag:
dataset_size = Vec3Int(dataset_size)
# The coarsest mag should have a size of ~ 100vx**2 per slice
coarsest_x_y = max(dataset_size[0], dataset_size[1])
# highest power of 2 larger (or equal) than coarsest_x_y divided by 100
# The calculated factor will be used for x, y and z here. If anisotropic downsampling takes place,
# the dimensions can still be downsampled independently according to the voxel_size.
return Mag(max(2 ** math.ceil(math.log(coarsest_x_y / 100, 2)), 4)) # at least 4
def test_import() -> None:
assert Vec3Int(1, 2, 3) == Vec3Int(1, 2, 3)
assert Vec3Int((1, 2, 3)) == Vec3Int(1, 2, 3)
assert Vec3Int([1, 2, 3]) == Vec3Int(1, 2, 3)
assert Vec3Int(i for i in [1, 2, 3]) == Vec3Int(1, 2, 3)
assert Vec3Int(np.array([1, 2, 3])) == Vec3Int(1, 2, 3)
assert Vec3Int(Mag(4)) == Vec3Int(4, 4, 4)
def downsample_unpadded_data(
buffer: np.ndarray, target_mag: Mag, interpolation_mode: InterpolationModes
) -> np.ndarray:
target_mag_np = np.array(target_mag.to_list())
current_dimension_size = np.array(buffer.shape[1:])
padding_size_for_downsampling = (
target_mag_np - (current_dimension_size % target_mag_np) % target_mag_np
)
padding_size_for_downsampling = list(zip([0, 0, 0], padding_size_for_downsampling))
buffer = np.pad(
buffer, pad_width=[(0, 0)] + padding_size_for_downsampling, mode="constant"
)
dimension_decrease = np.array([1] + target_mag.to_list())
downsampled_buffer_shape = np.array(buffer.shape) // dimension_decrease
downsampled_buffer = np.empty(dtype=buffer.dtype, shape=downsampled_buffer_shape)
for channel in range(buffer.shape[0]):
downsampled_buffer[channel] = downsample_cube(
buffer[channel], target_mag.to_list(), interpolation_mode
)
return downsampled_buffer
def get_mag(self, mag: Union[int, str, list, tuple, np.ndarray,
Mag]) -> MagView:
"""
Returns the `MagView` called `mag` of this layer. The return type is `webknossos.dataset.mag_view.MagView`.
This function raises an `IndexError` if the specified `mag` does not exist.
"""
mag = Mag(mag)
if mag not in self.mags.keys():
raise IndexError("The mag {} is not a mag of this layer".format(
mag.to_layer_name()))
return self.mags[mag]
def delete_mag(self, mag: Union[int, str, list, tuple, np.ndarray,
Mag]) -> None:
"""
Deletes the MagView from the `datasource-properties.json` and the data from disk.
This function raises an `IndexError` if the specified `mag` does not exist.
"""
self.dataset._ensure_writable()
mag = Mag(mag)
if mag not in self.mags.keys():
raise IndexError(
"Deleting mag {} failed. There is no mag with this name".
format(mag))
del self._mags[mag]
self._properties.mags = [
res for res in self._properties.mags if Mag(res.mag) != mag
]
self.dataset._export_as_json()
# delete files on disk
full_path = _find_mag_path_on_disk(self.dataset.path, self.name,
mag.to_layer_name())
rmtree(full_path)
def get_or_add_mag(
self,
mag: Union[int, str, list, tuple, np.ndarray, Mag],
chunk_size: Optional[Union[Vec3IntLike, int]] = None,
chunks_per_shard: Optional[Union[Vec3IntLike, int]] = None,
compress: Optional[bool] = None,
block_len: Optional[int] = None, # deprecated
file_len: Optional[int] = None, # deprecated
) -> MagView:
"""
Creates a new mag called and adds it to the dataset, in case it did not exist before.
Then, returns the mag.
See `add_mag` for more information.
"""
# normalize the name of the mag
mag = Mag(mag)
compression_mode = compress
chunk_size, chunks_per_shard = _get_sharding_parameters(
chunk_size=chunk_size,
chunks_per_shard=chunks_per_shard,
block_len=block_len,
file_len=file_len,
)
if mag in self._mags.keys():
assert (
chunk_size is None
or self._mags[mag].info.chunk_size == chunk_size
), f"Cannot get_or_add_mag: The mag {mag} already exists, but the chunk sizes do not match"
assert (
chunks_per_shard is None
or self._mags[mag].info.chunks_per_shard == chunks_per_shard
), f"Cannot get_or_add_mag: The mag {mag} already exists, but the chunks per shard do not match"
assert (
compression_mode is None
or self._mags[mag].info.compression_mode == compression_mode
), f"Cannot get_or_add_mag: The mag {mag} already exists, but the compression modes do not match"
return self.get_mag(mag)
else:
return self.add_mag(
mag,
chunk_size=chunk_size,
chunks_per_shard=chunks_per_shard,
compress=compression_mode or False,
)
def _setup_mag(self, mag: Mag) -> None:
# This method is used to initialize the mag when opening the Dataset. This does not create e.g. the wk_header.
mag_name = mag.to_layer_name()
self._assert_mag_does_not_exist_yet(mag)
try:
cls_array = BaseArray.get_class(self._properties.data_format)
info = cls_array.open(
_find_mag_path_on_disk(self.dataset.path, self.name,
mag_name)).info
self._mags[mag] = MagView(
self,
mag,
info.chunk_size,
info.chunks_per_shard,
info.compression_mode,
)
self._mags[mag]._read_only = self._dataset.read_only
except ArrayException as e:
logging.error(
f"Failed to setup magnification {mag_name}, which is specified in the datasource-properties.json. See {e}"
)
def upsample(
self,
from_mag: Mag,
finest_mag: Mag = Mag(1),
compress: bool = False,
sampling_mode: Union[str, SamplingModes] = SamplingModes.ANISOTROPIC,
align_with_other_layers: Union[bool, "Dataset"] = True,
buffer_shape: Optional[Vec3Int] = None,
buffer_edge_len: Optional[int] = None,
args: Optional[Namespace] = None,
*,
min_mag: Optional[Mag] = None,
) -> None:
"""
Upsamples the data starting from `from_mag` as long as the magnification is `>= finest_mag`.
There are three different `sampling_modes`:
- 'anisotropic' - The next magnification is chosen so that the width, height and depth of a downsampled voxel assimilate. For example, if the z resolution is worse than the x/y resolution, z won't be downsampled in the first downsampling step(s). As a basis for this method, the voxel_size from the datasource-properties.json is used.
- 'isotropic' - Each dimension is downsampled equally.
- 'constant_z' - The x and y dimensions are downsampled equally, but the z dimension remains the same.
`min_mag` is deprecated, please use `finest_mag` instead.
"""
assert (
from_mag in self.mags.keys()
), f"Failed to upsample data. The from_mag ({from_mag.to_layer_name()}) does not exist."
if min_mag is not None:
warn_deprecated("upsample(min_mag=…)", "upsample(finest_mag=…)")
assert finest_mag == Mag(
1
), "Cannot set both min_mag and finest_mag, please only use finest_mag."
finest_mag = min_mag
sampling_mode = SamplingModes.parse(sampling_mode)
voxel_size: Optional[Tuple[float, float, float]] = None
if sampling_mode == SamplingModes.ANISOTROPIC:
voxel_size = self.dataset.voxel_size
elif sampling_mode == SamplingModes.ISOTROPIC:
voxel_size = None
elif sampling_mode == SamplingModes.CONSTANT_Z:
finest_mag_with_fixed_z = finest_mag.to_list()
finest_mag_with_fixed_z[2] = from_mag.to_list()[2]
finest_mag = Mag(finest_mag_with_fixed_z)
voxel_size = None
else:
raise AttributeError(
f"Upsampling failed: {sampling_mode} is not a valid UpsamplingMode ({SamplingModes.ANISOTROPIC}, {SamplingModes.ISOTROPIC}, {SamplingModes.CONSTANT_Z})"
)
if buffer_shape is None and buffer_edge_len is not None:
buffer_shape = Vec3Int.full(buffer_edge_len)
dataset_to_align_with = self._get_dataset_from_align_with_other_layers(
align_with_other_layers)
mags_to_upsample = calculate_mags_to_upsample(from_mag, finest_mag,
dataset_to_align_with,
voxel_size)
for prev_mag, target_mag in zip([from_mag] + mags_to_upsample[:-1],
mags_to_upsample):
assert prev_mag > target_mag
assert target_mag not in self.mags
prev_mag_view = self.mags[prev_mag]
mag_factors = [
t / s
for (t, s) in zip(target_mag.to_list(), prev_mag.to_list())
]
# initialize the new mag
target_mag_view = self._initialize_mag_from_other_mag(
target_mag, prev_mag_view, compress)
# Get target view
target_view = target_mag_view.get_view()
# perform upsampling
with get_executor_for_args(args) as executor:
if buffer_shape is None:
buffer_shape = determine_buffer_shape(prev_mag_view.info)
func = named_partial(
upsample_cube_job,
mag_factors=mag_factors,
buffer_shape=buffer_shape,
)
prev_mag_view.get_view().for_zipped_chunks(
# this view is restricted to the bounding box specified in the properties
func,
target_view=target_view,
executor=executor,
progress_desc=
f"Upsampling from Mag {prev_mag} to Mag {target_mag}",
)
def download_dataset(
dataset_name: str,
organization_id: str,
sharing_token: Optional[str] = None,
bbox: Optional[BoundingBox] = None,
layers: Optional[List[str]] = None,
mags: Optional[List[Mag]] = None,
path: Optional[Union[PathLike, str]] = None,
exist_ok: bool = False,
) -> Dataset:
context = _get_context()
client = context.generated_client
dataset_info_response = dataset_info.sync_detailed(
organization_name=organization_id,
data_set_name=dataset_name,
client=client,
sharing_token=sharing_token,
)
assert dataset_info_response.status_code == 200, dataset_info_response
parsed = dataset_info_response.parsed
assert parsed is not None
datastore_client = context.get_generated_datastore_client(
parsed.data_store.url)
optional_datastore_token = sharing_token or context.datastore_token
actual_path = Path(dataset_name) if path is None else Path(path)
if actual_path.exists():
logger.warning(f"{actual_path} already exists, skipping download.")
return Dataset.open(actual_path)
voxel_size = cast(Tuple[float, float, float],
tuple(parsed.data_source.scale))
data_layers = parsed.data_source.data_layers
dataset = Dataset(actual_path,
name=parsed.name,
voxel_size=voxel_size,
exist_ok=exist_ok)
for layer_name in layers or [i.name for i in data_layers]:
response_layers = [i for i in data_layers if i.name == layer_name]
assert (
len(response_layers) > 0
), f"The provided layer name {layer_name} could not be found in the requested dataset."
assert (
len(response_layers) == 1
), f"The provided layer name {layer_name} was found multiple times in the requested dataset."
response_layer = response_layers[0]
category = cast(LayerCategoryType, response_layer.category)
layer = dataset.add_layer(
layer_name=layer_name,
category=category,
dtype_per_layer=response_layer.element_class,
num_channels=3 if response_layer.element_class == "uint24" else 1,
largest_segment_id=response_layer.additional_properties.get(
"largestSegmentId", None),
)
default_view_configuration_dict = None
if not isinstance(response_layer.default_view_configuration, Unset):
default_view_configuration_dict = (
response_layer.default_view_configuration.to_dict())
if default_view_configuration_dict is not None:
default_view_configuration = dataset_converter.structure(
default_view_configuration_dict, LayerViewConfiguration)
layer.default_view_configuration = default_view_configuration
if bbox is None:
response_bbox = response_layer.bounding_box
layer.bounding_box = BoundingBox(
response_bbox.top_left,
(response_bbox.width, response_bbox.height,
response_bbox.depth),
)
else:
assert isinstance(
bbox, BoundingBox
), f"Expected a BoundingBox object for the bbox parameter but got {type(bbox)}"
layer.bounding_box = bbox
if mags is None:
mags = [Mag(mag) for mag in response_layer.resolutions]
for mag in mags:
mag_view = layer.get_or_add_mag(
mag,
compress=True,
chunk_size=Vec3Int.full(32),
chunks_per_shard=_DOWNLOAD_CHUNK_SIZE // 32,
)
aligned_bbox = layer.bounding_box.align_with_mag(mag, ceil=True)
download_chunk_size_in_mag = _DOWNLOAD_CHUNK_SIZE * mag.to_vec3_int(
)
for chunk in track(
list(
aligned_bbox.chunk(download_chunk_size_in_mag,
download_chunk_size_in_mag)),
description=f"Downloading layer={layer.name} mag={mag}",
):
chunk_in_mag = chunk.in_mag(mag)
response = dataset_download.sync_detailed(
organization_name=organization_id,
data_set_name=dataset_name,
data_layer_name=layer_name,
mag=mag.to_long_layer_name(),
client=datastore_client,
token=optional_datastore_token,
x=chunk.topleft.x,
y=chunk.topleft.y,
z=chunk.topleft.z,
width=chunk_in_mag.size.x,
height=chunk_in_mag.size.y,
depth=chunk_in_mag.size.z,
)
assert response.status_code == 200, response
assert (
response.headers["missing-buckets"] == "[]"
), f"Download contained missing buckets {response.headers['missing-buckets']}."
data = np.frombuffer(response.content,
dtype=layer.dtype_per_channel).reshape(
layer.num_channels,
*chunk_in_mag.size,
order="F")
mag_view.write(data, absolute_offset=chunk.topleft)
return dataset
def add_mag(
self,
mag: Union[int, str, list, tuple, np.ndarray, Mag],
chunk_size: Optional[Union[Vec3IntLike,
int]] = None, # DEFAULT_CHUNK_SIZE,
chunks_per_shard: Optional[Union[
int, Vec3IntLike]] = None, # DEFAULT_CHUNKS_PER_SHARD,
compress: bool = False,
block_len: Optional[int] = None, # deprecated
file_len: Optional[int] = None, # deprecated
) -> MagView:
"""
Creates a new mag called and adds it to the layer.
The parameter `chunk_size`, `chunks_per_shard` and `compress` can be
specified to adjust how the data is stored on disk.
Note that writing compressed data which is not aligned with the blocks on disk may result in
diminished performance, as full blocks will automatically be read to pad the write actions. Alternatively,
you can call mag.compress() after all the data was written
The return type is `webknossos.dataset.mag_view.MagView`.
Raises an IndexError if the specified `mag` already exists.
"""
self.dataset._ensure_writable()
# normalize the name of the mag
mag = Mag(mag)
compression_mode = compress
chunk_size, chunks_per_shard = _get_sharding_parameters(
chunk_size=chunk_size,
chunks_per_shard=chunks_per_shard,
block_len=block_len,
file_len=file_len,
)
if chunk_size is None:
chunk_size = DEFAULT_CHUNK_SIZE
if chunks_per_shard is None:
if self.data_format == DataFormat.Zarr:
chunks_per_shard = DEFAULT_CHUNKS_PER_SHARD_ZARR
else:
chunks_per_shard = DEFAULT_CHUNKS_PER_SHARD
if chunk_size not in (Vec3Int.full(32), Vec3Int.full(64)):
warnings.warn(
"[INFO] `chunk_size` of `32, 32, 32` or `64, 64, 64` is recommended for optimal "
+ f"performance in webKnossos. Got {chunk_size}.")
self._assert_mag_does_not_exist_yet(mag)
self._create_dir_for_mag(mag)
mag_view = MagView(
self,
mag,
chunk_size=chunk_size,
chunks_per_shard=chunks_per_shard,
compression_mode=compression_mode,
create=True,
)
mag_view._array.ensure_size(
self.bounding_box.align_with_mag(mag).in_mag(mag).bottomright)
self._mags[mag] = mag_view
mag_array_info = mag_view.info
self._properties.mags += [
MagViewProperties(
mag=Mag(mag_view.name),
cube_length=(mag_array_info.shard_size.x
if mag_array_info.data_format == DataFormat.WKW
else None),
)
]
self.dataset._export_as_json()
return self._mags[mag]
def calculate_mags_to_downsample(
from_mag: Mag,
coarsest_mag: Mag,
dataset_to_align_with: Optional["Dataset"],
voxel_size: Optional[Tuple[float, float, float]],
) -> List[Mag]:
assert np.all(from_mag.to_np() <= coarsest_mag.to_np())
mags = []
current_mag = from_mag
if dataset_to_align_with is None:
mags_to_align_with = set()
else:
mags_to_align_with = set(
mag
for layer in dataset_to_align_with.layers.values()
for mag in layer.mags.keys()
)
mags_to_align_with_by_max_dim = {mag.max_dim: mag for mag in mags_to_align_with}
assert len(mags_to_align_with) == len(
mags_to_align_with_by_max_dim
), "Some layers contain different values for the same mag, this is not allowed."
while current_mag < coarsest_mag:
if current_mag.max_dim * 2 in mags_to_align_with_by_max_dim:
current_mag = mags_to_align_with_by_max_dim[current_mag.max_dim * 2]
if current_mag > coarsest_mag:
warnings.warn(
"The mag taken from another layer is larger in some dimensions than coarsest_mag."
)
elif voxel_size is None:
# In case the sampling mode is CONSTANT_Z or ISOTROPIC:
current_mag = Mag(np.minimum(current_mag.to_np() * 2, coarsest_mag.to_np()))
else:
# In case the sampling mode is ANISOTROPIC:
current_size = current_mag.to_np() * np.array(voxel_size)
min_value = np.min(current_size)
min_value_bitmask = np.array(current_size == min_value)
factor = min_value_bitmask + 1
# Calculate the two potential magnifications.
# Either, double all components or only double the smallest component.
all_voxel_sized = current_size * 2
min_voxel_sized = (
current_size * factor
) # only multiply the smallest dimension
all_voxel_sized_ratio = np.max(all_voxel_sized) / np.min(all_voxel_sized)
min_voxel_sized_ratio = np.max(min_voxel_sized) / np.min(min_voxel_sized)
# The smaller the ratio between the smallest dimension and the largest dimension, the better.
if all_voxel_sized_ratio < min_voxel_sized_ratio:
# Multiply all dimensions with "2"
new_mag = Mag(np.minimum(current_mag.to_np() * 2, coarsest_mag.to_np()))
else:
# Multiply only the minimal dimension by "2".
new_mag = Mag(
np.minimum(current_mag.to_np() * factor, coarsest_mag.to_np())
)
# In case of isotropic resolution but anisotropic mag we need to ensure unique max dims.
# current mag: 4-4-2, voxel_size: 1,1,1 -> new_mag: 4-4-4, therefore we skip this entry.
if new_mag.max_dim == current_mag.max_dim:
current_mag = new_mag
continue
if new_mag == current_mag:
raise RuntimeError(
f"The coarsest mag {coarsest_mag} can not be reached from {current_mag} with voxel_size {voxel_size}!"
)
current_mag = new_mag
mags += [current_mag]
return mags
def _create_dir_for_mag(
self, mag: Union[int, str, list, tuple, np.ndarray, Mag]) -> None:
mag = Mag(mag).to_layer_name()
full_path = self.path / mag
full_path.mkdir(parents=True, exist_ok=True)
dataset_converter = cattr.Converter()
# register (un-)structure hooks for non-attr-classes
bbox_to_wkw: Callable[[BoundingBox], dict] = lambda o: o.to_wkw_dict()
dataset_converter.register_unstructure_hook(BoundingBox, bbox_to_wkw)
dataset_converter.register_structure_hook(
BoundingBox, lambda d, _: BoundingBox.from_wkw_dict(d))
def mag_unstructure(mag: Mag) -> List[int]:
return mag.to_list()
dataset_converter.register_unstructure_hook(Mag, mag_unstructure)
dataset_converter.register_structure_hook(Mag, lambda d, _: Mag(d))
vec3int_to_array: Callable[[Vec3Int], List[int]] = lambda o: o.to_list()
dataset_converter.register_unstructure_hook(Vec3Int, vec3int_to_array)
dataset_converter.register_structure_hook(
Vec3Int, lambda d, _: Vec3Int.full(d)
if isinstance(d, int) else Vec3Int(d))
dataset_converter.register_structure_hook_func(
lambda d: d == LayerCategoryType, lambda d, _: str(d))
# Register (un-)structure hooks for attr-classes to bring the data into the expected format.
# The properties on disk (in datasource-properties.json) use camel case for the names of the attributes.
# However, we use snake case for the attribute names in python.
# This requires that the names of the attributes are renamed during (un-)structuring.
# Additionally we only want to unstructure attributes which don't have the default value
def downsample(
self,
from_mag: Optional[Mag] = None,
coarsest_mag: Optional[Mag] = None,
interpolation_mode: str = "default",
compress: bool = True,
sampling_mode: Union[str, SamplingModes] = SamplingModes.ANISOTROPIC,
align_with_other_layers: Union[bool, "Dataset"] = True,
buffer_shape: Optional[Vec3Int] = None,
force_sampling_scheme: bool = False,
args: Optional[Namespace] = None,
allow_overwrite: bool = False,
only_setup_mags: bool = False,
) -> None:
"""
Downsamples the data starting from `from_mag` until a magnification is `>= max(coarsest_mag)`.
There are three different `sampling_modes`:
- 'anisotropic' - The next magnification is chosen so that the width, height and depth of a downsampled voxel assimilate. For example, if the z resolution is worse than the x/y resolution, z won't be downsampled in the first downsampling step(s). As a basis for this method, the voxel_size from the datasource-properties.json is used.
- 'isotropic' - Each dimension is downsampled equally.
- 'constant_z' - The x and y dimensions are downsampled equally, but the z dimension remains the same.
See `downsample_mag` for more information.
Example:
```python
from webknossos import SamplingModes
# ...
# let 'layer' be a `Layer` with only `Mag(1)`
assert "1" in self.mags.keys()
layer.downsample(
coarsest_mag=Mag(4),
sampling_mode=SamplingModes.ISOTROPIC
)
assert "2" in self.mags.keys()
assert "4" in self.mags.keys()
```
"""
if from_mag is None:
assert (
len(self.mags.keys()) > 0
), "Failed to downsample data because no existing mag was found."
from_mag = max(self.mags.keys())
assert (
from_mag in self.mags.keys()
), f"Failed to downsample data. The from_mag ({from_mag.to_layer_name()}) does not exist."
if coarsest_mag is None:
coarsest_mag = calculate_default_coarsest_mag(
self.bounding_box.size)
sampling_mode = SamplingModes.parse(sampling_mode)
if self._properties.bounding_box.size.z == 1:
if sampling_mode != SamplingModes.CONSTANT_Z:
warnings.warn(
"The sampling_mode was changed to 'CONSTANT_Z'. Downsampling 2D data with a different sampling mode mixes in black and thus leads to darkened images."
)
sampling_mode = SamplingModes.CONSTANT_Z
voxel_size: Optional[Tuple[float, float, float]] = None
if sampling_mode == SamplingModes.ANISOTROPIC:
voxel_size = self.dataset.voxel_size
elif sampling_mode == SamplingModes.ISOTROPIC:
voxel_size = None
elif sampling_mode == SamplingModes.CONSTANT_Z:
coarsest_mag_with_fixed_z = coarsest_mag.to_list()
coarsest_mag_with_fixed_z[2] = from_mag.to_list()[2]
coarsest_mag = Mag(coarsest_mag_with_fixed_z)
voxel_size = None
else:
raise AttributeError(
f"Downsampling failed: {sampling_mode} is not a valid SamplingMode ({SamplingModes.ANISOTROPIC}, {SamplingModes.ISOTROPIC}, {SamplingModes.CONSTANT_Z})"
)
dataset_to_align_with = self._get_dataset_from_align_with_other_layers(
align_with_other_layers)
mags_to_downsample = calculate_mags_to_downsample(
from_mag, coarsest_mag, dataset_to_align_with, voxel_size)
if len(set([max(m.to_list())
for m in mags_to_downsample])) != len(mags_to_downsample):
msg = (
f"The downsampling scheme contains multiple magnifications with the same maximum value. This is not supported by webknossos. "
f"Consider using a different sampling mode (e.g. {SamplingModes.ISOTROPIC}). "
f"The calculated downsampling scheme is: {[m.to_layer_name() for m in mags_to_downsample]}"
)
if force_sampling_scheme:
warnings.warn(msg)
else:
raise RuntimeError(msg)
for prev_mag, target_mag in zip([from_mag] + mags_to_downsample[:-1],
mags_to_downsample):
self.downsample_mag(
from_mag=prev_mag,
target_mag=target_mag,
interpolation_mode=interpolation_mode,
compress=compress,
buffer_shape=buffer_shape,
args=args,
allow_overwrite=allow_overwrite,
only_setup_mag=only_setup_mags,
)
def mag_unstructure(mag: Mag) -> List[int]:
return mag.to_list()
def downsample_mag(
self,
from_mag: Mag,
target_mag: Mag,
interpolation_mode: str = "default",
compress: bool = True,
buffer_shape: Optional[Vec3Int] = None,
args: Optional[Namespace] = None,
allow_overwrite: bool = False,
only_setup_mag: bool = False,
) -> None:
"""
Performs a single downsampling step from `from_mag` to `target_mag`.
The supported `interpolation_modes` are:
- "median"
- "mode"
- "nearest"
- "bilinear"
- "bicubic"
The `args` can contain information to distribute the computation.
If allow_overwrite is True, an existing Mag may be overwritten.
If only_setup_mag is True, the magnification is created, but left
empty. This parameter can be used to prepare for parallel downsampling
of multiple layers while avoiding parallel writes with outdated updates
to the datasource-properties.json file.
"""
assert (
from_mag in self.mags.keys()
), f"Failed to downsample data. The from_mag ({from_mag.to_layer_name()}) does not exist."
parsed_interpolation_mode = parse_interpolation_mode(
interpolation_mode, self.category)
assert from_mag <= target_mag
assert (
allow_overwrite or target_mag not in self.mags
), "The target mag already exists. Pass allow_overwrite=True if you want to overwrite it."
prev_mag_view = self.mags[from_mag]
mag_factors = target_mag.to_vec3_int() // from_mag.to_vec3_int()
if target_mag in self.mags.keys() and allow_overwrite:
target_mag_view = self.get_mag(target_mag)
else:
# initialize the new mag
target_mag_view = self._initialize_mag_from_other_mag(
target_mag, prev_mag_view, compress)
if only_setup_mag:
return
bb_mag1 = self.bounding_box.align_with_mag(target_mag, ceil=True)
# Get target view
target_view = target_mag_view.get_view(
absolute_offset=bb_mag1.topleft,
size=bb_mag1.size,
)
source_view = prev_mag_view.get_view(
absolute_offset=bb_mag1.topleft,
size=bb_mag1.size,
read_only=True,
)
# perform downsampling
with get_executor_for_args(args) as executor:
if buffer_shape is None:
buffer_shape = determine_buffer_shape(prev_mag_view.info)
func = named_partial(
downsample_cube_job,
mag_factors=mag_factors,
interpolation_mode=parsed_interpolation_mode,
buffer_shape=buffer_shape,
)
source_view.for_zipped_chunks(
# this view is restricted to the bounding box specified in the properties
func,
target_view=target_view,
executor=executor,
progress_desc=
f"Downsampling layer {self.name} from Mag {from_mag} to Mag {target_mag}",
)
def merge_with_fallback_layer(
output_path: Path,
volume_annotation_path: Path,
segmentation_layer_path: Path,
) -> MagView:
assert not output_path.exists(), f"Dataset at {output_path} already exists"
# Prepare output dataset by creatign a shallow copy of the dataset
# determined by segmentation_layer_path, but do a deep copy of
# segmentation_layer_path itself (so that we can mutate it).
input_segmentation_dataset = wk.Dataset.open(
segmentation_layer_path.parent)
time_start("Prepare output dataset")
output_dataset = input_segmentation_dataset.shallow_copy_dataset(
output_path,
name=output_path.name,
make_relative=True,
layers_to_ignore=[segmentation_layer_path.name],
)
output_layer = output_dataset.add_copy_layer(segmentation_layer_path,
segmentation_layer_path.name)
time_stop("Prepare output dataset")
input_segmentation_mag = input_segmentation_dataset.get_layer(
segmentation_layer_path.name).get_finest_mag()
with temporary_annotation_view(
volume_annotation_path) as input_annotation_layer:
input_annotation_mag = input_annotation_layer.get_finest_mag()
bboxes = [
bbox.in_mag(input_annotation_mag._mag)
for bbox in input_annotation_mag.get_bounding_boxes_on_disk()
]
output_mag = output_layer.get_mag(input_segmentation_mag.mag)
cube_size = output_mag.info.chunk_size[
0] * output_mag.info.chunks_per_shard[0]
chunks_with_bboxes = BoundingBox.group_boxes_with_aligned_mag(
bboxes, Mag(cube_size))
assert (input_annotation_mag.info.chunks_per_shard == Vec3Int.ones()
), "volume annotation must have file_len=1"
assert (input_annotation_mag.info.voxel_type ==
input_segmentation_mag.info.voxel_type
), "Volume annotation must have same dtype as fallback layer"
chunk_count = 0
for chunk, bboxes in chunks_with_bboxes.items():
chunk_count += 1
logger.info(f"Processing chunk {chunk_count}...")
time_start("Read chunk")
data_buffer = output_mag.read(chunk.topleft,
chunk.size)[0, :, :, :]
time_stop("Read chunk")
time_start("Read/merge bboxes")
for bbox in bboxes:
read_data = input_annotation_mag.read(bbox.topleft, bbox.size)
data_buffer[bbox.offset(
-chunk.topleft).to_slices()] = read_data
time_stop("Read/merge bboxes")
time_start("Write chunk")
output_mag.write(data_buffer, chunk.topleft)
time_stop("Write chunk")
return output_mag
def test_mag_constructor() -> None:
mag = Mag(16)
assert mag.to_list() == [16, 16, 16]
mag = Mag("256")
assert mag.to_list() == [256, 256, 256]
mag = Mag("16-2-4")
assert mag.to_list() == [16, 2, 4]
mag1 = Mag("16-2-4")
mag2 = Mag("8-2-4")
assert mag1 > mag2
assert mag1.to_layer_name() == "16-2-4"
assert np.all(mag1.to_np() == np.array([16, 2, 4]))
assert mag1 == Mag(mag1)
assert mag1 == Mag(mag1.to_np())
