def downsample_cube_job(
args: Tuple[View, View, int],
mag_factors: Vec3Int,
interpolation_mode: InterpolationModes,
buffer_shape: Vec3Int,
) -> None:
(source_view, target_view, _i) = args
try:
num_channels = target_view.info.num_channels
shape = (num_channels,) + target_view.bounding_box.in_mag(
target_view.mag
).size.to_tuple()
file_buffer = np.zeros(shape, target_view.get_dtype())
tiles = product(
*(
list(range(0, math.ceil(len / buffer_edge_len)))
for len, buffer_edge_len in zip(shape[-3:], buffer_shape)
)
)
for tile in tiles:
target_offset = Vec3Int(tile) * buffer_shape
source_offset = mag_factors * target_offset
source_size = source_view.bounding_box.in_mag(source_view.mag).size
source_size = (mag_factors * buffer_shape).pairmin(
source_size - source_offset
)
bbox = BoundingBox(source_offset, source_size)
cube_buffer_channels = source_view.read(
relative_bounding_box=bbox.from_mag_to_mag1(source_view.mag),
)
for channel_index in range(num_channels):
cube_buffer = cube_buffer_channels[channel_index]
if not np.all(cube_buffer == 0):
# Downsample the buffer
data_cube = downsample_cube(
cube_buffer,
mag_factors.to_list(),
interpolation_mode,
)
buffer_bbox = BoundingBox(target_offset, data_cube.shape)
file_buffer[(channel_index,) + buffer_bbox.to_slices()] = data_cube
# Write the downsampled buffer to target
if source_view.info.num_channels == 1:
file_buffer = file_buffer[0] # remove channel dimension
target_view.write(file_buffer)
except Exception as exc:
logging.error(
f"Downsampling of target {target_view.bounding_box} failed with {exc}"
)
raise exc
def test_buffered_slice_reader_along_different_axis(tmp_path: Path) -> None:
test_cube = (np.random.random((3, 13, 13, 13)) * 100).astype(np.uint8)
cube_size_without_channel = Vec3Int(test_cube.shape[1:])
offset = Vec3Int(5, 10, 20)
for dim in [0, 1, 2]:
ds = Dataset(tmp_path / f"buffered_slice_reader_{dim}",
voxel_size=(1, 1, 1))
mag_view = ds.add_layer("color", COLOR_CATEGORY,
num_channels=3).add_mag(1)
mag_view.write(test_cube, absolute_offset=offset)
with mag_view.get_buffered_slice_reader(
buffer_size=5,
dimension=dim) as reader_a, mag_view.get_buffered_slice_reader(
absolute_bounding_box=BoundingBox(
offset, cube_size_without_channel),
buffer_size=5,
dimension=dim,
) as reader_b:
i = 0
for slice_data_a, slice_data_b in zip(reader_a, reader_b):
if dim == 0:
original_slice = test_cube[:, i, :, :]
elif dim == 1:
original_slice = test_cube[:, :, i, :]
else: # dim == 2
original_slice = test_cube[:, :, :, i]
i += 1
assert np.array_equal(slice_data_a, original_slice)
assert np.array_equal(slice_data_b, original_slice)
def upsample_cube_job(
args: Tuple[View, View, int],
mag_factors: List[float],
buffer_shape: Vec3Int,
) -> None:
(source_view, target_view, _i) = args
assert all(
1 >= f for f in mag_factors
), f"mag_factors ({mag_factors}) for upsampling must be smaller than 1"
try:
num_channels = target_view.info.num_channels
target_size = target_view.bounding_box.in_mag(target_view.mag).size
shape = (num_channels, ) + target_size.to_tuple()
file_buffer = np.zeros(shape, target_view.get_dtype())
tiles = product(*list([
list(range(0, math.ceil(len)))
for len in target_size.to_np() / buffer_shape.to_np()
]))
for tile in tiles:
target_offset = Vec3Int(tile) * buffer_shape
source_offset = _vec3int_mulf(target_offset, mag_factors)
source_size = source_view.bounding_box.in_mag(source_view.mag).size
source_size = _vec3int_mulf(
buffer_shape, mag_factors).pairmin(source_size - source_offset)
bbox = BoundingBox(source_offset, source_size)
cube_buffer_channels = source_view.read(
relative_bounding_box=bbox.from_mag_to_mag1(source_view.mag), )
for channel_index in range(num_channels):
cube_buffer = cube_buffer_channels[channel_index]
if not np.all(cube_buffer == 0):
# Upsample the buffer
inverse_factors = [int(1 / f) for f in mag_factors]
data_cube = upsample_cube(cube_buffer, inverse_factors)
buffer_offset = target_offset
buffer_end = buffer_offset + data_cube.shape
file_buffer[channel_index, buffer_offset[0]:buffer_end[0],
buffer_offset[1]:buffer_end[1],
buffer_offset[2]:buffer_end[2], ] = data_cube
# Write the upsampled buffer to target
if source_view.info.num_channels == 1:
file_buffer = file_buffer[0] # remove channel dimension
target_view.write(file_buffer)
except Exception as exc:
logging.error(
f"Upsampling of target {target_view.bounding_box} failed with {exc}"
)
raise exc
def test_export() -> None:
assert Vec3Int(1, 2, 3).x == 1
assert Vec3Int(1, 2, 3).y == 2
assert Vec3Int(1, 2, 3).z == 3
assert Vec3Int(1, 2, 3)[0] == 1
assert Vec3Int(1, 2, 3)[1] == 2
assert Vec3Int(1, 2, 3)[2] == 3
assert np.array_equal(Vec3Int(1, 2, 3).to_np(), np.array([1, 2, 3]))
assert Vec3Int(1, 2, 3).to_list() == [1, 2, 3]
assert Vec3Int(1, 2, 3).to_tuple() == (1, 2, 3)
def _nml_tree_to_wk_tree(
new_tree: "Tree",
nml_tree: wknml.Tree,
) -> None:
optional_attribute_list = [
"rotation",
"inVp",
"inMag",
"bitDepth",
"interpolation",
"time",
]
for nml_node in nml_tree.nodes:
node_id = nml_node.id
current_node = new_tree.add_node(
position=Vec3Int.from_vec3_float(nml_node.position),
_enforced_id=node_id,
radius=nml_node.radius,
)
for optional_attribute in optional_attribute_list:
if getattr(nml_node, optional_attribute) is not None:
setattr(
current_node,
optional_attribute,
getattr(nml_node, optional_attribute),
)
for edge in nml_tree.edges:
new_tree.add_edge(edge.source, edge.target)
def test_buffered_slice_writer_along_different_axis(tmp_path: Path) -> None:
test_cube = (np.random.random((3, 13, 13, 13)) * 100).astype(np.uint8)
cube_size_without_channel = test_cube.shape[1:]
offset = Vec3Int(5, 10, 20)
for dim in [0, 1, 2]:
ds = Dataset(tmp_path / f"buffered_slice_writer_{dim}",
voxel_size=(1, 1, 1))
mag_view = ds.add_layer("color", COLOR_CATEGORY,
num_channels=3).add_mag(1)
with mag_view.get_buffered_slice_writer(absolute_offset=offset,
buffer_size=5,
dimension=dim) as writer:
for i in range(cube_size_without_channel[dim]):
if dim == 0:
current_slice = test_cube[:, i, :, :]
elif dim == 1:
current_slice = test_cube[:, :, i, :]
else: # dim == 2
current_slice = test_cube[:, :, :, i]
writer.send(current_slice)
assert np.array_equal(
mag_view.read(absolute_offset=offset,
size=cube_size_without_channel),
test_cube,
)
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_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 _get_sharding_parameters(
chunk_size: Optional[Union[Vec3IntLike, int]],
chunks_per_shard: Optional[Union[Vec3IntLike, int]],
block_len: Optional[int],
file_len: Optional[int],
) -> Tuple[Optional[Vec3Int], Optional[Vec3Int]]:
if chunk_size is not None:
chunk_size = Vec3Int.from_vec_or_int(chunk_size)
elif block_len is not None:
warn_deprecated("block_len", "chunk_size")
chunk_size = Vec3Int.full(block_len)
if chunks_per_shard is not None:
chunks_per_shard = Vec3Int.from_vec_or_int(chunks_per_shard)
elif file_len is not None:
warn_deprecated("file_len", "chunks_per_shard")
chunks_per_shard = Vec3Int.full(file_len)
return (chunk_size, chunks_per_shard)
def create(
cls,
task_type_id: str,
project_name: str,
dataset_name: Union[str, RemoteDataset],
needed_experience_domain: str,
needed_experience_value: int,
starting_position: Vec3Int,
starting_rotation: Optional[Vec3Int] = Vec3Int(0, 0, 0),
instances: int = 1,
script_id: Optional[str] = None,
bounding_box: Optional[BoundingBox] = None,
) -> List["Task"]:
"""Submits tasks in webKnossos based on a dataset, starting position + rotation, and returns the Task objects"""
client = _get_generated_client(enforce_auth=True)
url = f"{client.base_url}/api/tasks"
if isinstance(dataset_name, RemoteDataset):
dataset_name = dataset_name._dataset_name
task_parameters = {
"taskTypeId":
task_type_id,
"neededExperience": {
"domain": needed_experience_domain,
"value": needed_experience_value,
},
"openInstances":
instances,
"projectName":
project_name,
"scriptId":
script_id,
"dataSet":
dataset_name,
"editPosition":
starting_position,
"editRotation":
starting_rotation,
"boundingBox":
bounding_box.to_wkw_dict() if bounding_box is not None else None,
}
response = httpx.post(
url=url,
headers=client.get_headers(),
cookies=client.get_cookies(),
timeout=client.get_timeout(),
json=[task_parameters],
)
assert (
response.status_code == 200
), f"Failed to create tasks: {response.status_code}: {response.text}"
return cls._handle_task_creation_response(response)
def __init__(
self,
view: "View",
offset: Optional[Vec3IntLike] = None,
# buffer_size specifies, how many slices should be aggregated until they are flushed.
buffer_size: int = 32,
dimension: int = 2, # z
*,
relative_offset: Optional[Vec3IntLike] = None, # in mag1
absolute_offset: Optional[Vec3IntLike] = None, # in mag1
use_logging: bool = False,
) -> None:
"""see `View.get_buffered_slice_writer()`"""
self.view = view
self.buffer_size = buffer_size
self.dtype = self.view.get_dtype()
self.use_logging = use_logging
if offset is None and relative_offset is None and absolute_offset is None:
relative_offset = Vec3Int.zeros()
if offset is not None:
warnings.warn(
"[DEPRECATION] Using offset for a buffered slice writer is deprecated. "
+ "Please use the parameter relative_offset or absolute_offset in Mag(1) instead.",
DeprecationWarning,
)
self.offset = None if offset is None else Vec3Int(offset)
self.relative_offset = (
None if relative_offset is None else Vec3Int(relative_offset)
)
self.absolute_offset = (
None if absolute_offset is None else Vec3Int(absolute_offset)
)
self.dimension = dimension
assert 0 <= dimension <= 2
self.buffer: List[np.ndarray] = []
self.current_slice: Optional[int] = None
self.buffer_start_slice: Optional[int] = None
def test_annotation_from_zip_file() -> None:
annotation = wk.Annotation.load(
TESTDATA_DIR / "annotations" /
"l4dense_motta_et_al_demo_v2__explorational__4a6356.zip")
assert annotation.dataset_name == "l4dense_motta_et_al_demo_v2"
assert annotation.organization_id == "scalable_minds"
assert annotation.owner_name == "Philipp Otto"
assert annotation.annotation_id == "61c20205010000cc004a6356"
assert ("timestamp" in annotation.metadata # pylint: disable=unsupported-membership-test
)
assert len(list(annotation.get_volume_layer_names())) == 1
assert len(list(annotation.skeleton.flattened_trees())) == 1
annotation.save(TESTOUTPUT_DIR / "test_dummy.zip")
copied_annotation = wk.Annotation.load(TESTOUTPUT_DIR / "test_dummy.zip")
assert copied_annotation.dataset_name == "l4dense_motta_et_al_demo_v2"
assert copied_annotation.organization_id == "scalable_minds"
assert copied_annotation.owner_name == "Philipp Otto"
assert copied_annotation.annotation_id == "61c20205010000cc004a6356"
assert ("timestamp" in copied_annotation.metadata # pylint: disable=unsupported-membership-test
)
assert len(list(copied_annotation.get_volume_layer_names())) == 1
assert len(list(copied_annotation.skeleton.flattened_trees())) == 1
copied_annotation.add_volume_layer(name="new_volume_layer")
assert len(list(copied_annotation.get_volume_layer_names())) == 2
copied_annotation.delete_volume_layer(volume_layer_name="new_volume_layer")
assert len(list(copied_annotation.get_volume_layer_names())) == 1
with annotation.temporary_volume_layer_copy() as volume_layer:
input_annotation_mag = volume_layer.get_finest_mag()
voxel_id = input_annotation_mag.read(absolute_offset=Vec3Int(
2830, 4356, 1792),
size=Vec3Int.full(1))
assert voxel_id == 2504698
def main(args: argparse.Namespace) -> None:
# Use the skeleton API to read the bounding boxes once
# https://github.com/scalableminds/webknossos-libs/issues/482 is done.
nml_regex = re.compile(
r'<userBoundingBox .*name="Limits of flood-fill \(source_id=(\d+), target_id=(\d+), seed=([\d,]+), timestamp=(\d+)\)".*topLeftX="(\d+)" topLeftY="(\d+)" topLeftZ="(\d+)" width="(\d+)" height="(\d+)" depth="(\d+)" />'
)
bboxes: List[FloodFillBbox] = []
nml_file = open(args.nml_path, "r", encoding="utf-8")
lines = nml_file.readlines()
nml_file.close()
for line in lines:
matches = nml_regex.findall(line)
for match in matches:
# each match is a tuple of (source_id, target_id, seed, timestamp, top_left_x, top_left_y, top_left_z, width, height, depth
bboxes.append(
FloodFillBbox(
bounding_box=BoundingBox((match[4], match[5], match[6]),
(match[7], match[8], match[9])),
seed_position=Vec3Int(match[2].split(",")),
source_id=int(match[0]),
target_id=int(match[1]),
timestamp=int(match[3]),
))
bboxes = sorted(bboxes, key=lambda x: x.timestamp)
time_start("Merge with fallback layer")
data_mag = merge_with_fallback_layer(
args.output_path,
args.volume_path,
args.segmentation_layer_path,
)
time_stop("Merge with fallback layer")
time_start("All floodfills")
for floodfill in bboxes:
time_start("Floodfill")
execute_floodfill(
data_mag,
floodfill.seed_position,
floodfill.bounding_box,
floodfill.source_id,
floodfill.target_id,
)
time_stop("Floodfill")
time_stop("All floodfills")
time_start("Recompute downsampled mags")
data_mag.layer.redownsample()
time_stop("Recompute downsampled mags")
def add_node( # pylint: disable=arguments-differ arguments-renamed
self,
position: Vec3IntLike,
comment: Optional[str] = None,
radius: Optional[float] = None,
rotation: Optional[Vector3] = None,
inVp: Optional[int] = None,
inMag: Optional[int] = None,
bitDepth: Optional[int] = None,
interpolation: Optional[bool] = None,
time: Optional[int] = None,
is_branchpoint: bool = False,
branchpoint_time: Optional[int] = None,
_enforced_id: Optional[int] = None,
) -> Node:
"""
Adds a node to the tree. Apart from the mandatory `position` parameter,
there are several optional parameters which can be used to encode
additional information. For example, the comment will be shown by the
webKnossos UI.
"""
node = Node(
position=Vec3Int(position),
comment=comment,
radius=radius,
rotation=rotation,
inVp=inVp,
inMag=inMag,
bitDepth=bitDepth,
interpolation=interpolation,
time=time,
is_branchpoint=is_branchpoint,
branchpoint_time=branchpoint_time,
enforced_id=_enforced_id,
skeleton=self._skeleton,
)
super().add_node(node)
return node
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
# (e.g. Layer.default_view_configuration has many attributes which are all optionally).
for cls in [
DatasetProperties,
MagViewProperties,
DatasetViewConfiguration,
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_repr() -> None:
assert str(Vec3Int(1, 2, 3)) == "Vec3Int(1,2,3)"
import webknossos as wk
from webknossos.dataset import Layer, MagView
from webknossos.geometry import BoundingBox, Mag, Vec3Int
from webknossos.utils import (
add_verbose_flag,
setup_logging,
setup_warnings,
time_start,
time_stop,
)
logger = logging.getLogger(__name__)
NEIGHBORS = [
Vec3Int(1, 0, 0),
Vec3Int(-1, 0, 0),
Vec3Int(0, 1, 0),
Vec3Int(0, -1, 0),
Vec3Int(0, 0, 1),
Vec3Int(0, 0, -1),
]
FloodFillBbox = namedtuple(
"FloodFillBbox",
["bounding_box", "seed_position", "source_id", "target_id", "timestamp"],
)
def create_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
def execute_floodfill(
data_mag: MagView,
seed_position: Vec3Int,
already_processed_bbox: BoundingBox,
source_id: int,
target_id: int,
) -> None:
cube_size = data_mag.info.shard_size
cube_bbox = BoundingBox(Vec3Int(0, 0, 0), cube_size)
chunk_with_relative_seed: List[Tuple[Vec3Int, Vec3Int]] = [
get_chunk_pos_and_offset(seed_position, cube_size)
]
# The `is_visited` variable is used to know which parts of the already processed bbox
# were already traversed. Outside of that bounding box, the actual data already
# is an indicator of whether the flood-fill has reached a voxel.
is_visited = np.zeros(already_processed_bbox.size.to_tuple(),
dtype=np.uint8)
chunk_count = 0
while len(chunk_with_relative_seed) > 0:
chunk_count += 1
if chunk_count % 10000 == 0:
logger.info(f"Handled seed positions {chunk_count}")
dirty_bucket = False
current_cube, relative_seed = chunk_with_relative_seed.pop()
global_seed = current_cube + relative_seed
# Only reading one voxel for the seed can be up to 30,000 times faster
# which is very relevent, since the chunk doesn't need to be traversed
# if the seed voxel was already covered.
value_at_seed_position = data_mag.read(current_cube + relative_seed,
(1, 1, 1))
if value_at_seed_position == source_id or (
already_processed_bbox.contains(global_seed)
and value_at_seed_position == target_id and
not is_visited[global_seed - already_processed_bbox.topleft]):
logger.info(
f"Handling chunk {chunk_count} with current cube {current_cube}"
)
time_start("read data")
cube_data = data_mag.read(current_cube, cube_size)
cube_data = cube_data[0, :, :, :]
time_stop("read data")
seeds_in_current_chunk: Set[Vec3Int] = set()
seeds_in_current_chunk.add(relative_seed)
time_start("traverse cube")
while len(seeds_in_current_chunk) > 0:
current_relative_seed = seeds_in_current_chunk.pop()
current_global_seed = current_cube + current_relative_seed
if already_processed_bbox.contains(current_global_seed):
is_visited[current_global_seed -
already_processed_bbox.topleft] = 1
if cube_data[current_relative_seed] != target_id:
cube_data[current_relative_seed] = target_id
dirty_bucket = True
# check neighbors
for neighbor in NEIGHBORS:
neighbor_pos = current_relative_seed + neighbor
global_neighbor_pos = current_cube + neighbor_pos
if already_processed_bbox.contains(global_neighbor_pos):
if is_visited[global_neighbor_pos -
already_processed_bbox.topleft]:
continue
if cube_bbox.contains(neighbor_pos):
if cube_data[neighbor_pos] == source_id or (
already_processed_bbox.contains(
global_neighbor_pos)
and cube_data[neighbor_pos] == target_id):
seeds_in_current_chunk.add(neighbor_pos)
else:
chunk_with_relative_seed.append(
get_chunk_pos_and_offset(global_neighbor_pos,
cube_size))
time_stop("traverse cube")
if dirty_bucket:
time_start("write chunk")
data_mag.write(cube_data, current_cube)
time_stop("write chunk")
def _vec3int_mulf(vec: Vec3Int, factors: List[float]) -> Vec3Int:
return Vec3Int(int(vec.x * factors[0]), int(vec.y * factors[1]),
int(vec.z * factors[2]))
def test_prod() -> None:
assert Vec3Int(1, 2, 3).prod() == 6
def test_contains() -> None:
assert Vec3Int(1, 2, 3).contains(1)
assert not Vec3Int(1, 2, 3).contains(4)
from webknossos.geometry import Vec3Int DEFAULT_WKW_FILE_LEN = 32 DEFAULT_CHUNK_SIZE = Vec3Int.full(32) DEFAULT_CHUNKS_PER_SHARD = Vec3Int.full(32) DEFAULT_CHUNKS_PER_SHARD_ZARR = Vec3Int.full(1)
def test_custom_initialization() -> None:
assert Vec3Int.zeros() == Vec3Int(0, 0, 0)
assert Vec3Int.ones() == Vec3Int(1, 1, 1)
assert Vec3Int.full(4) == Vec3Int(4, 4, 4)
assert Vec3Int.ones() - Vec3Int.ones() == Vec3Int.zeros()
assert Vec3Int.full(4) == Vec3Int.ones() * 4
def _write_buffer(self) -> None:
if len(self.buffer) == 0:
return
assert (
len(self.buffer) <= self.buffer_size
), "The WKW buffer is larger than the defined batch_size. The buffer should have been flushed earlier. This is probably a bug in the BufferedSliceWriter."
uniq_dtypes = set(map(lambda _slice: _slice.dtype, self.buffer))
assert (
len(uniq_dtypes) == 1
), "The buffer of BufferedSliceWriter contains slices with differing dtype."
assert uniq_dtypes.pop() == self.dtype, (
"The buffer of BufferedSliceWriter contains slices with a dtype "
"which differs from the dtype with which the BufferedSliceWriter was instantiated."
)
if self.use_logging:
info(
"({}) Writing {} slices at position {}.".format(
getpid(), len(self.buffer), self.buffer_start_slice
)
)
log_memory_consumption()
try:
assert (
self.buffer_start_slice is not None
), "Failed to write buffer: The buffer_start_slice is not set."
max_width = max(slice.shape[-2] for slice in self.buffer)
max_height = max(slice.shape[-1] for slice in self.buffer)
self.buffer = [
np.pad(
slice,
mode="constant",
pad_width=[
(0, 0),
(0, max_width - slice.shape[-2]),
(0, max_height - slice.shape[-1]),
],
)
for slice in self.buffer
]
data = np.concatenate(
[np.expand_dims(slice, self.dimension + 1) for slice in self.buffer],
axis=self.dimension + 1,
)
buffer_start_list = [0, 0, 0]
buffer_start_list[self.dimension] = self.buffer_start_slice
buffer_start = Vec3Int(buffer_start_list)
buffer_start_mag1 = buffer_start * self.view.mag.to_vec3_int()
self.view.write(
data,
offset=buffer_start.add_or_none(self.offset),
relative_offset=buffer_start_mag1.add_or_none(self.relative_offset),
absolute_offset=buffer_start_mag1.add_or_none(self.absolute_offset),
)
except Exception as exc:
error(
"({}) An exception occurred in BufferedSliceWriter._write_buffer with {} "
"slices at position {}. Original error is:\n{}:{}\n\nTraceback:".format(
getpid(),
len(self.buffer),
self.buffer_start_slice,
type(exc).__name__,
exc,
)
)
traceback.print_tb(exc.__traceback__)
error("\n")
raise exc
finally:
self.buffer = []
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
import numpy as np
from rich.progress import track
from webknossos.client._generated.api.datastore import dataset_download
from webknossos.client._generated.api.default import dataset_info
from webknossos.client._generated.types import Unset
from webknossos.client.context import _get_context
from webknossos.dataset import Dataset, LayerCategoryType
from webknossos.dataset.properties import LayerViewConfiguration, dataset_converter
from webknossos.geometry import BoundingBox, Mag, Vec3Int
logger = logging.getLogger(__name__)
T = TypeVar("T")
_DOWNLOAD_CHUNK_SIZE = Vec3Int(512, 512, 512)
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
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 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}",
)
from ._array import ArrayInfo
from .layer_categories import LayerCategoryType
from .view import View
class InterpolationModes(Enum):
MEDIAN = 0
MODE = 1
NEAREST = 2
BILINEAR = 3
BICUBIC = 4
MAX = 5
MIN = 6
DEFAULT_BUFFER_SHAPE = Vec3Int.full(256)
def determine_buffer_shape(array_info: ArrayInfo) -> Vec3Int:
return min(DEFAULT_BUFFER_SHAPE, array_info.shard_size)
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
