Refactor drop t dimension task to ngio 0.2.2

Author: jluethi

src/fractal_helper_tasks/drop_t_dimension.py

@@ -8,37 +8,17 @@
 """Task to remove singleton T dimension from an OME-Zarr."""
 
 import logging
+import os
+import shutil
 from typing import Any
 
 import dask.array as da
-import zarr
-from fractal_tasks_core.ngff import load_NgffImageMeta
-from fractal_tasks_core.pyramids import build_pyramid
+import ngio
 from pydantic import validate_call
 
 logger = logging.getLogger(__name__)
 
 
-def get_attrs_without_t(zarr_url: str):
-    """Generate zattrs without the t dimension.
-
-    Args:
-        zarr_url: Path to the zarr image
-    """
-    image_group = zarr.open_group(zarr_url)
-    zattrs = image_group.attrs.asdict()
-    # print(zattrs)
-    for multiscale in zattrs["multiscales"]:
-        # Update axes
-        multiscale["axes"] = multiscale["axes"][1:]
-        # Update coordinate Transforms
-        for dataset in multiscale["datasets"]:
-            for transform in dataset["coordinateTransformations"]:
-                if transform["type"] == "scale":
-                    transform["scale"] = transform["scale"][1:]
-    return zattrs
-
-
 @validate_call
 def drop_t_dimension(
     *,
@@ -58,60 +38,49 @@ def drop_t_dimension(
     """
     # Normalize zarr_url
     zarr_url_old = zarr_url.rstrip("/")
-    if overwrite_input:
-        zarr_url_new = zarr_url_old
-    else:
-        zarr_url_new = f"{zarr_url_old}_{suffix}"
+    zarr_url_new = f"{zarr_url_old}_{suffix}"
 
     logger.info(f"{zarr_url_old=}")
     logger.info(f"{zarr_url_new=}")
 
-    # Read some parameters from metadata
-    ngff_image = load_NgffImageMeta(zarr_url_old)
-
-    # Check that T axis is the first axis:
-    if not ngff_image.multiscale.axes[0].name == "t":
-        logger.warning(
-            f"The Zarr image {zarr_url_old} did not contain a T axis as its "
-            f"first axis. The axes were: {ngff_image.multiscale.axes} \n"
-            "The Drop T axis task is skipped"
+    old_ome_zarr = ngio.open_ome_zarr_container(zarr_url_old)
+    old_ome_zarr_img = old_ome_zarr.get_image()
+    if not old_ome_zarr_img.has_axis("t"):
+        raise ValueError(
+            f"The Zarr image {zarr_url_old} does not contain a T axis. "
+            "Thus, the drop T dimension task can't be applied to it."
         )
-        return {}
-
-    # Load 0-th level
-    data_tczyx = da.from_zarr(zarr_url_old + "/0")
-    # TODO: Check that T dimension is actually a singleton.
-    new_data = data_tczyx[0, ...]
+    # TODO: Check if T dimension not singleton
+    image = old_ome_zarr_img.get_array(mode="dask")
+    t_index = old_ome_zarr_img.meta.axes_mapper.get_index("t")
+    new_img = da.squeeze(image, axis=t_index)
+    pixel_size = old_ome_zarr_img.pixel_size
+    new_pixel_size = ngio.PixelSize(x=pixel_size.x, y=pixel_size.y, z=pixel_size.z)
+    axes_names = old_ome_zarr_img.meta.axes_mapper.on_disk_axes_names
+    del axes_names[t_index]
+    chunk_sizes = old_ome_zarr_img.chunks
+    new_chunk_sizes = chunk_sizes[:t_index] + chunk_sizes[t_index + 1 :]
+    new_ome_zarr_container = old_ome_zarr.derive_image(
+        store=zarr_url_new,
+        shape=new_img.shape,
+        chunks=new_chunk_sizes,
+        dtype=old_ome_zarr_img.dtype,
+        pixel_size=new_pixel_size,
+        axes_names=axes_names,
+    )
+    new_image_container = new_ome_zarr_container.get_image()
+    new_image_container.set_array(new_img)
+    new_image_container.consolidate()
 
     if overwrite_input:
         image_list_update = dict(zarr_url=zarr_url_old, types=dict(has_t=False))
+        os.rename(zarr_url_old, f"{zarr_url_old}_tmp")
+        os.rename(zarr_url_new, zarr_url_old)
+        shutil.rmtree(f"{zarr_url}_tmp")
     else:
-        # Generate attrs without the T dimension
-        new_attrs = get_attrs_without_t(zarr_url_old)
-        new_image_group = zarr.group(zarr_url_new)
-        new_image_group.attrs.put(new_attrs)
         image_list_update = dict(
             zarr_url=zarr_url_new, origin=zarr_url_old, types=dict(has_t=False)
         )
-        # TODO: Check if image contains labels & raise error (or even copy them)
-        # FIXME: Check if image contains ROI tables & copy them
-
-    # Write to disk (triggering execution)
-    logger.debug(f"Writing Zarr without T dimension to {zarr_url_new}")
-    new_data.to_zarr(
-        f"{zarr_url_new}/0",
-        overwrite=True,
-        dimension_separator="/",
-        write_empty_chunks=False,
-    )
-    logger.debug(f"Finished writing Zarr without T dimension to {zarr_url_new}")
-    build_pyramid(
-        zarrurl=zarr_url_new,
-        overwrite=True,
-        num_levels=ngff_image.num_levels,
-        coarsening_xy=ngff_image.coarsening_xy,
-        chunksize=new_data.chunksize,
-    )
 
     return {"image_list_updates": [image_list_update]}
 

tests/test_drop_t_dimension.py

@@ -0,0 +1,38 @@
+"""Test drop t dimension task."""
+
+from pathlib import Path
+
+import ngio
+import numpy as np
+
+from fractal_helper_tasks.drop_t_dimension import (
+    drop_t_dimension,
+)
+
+
+def test_drop_t_dimension(
+    tmp_path: Path,
+):
+    zarr_url = str(tmp_path / "my_zarr.zarr")
+
+    ngio.create_ome_zarr_from_array(
+        store=zarr_url,
+        array=np.zeros((1, 1, 1, 100, 100)),
+        xy_pixelsize=0.5,
+        z_spacing=1.0,
+        axes_names="tczyx",
+        overwrite=True,
+    )
+
+    drop_t_dimension(
+        zarr_url=zarr_url,
+        overwrite_input=False,
+    )
+    new_zarr_url = f"{zarr_url}_no_T"
+    new_ome_zarr_container = ngio.open_ome_zarr_container(new_zarr_url)
+    assert new_ome_zarr_container.image_meta.axes_mapper.on_disk_axes_names == [
+        "c",
+        "z",
+        "y",
+        "x",
+    ]