Update rechunk task & add tests

Author: jluethi

src/fractal_helper_tasks/drop_t_dimension.py

@@ -53,7 +53,8 @@ def drop_t_dimension(
             (standard argument for Fractal tasks, managed by Fractal server).
         suffix: Suffix to be used for the new Zarr image. If overwrite_input
             is True, this file is only temporary.
-        overwrite_input: Whether
+        overwrite_input: Whether the existing iamge should be overwritten with
+            the new OME-Zarr without the T dimension.
     """
     # Normalize zarr_url
     zarr_url_old = zarr_url.rstrip("/")

src/fractal_helper_tasks/rechunk_zarr.py

@@ -1,4 +1,4 @@
-# Copyright 2024 (C) BioVisionCenter, University of Zurich
+# Copyright 2025 (C) BioVisionCenter, University of Zurich
 #
 # Original authors:
 # Joel Lüthi <[email protected]>
@@ -10,9 +10,10 @@
 from typing import Any, Optional
 
 import ngio
-from ngio.core.utils import create_empty_ome_zarr_label
 from pydantic import validate_call
 
+from fractal_helper_tasks.utils import normalize_chunk_size_dict, rechunk_label
+
 logger = logging.getLogger(__name__)
 
 
@@ -22,6 +23,7 @@ def rechunk_zarr(
     zarr_url: str,
     chunk_sizes: Optional[dict[str, Optional[int]]] = None,
     suffix: str = "rechunked",
+    rechunk_labels: bool = True,
     rebuild_pyramids: bool = True,
     overwrite_input: bool = True,
     overwrite: bool = False,
@@ -38,6 +40,8 @@ def rechunk_zarr(
             sizes. {"z": 10} will just change the Z chunking while keeping
             all other chunk sizes the same as the input.
         suffix: Suffix of the rechunked image.
+        rechunk_labels: Whether to apply the same rechunking to all label
+            images of the OME-Zarr as well.
         rebuild_pyramids: Whether pyramids are built fresh in the rechunked
             image. This has a small performance overhead, but ensures that
             this task is save against off-by-one issues when pyramid levels
@@ -47,11 +51,9 @@ def rechunk_zarr(
         overwrite: Whether to overwrite potential pre-existing output with the
             name zarr_url_suffix.
     """
-    chunk_sizes = chunk_sizes or {}
-    valid_axes = ["t", "c", "z", "y", "x"]
-    for axis in valid_axes:
-        if axis not in chunk_sizes:
-            chunk_sizes[axis] = None
+    logger.info(f"Running `rechunk_zarr` on {zarr_url=} with {chunk_sizes=}.")
+
+    chunk_sizes = normalize_chunk_size_dict(chunk_sizes)
 
     rechunked_zarr_url = zarr_url + f"_{suffix}"
     ngff_image = ngio.NgffImage(zarr_url)
@@ -62,19 +64,27 @@ def rechunk_zarr(
 
     # Compute the chunksize tuple
     new_chunksize = [c[0] for c in chunks]
+    logger.info(f"Initial chunk sizes were: {chunks}")
     # Overwrite chunk_size with user-set chunksize
     for i, axis in enumerate(axes_names):
         if axis in chunk_sizes:
             if chunk_sizes[axis] is not None:
                 new_chunksize[i] = chunk_sizes[axis]
 
-    # TODO: Check for extra axes specified
+    for axis in chunk_sizes:
+        if axis not in axes_names:
+            raise NotImplementedError(
+                f"Rechunking with {axis=} is specified, but the OME-Zarr only "
+                f"has the following axes: {axes_names}"
+            )
+
+    logger.info(f"Chunk sizes after rechunking will be: {new_chunksize=}")
 
     new_ngff_image = ngff_image.derive_new_image(
         store=rechunked_zarr_url,
         name=ngff_image.image_meta.name,
         overwrite=overwrite,
-        copy_labels=False,  # Copy if rechunk labels is not selected?
+        copy_labels=not rechunk_labels,
         copy_tables=True,
         chunks=new_chunksize,
     )
@@ -93,79 +103,37 @@ def rechunk_zarr(
                 ngff_image.get_image(path=path).on_disk_dask_array
             )
 
-    # Copy labels: Loop over them
-    # Labels don't have a channel dimension
-    chunk_sizes["c"] = None
-    label_names = ngff_image.labels.list()
-    for label in label_names:
-        old_label = ngff_image.labels.get_label(name=label)
-        label_level_paths = ngff_image.labels.levels_paths(name=label)
-        # Compute the chunksize tuple
-        chunks = old_label.on_disk_dask_array.chunks
-        new_chunksize = [c[0] for c in chunks]
-        # Overwrite chunk_size with user-set chunksize
-        for i, axis in enumerate(old_label.dataset.on_disk_axes_names):
-            if axis in chunk_sizes:
-                if chunk_sizes[axis] is not None:
-                    new_chunksize[i] = chunk_sizes[axis]
-        create_empty_ome_zarr_label(
-            store=new_ngff_image.store
-            + "/"
-            + "labels"
-            + "/"
-            + label,  # FIXME: Set this better?
-            on_disk_shape=old_label.on_disk_shape,
-            chunks=new_chunksize,
-            dtype=old_label.on_disk_dask_array.dtype,
-            on_disk_axis=old_label.dataset.on_disk_axes_names,
-            pixel_sizes=old_label.dataset.pixel_size,
-            xy_scaling_factor=old_label.metadata.xy_scaling_factor,
-            z_scaling_factor=old_label.metadata.z_scaling_factor,
-            time_spacing=old_label.dataset.time_spacing,
-            time_units=old_label.dataset.time_axis_unit,
-            levels=label_level_paths,
-            name=label,
-            overwrite=overwrite,
-            version=old_label.metadata.version,
-        )
-
-        # Fill in labels .attrs to contain the label name
-        list_of_labels = new_ngff_image.labels.list()
-        if label not in list_of_labels:
-            new_ngff_image.labels._label_group.attrs["labels"] = [
-                *list_of_labels,
-                label,
-            ]
-
-        if rebuild_pyramids:
-            # Set the highest resolution, then consolidate to build a new pyramid
-            new_ngff_image.labels.get_label(
-                name=label, highest_resolution=True
-            ).set_array(
-                ngff_image.labels.get_label(
-                    name=label, highest_resolution=True
-                ).on_disk_dask_array
+    # Copy labels
+    if rechunk_labels:
+        chunk_sizes["c"] = None
+        label_names = ngff_image.labels.list()
+        for label in label_names:
+            rechunk_label(
+                orig_ngff_image=ngff_image,
+                new_ngff_image=new_ngff_image,
+                label=label,
+                chunk_sizes=chunk_sizes,
+                overwrite=overwrite,
+                rebuild_pyramids=rebuild_pyramids,
             )
-            new_ngff_image.labels.get_label(
-                name=label, highest_resolution=True
-            ).consolidate()
-        else:
-            for label_path in label_level_paths:
-                new_ngff_image.labels.get_label(name=label, path=label_path).set_array(
-                    ngff_image.labels.get_label(
-                        name=label, path=label_path
-                    ).on_disk_dask_array
-                )
+
     if overwrite_input:
         os.rename(zarr_url, f"{zarr_url}_tmp")
         os.rename(rechunked_zarr_url, zarr_url)
         shutil.rmtree(f"{zarr_url}_tmp")
         return
     else:
-        image_list_updates = dict(
-            image_list_updates=[dict(zarr_url=rechunked_zarr_url, origin=zarr_url)]
+        output = dict(
+            image_list_updates=[
+                dict(
+                    zarr_url=rechunked_zarr_url,
+                    origin=zarr_url,
+                    types=dict(rechunked=True),
+                )
+            ],
+            filters=dict(types=dict(rechunked=True)),
         )
-        return image_list_updates
+        return output
 
 
 if __name__ == "__main__":

src/fractal_helper_tasks/utils.py

@@ -0,0 +1,118 @@
+# Copyright 2025 (C) BioVisionCenter, University of Zurich
+#
+# Original authors:
+# Joel Lüthi <[email protected]>
+"""Utils for helper tasks."""
+
+from typing import Optional
+
+import ngio
+from ngio.core.utils import create_empty_ome_zarr_label
+
+
+def normalize_chunk_size_dict(chunk_sizes: dict[str, Optional[int]]):
+    """Converts all chunk_size axes names to lower case and assert validity.
+
+    Args:
+        chunk_sizes: Dictionary of chunk sizes that should be adapted. Can
+            contain new chunk sizes for t, c, z, y & x.
+
+    Returns:
+        chunk_sizes_norm: Normalized chunk_sizes dict.
+    """
+    chunk_sizes = chunk_sizes or {}
+    chunk_sizes_norm = {}
+    for key, value in chunk_sizes.items():
+        chunk_sizes_norm[key.lower()] = value
+
+    valid_axes = ["t", "c", "z", "y", "x"]
+    for axis in chunk_sizes_norm:
+        if axis not in valid_axes:
+            raise ValueError(
+                f"Axis {axis} is not supported. Valid axes choices are "
+                f"{valid_axes}."
+            )
+    return chunk_sizes_norm
+
+
+def rechunk_label(
+    orig_ngff_image: ngio.NgffImage,
+    new_ngff_image: ngio.NgffImage,
+    label: str,
+    chunk_sizes: list[int],
+    overwrite: bool = False,
+    rebuild_pyramids: bool = True,
+):
+    """Saves a rechunked label image into a new OME-Zarr
+
+    The label image is based on an existing label image in another OME-Zarr.
+
+    Args:
+        orig_ngff_image: Original OME-Zarr that contains the label image
+        new_ngff_image: OME-Zarr to which the rechunked label image should be
+            added.
+        label: Name of the label image.
+        chunk_sizes: New chunk sizes that should be applied
+        overwrite: Whether the label image in `new_ngff_image` should be
+            overwritten if it already exists.
+        rebuild_pyramids: Whether pyramids are built fresh in the rechunked
+            label image. This has a small performance overhead, but ensures
+            that this task is save against off-by-one issues when pyramid
+            levels aren't easily downsampled by 2.
+    """
+    old_label = orig_ngff_image.labels.get_label(name=label)
+    label_level_paths = orig_ngff_image.labels.levels_paths(name=label)
+    # Compute the chunksize tuple
+    chunks = old_label.on_disk_dask_array.chunks
+    new_chunksize = [c[0] for c in chunks]
+    # Overwrite chunk_size with user-set chunksize
+    for i, axis in enumerate(old_label.dataset.on_disk_axes_names):
+        if axis in chunk_sizes:
+            if chunk_sizes[axis] is not None:
+                new_chunksize[i] = chunk_sizes[axis]
+    create_empty_ome_zarr_label(
+        store=new_ngff_image.store
+        + "/"
+        + "labels"
+        + "/"
+        + label,  # FIXME: Set this better?
+        on_disk_shape=old_label.on_disk_shape,
+        chunks=new_chunksize,
+        dtype=old_label.on_disk_dask_array.dtype,
+        on_disk_axis=old_label.dataset.on_disk_axes_names,
+        pixel_sizes=old_label.dataset.pixel_size,
+        xy_scaling_factor=old_label.metadata.xy_scaling_factor,
+        z_scaling_factor=old_label.metadata.z_scaling_factor,
+        time_spacing=old_label.dataset.time_spacing,
+        time_units=old_label.dataset.time_axis_unit,
+        levels=label_level_paths,
+        name=label,
+        overwrite=overwrite,
+        version=old_label.metadata.version,
+    )
+
+    # Fill in labels .attrs to contain the label name
+    list_of_labels = new_ngff_image.labels.list()
+    if label not in list_of_labels:
+        new_ngff_image.labels._label_group.attrs["labels"] = [
+            *list_of_labels,
+            label,
+        ]
+
+    if rebuild_pyramids:
+        # Set the highest resolution, then consolidate to build a new pyramid
+        new_ngff_image.labels.get_label(name=label, highest_resolution=True).set_array(
+            orig_ngff_image.labels.get_label(
+                name=label, highest_resolution=True
+            ).on_disk_dask_array
+        )
+        new_ngff_image.labels.get_label(
+            name=label, highest_resolution=True
+        ).consolidate()
+    else:
+        for label_path in label_level_paths:
+            new_ngff_image.labels.get_label(name=label, path=label_path).set_array(
+                orig_ngff_image.labels.get_label(
+                    name=label, path=label_path
+                ).on_disk_dask_array
+            )