Add a Add Z singleton dimension task and fix table copying in drop t task

Author: jluethi

src/fractal_helper_tasks/__FRACTAL_MANIFEST__.json

@@ -9,7 +9,7 @@
       "tags": [
         "Singleton time dimension"
       ],
-      "docs_info": "### Purpose\n- Removes a **singleton time (T) dimension** from an OME-Zarr image.  \n- Creates a new OME-Zarr image with updated metadata and dimensions.\n- Optionally overwrites the input image if `overwrite_input` is set to True.\n\n### Outputs\n- A **new Zarr image** without the singleton T-dimension, stored with a configurable suffix.  \n\n### Limitations\n- Only processes OME-Zarr images where the **T-axis is the first axis**.  \n- Assumes the T-dimension is **singleton**; does not process non-singleton time axes.  \n- Does not copy associated **label images** or **ROI tables** to the new Zarr structure.  ",
+      "docs_info": "### Purpose\n- Removes a **singleton time (T) dimension** from an OME-Zarr image.  \n- Creates a new OME-Zarr image with updated metadata and dimensions.\n- Optionally overwrites the input image if `overwrite_input` is set to True.\n\n### Outputs\n- A **new Zarr image** without the singleton T-dimension, stored with a configurable suffix.  \n\n### Limitations\n- Only processes OME-Zarr images where the **T-axis is the first axis**.  \n- Assumes the T-dimension is **singleton**; does not process non-singleton time axes.  \n- Does not copy associated **label images** to the new Zarr structure.  ",
       "type": "parallel",
       "executable_parallel": "drop_t_dimension.py",
       "meta_parallel": {
@@ -205,6 +205,50 @@
         "title": "RechunkZarr"
       },
       "docs_link": "https://github.com/jluethi/fractal-helper-tasks"
+    },
+    {
+      "name": "Add Z Singleton Dimension",
+      "input_types": {
+        "is_3D": false
+      },
+      "tags": [
+        "Singleton Z dimension"
+      ],
+      "docs_info": "### Purpose\n- Removes a **singleton time (T) dimension** from an OME-Zarr image.  \n- Creates a new OME-Zarr image with updated metadata and dimensions.\n- Optionally overwrites the input image if `overwrite_input` is set to True.\n\n### Outputs\n- A **new Zarr image** without the singleton T-dimension, stored with a configurable suffix.  \n\n### Limitations\n- Only processes OME-Zarr images where the **T-axis is the first axis**.  \n- Assumes the T-dimension is **singleton**; does not process non-singleton time axes.  \n- Does not copy associated **label images** to the new Zarr structure.  ",
+      "type": "parallel",
+      "executable_parallel": "add_z_singleton.py",
+      "meta_parallel": {
+        "cpus_per_task": 1,
+        "mem": 4000
+      },
+      "args_schema_parallel": {
+        "additionalProperties": false,
+        "properties": {
+          "zarr_url": {
+            "title": "Zarr Url",
+            "type": "string",
+            "description": "Path or url to the individual OME-Zarr image to be processed. (standard argument for Fractal tasks, managed by Fractal server)."
+          },
+          "suffix": {
+            "default": "z_singleton",
+            "title": "Suffix",
+            "type": "string",
+            "description": "Suffix to be used for the new Zarr image. If overwrite_input is True, this file is only temporary."
+          },
+          "overwrite_input": {
+            "default": true,
+            "title": "Overwrite Input",
+            "type": "boolean",
+            "description": "Whether the existing iamge should be overwritten with the new OME-Zarr with the Z singleton dimension."
+          }
+        },
+        "required": [
+          "zarr_url"
+        ],
+        "type": "object",
+        "title": "AddZSingleton"
+      },
+      "docs_link": "https://github.com/jluethi/fractal-helper-tasks"
     }
   ],
   "has_args_schemas": true,

src/fractal_helper_tasks/add_z_singleton.py

@@ -0,0 +1,107 @@
+# Copyright 2024 (C) BioVisionCenter, University of Zurich
+#
+# Original authors:
+# Joel Lüthi <[email protected]>
+#
+# This file is derived from a Fractal task core task developed by
+# Tommaso Comparin & Marco Franzon
+"""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 ngio
+from pydantic import validate_call
+
+logger = logging.getLogger(__name__)
+
+
+@validate_call
+def add_z_singleton(
+    *,
+    zarr_url: str,
+    suffix: str = "z_singleton",
+    overwrite_input: bool = True,
+) -> dict[str, Any]:
+    """Add a singleton Z dimension to a 2D OME-Zarr.
+
+    Args:
+        zarr_url: Path or url to the individual OME-Zarr image to be processed.
+            (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 the existing iamge should be overwritten with
+            the new OME-Zarr with the Z singleton dimension.
+    """
+    # Normalize zarr_url
+    zarr_url_old = zarr_url.rstrip("/")
+    zarr_url_new = f"{zarr_url_old}_{suffix}"
+
+    logger.info(f"{zarr_url_old=}")
+    logger.info(f"{zarr_url_new=}")
+
+    old_ome_zarr = ngio.open_ome_zarr_container(zarr_url_old)
+    old_ome_zarr_img = old_ome_zarr.get_image()
+    if old_ome_zarr_img.has_axis("z"):
+        raise ValueError(
+            f"The Zarr image {zarr_url_old} already contains a Z axis. "
+            "Thus, the add Z singleton dimension task can't be applied to it."
+        )
+    image = old_ome_zarr_img.get_array(mode="dask")
+    axes_names = old_ome_zarr_img.meta.axes_mapper.on_disk_axes_names
+    ndim = image.ndim
+    insert_index = ndim - 2
+    if insert_index < 0:
+        raise ValueError(
+            f"Cannot insert a Z axis at position {insert_index} in an array"
+            f" with {ndim} dimensions."
+        )
+    # Insert singleton Z dimension
+    image_with_z = da.expand_dims(image, axis=insert_index)
+    logger.info(f"Original shape: {image.shape}, new shape: {image_with_z.shape}")
+    axes_names_with_z = axes_names[:insert_index] + ["z"] + axes_names[insert_index:]
+
+    pixel_size = old_ome_zarr_img.pixel_size
+    new_pixel_size = ngio.PixelSize(x=pixel_size.x, y=pixel_size.y, z=1.0)
+
+    chunk_sizes = old_ome_zarr_img.chunks
+    new_chunk_sizes = chunk_sizes[:insert_index] + (1,) + chunk_sizes[insert_index:]
+
+    new_ome_zarr_container = old_ome_zarr.derive_image(
+        store=zarr_url_new,
+        shape=image_with_z.shape,
+        chunks=new_chunk_sizes,
+        dtype=old_ome_zarr_img.dtype,
+        pixel_size=new_pixel_size,
+        axes_names=axes_names_with_z,
+        copy_tables=True,
+    )
+    new_image_container = new_ome_zarr_container.get_image()
+    new_image_container.set_array(image_with_z)
+    new_image_container.consolidate()
+
+    # TODO: Also handle copying over & adding Z dimension to label images?
+
+    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:
+        image_list_update = dict(
+            zarr_url=zarr_url_new, origin=zarr_url_old, types=dict(has_t=False)
+        )
+
+    return {"image_list_updates": [image_list_update]}
+
+
+if __name__ == "__main__":
+    from fractal_task_tools.task_wrapper import run_fractal_task
+
+    run_fractal_task(
+        task_function=add_z_singleton,
+        logger_name=logger.name,
+    )

src/fractal_helper_tasks/dev/docs_info/add_z_singleton.md

@@ -0,0 +1,10 @@
+### Purpose
+- Creates a **singleton Z dimension** in a 2D OME-Zarr image. Useful when 2D images don't have a singleton Z dimension but downstream tasks require it.
+- Optionally overwrites the input image if `overwrite_input` is set to True.
+
+### Outputs
+- A **new Zarr image** with the singleton Z dimension 
+
+### Limitations
+- Only processes 2D OME-Zarr images without a **Z-axis**.  
+- Does not copy associated **label images** to the new Zarr structure.  

src/fractal_helper_tasks/dev/docs_info/drop_t_dimension.md

@@ -9,4 +9,4 @@
 ### Limitations
 - Only processes OME-Zarr images where the **T-axis is the first axis**.  
 - Assumes the T-dimension is **singleton**; does not process non-singleton time axes.  
-- Does not copy associated **label images** or **ROI tables** to the new Zarr structure.  
+- Does not copy associated **label images** to the new Zarr structure.  

src/fractal_helper_tasks/dev/task_list.py

@@ -35,4 +35,12 @@
         ],
         docs_info="file:docs_info/rechunk_zarr.md",
     ),
+    ParallelTask(
+        name="Add Z Singleton Dimension",
+        executable="add_z_singleton.py",
+        meta={"cpus_per_task": 1, "mem": 4000},
+        input_types=dict(is_3D=False),
+        tags=["Singleton Z dimension"],
+        docs_info="file:docs_info/drop_t_dimension.md",
+    ),
 ]

src/fractal_helper_tasks/drop_t_dimension.py

@@ -67,6 +67,7 @@ def drop_t_dimension(
         dtype=old_ome_zarr_img.dtype,
         pixel_size=new_pixel_size,
         axes_names=axes_names,
+        copy_tables=True,
     )
     new_image_container = new_ome_zarr_container.get_image()
     new_image_container.set_array(new_img)

tests/test_add_z_singleton.py

@@ -0,0 +1,67 @@
+"""Test drop t dimension task."""
+
+from pathlib import Path
+
+import ngio
+import numpy as np
+import pytest
+
+from fractal_helper_tasks.add_z_singleton import (
+    add_z_singleton,
+)
+
+
+@pytest.mark.parametrize(
+    "orig_axes_names, target_axes_names, orig_dimensions, "
+    "target_dimensions, overwrite_input",
+    [
+        ("cyx", ["c", "z", "y", "x"], (5, 100, 100), (5, 1, 100, 100), True),
+        ("cyx", ["c", "z", "y", "x"], (5, 100, 100), (5, 1, 100, 100), False),
+        ("yx", ["z", "y", "x"], (100, 100), (1, 100, 100), True),
+        (
+            "tcyx",
+            ["t", "c", "z", "y", "x"],
+            (3, 5, 100, 100),
+            (3, 5, 1, 100, 100),
+            True,
+        ),
+        ("tyx", ["t", "z", "y", "x"], (3, 100, 100), (3, 1, 100, 100), True),
+    ],
+)
+def test_add_singleton(
+    tmp_path: Path,
+    orig_axes_names: str,
+    target_axes_names: list[str],
+    orig_dimensions: tuple[int],
+    target_dimensions: tuple[int],
+    overwrite_input: bool,
+):
+    zarr_url = str(tmp_path / "my_zarr.zarr")
+
+    ngio.create_ome_zarr_from_array(
+        store=zarr_url,
+        array=np.zeros(orig_dimensions),
+        xy_pixelsize=0.5,
+        axes_names=orig_axes_names,
+        overwrite=True,
+    )
+
+    suffix = "z_singleton"
+    add_z_singleton(
+        zarr_url=zarr_url,
+        suffix=suffix,
+        overwrite_input=overwrite_input,
+    )
+
+    if overwrite_input:
+        new_zarr_url = zarr_url
+    else:
+        new_zarr_url = f"{zarr_url}_{suffix}"
+
+    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
+        == target_axes_names
+    )
+    assert new_ome_zarr_container.get_image().pixel_size.z == 1.0
+    assert new_ome_zarr_container.get_image().shape == target_dimensions