diff --git a/src/fractal_helper_tasks/__FRACTAL_MANIFEST__.json b/src/fractal_helper_tasks/__FRACTAL_MANIFEST__.json index 5567db5..8aae4db 100644 --- a/src/fractal_helper_tasks/__FRACTAL_MANIFEST__.json +++ b/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, diff --git a/src/fractal_helper_tasks/add_z_singleton.py b/src/fractal_helper_tasks/add_z_singleton.py new file mode 100644 index 0000000..d52579d --- /dev/null +++ b/src/fractal_helper_tasks/add_z_singleton.py @@ -0,0 +1,107 @@ +# Copyright 2024 (C) BioVisionCenter, University of Zurich +# +# Original authors: +# Joel Lüthi +# +# 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, + ) diff --git a/src/fractal_helper_tasks/dev/docs_info/add_z_singleton.md b/src/fractal_helper_tasks/dev/docs_info/add_z_singleton.md new file mode 100644 index 0000000..a8a825e --- /dev/null +++ b/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. \ No newline at end of file diff --git a/src/fractal_helper_tasks/dev/docs_info/drop_t_dimension.md b/src/fractal_helper_tasks/dev/docs_info/drop_t_dimension.md index d0e548e..6d61c8c 100644 --- a/src/fractal_helper_tasks/dev/docs_info/drop_t_dimension.md +++ b/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. \ No newline at end of file +- Does not copy associated **label images** to the new Zarr structure. \ No newline at end of file diff --git a/src/fractal_helper_tasks/dev/task_list.py b/src/fractal_helper_tasks/dev/task_list.py index 4147767..4b5b461 100644 --- a/src/fractal_helper_tasks/dev/task_list.py +++ b/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", + ), ] diff --git a/src/fractal_helper_tasks/drop_t_dimension.py b/src/fractal_helper_tasks/drop_t_dimension.py index 2062553..3cbc1e1 100644 --- a/src/fractal_helper_tasks/drop_t_dimension.py +++ b/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) diff --git a/tests/test_add_z_singleton.py b/tests/test_add_z_singleton.py new file mode 100644 index 0000000..bc462d4 --- /dev/null +++ b/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