Refactor to ngio 0.4

Author: jluethi

pyproject.toml

@@ -27,14 +27,15 @@ authors = [
 # Required Python version and dependencies
 requires-python = ">=3.11"
 dependencies = [
-    "ngio>=0.2.8, <0.3.0",
+    "ngio>=0.4.1, <0.5.0",
     "fractal-task-tools==0.0.12",
+    "pyarrow<21.0.0",
 ]
 
 # Optional dependencies (e.g. for `pip install -e ".[dev]"`, see
 # https://peps.python.org/pep-0621/#dependencies-optional-dependencies)
 [project.optional-dependencies]
-dev = ["devtools", "hatch", "pytest", "requests", "jsonschema", "ruff", "pre-commit", "pooch"]
+dev = ["devtools", "hatch", "pytest", "requests", "jsonschema", "ruff", "pre-commit", "pooch", "pillow"]
 test = ["pytest", "pytest-cov"]
 
 # https://docs.astral.sh/ruff

src/fractal_helper_tasks/add_z_singleton.py

@@ -51,7 +51,7 @@ def add_z_singleton(
             "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
+    axes_names = list(old_ome_zarr_img.meta.axes_handler.axes_names)
     ndim = image.ndim
     insert_index = ndim - 2
     if insert_index < 0:

src/fractal_helper_tasks/convert_2D_segmentation_to_3D.py

@@ -124,7 +124,7 @@ def convert_2D_segmentation_to_3D(
             raise ValueError(
                 f"Label image {label_name} is not 2D. It has a shape of "
                 f"{label_img.shape} and the axes "
-                f"{label_img.axes_mapper.on_disk_axes_names}."
+                f"{label_img.axes}."
             )
 
         chunks = list(label_img.chunks)
@@ -135,10 +135,10 @@ def convert_2D_segmentation_to_3D(
             pixel_size=label_img.pixel_size,
         )
 
-        z_index = label_img.axes_mapper.get_index("z")
-        y_index = label_img.axes_mapper.get_index("y")
-        x_index = label_img.axes_mapper.get_index("x")
-        z_index_3d_reference = ref_image_3d.axes_mapper.get_index("z")
+        z_index = label_img.axes_handler.get_index("z")
+        y_index = label_img.axes_handler.get_index("y")
+        x_index = label_img.axes_handler.get_index("x")
+        z_index_3d_reference = ref_image_3d.axes_handler.get_index("z")
         if z_chunks:
             chunks[z_index] = z_chunks
         else:
@@ -151,7 +151,7 @@ def convert_2D_segmentation_to_3D(
 
         pixel_size = label_img.pixel_size
         pixel_size.z = ref_image_3d.pixel_size.z
-        axes_names = label_img.axes_mapper.on_disk_axes_names
+        axes_names = label_img.meta.axes_handler.axes_names
 
         z_extent = nb_z_planes * pixel_size.z
 
@@ -191,8 +191,10 @@ def convert_2D_segmentation_to_3D(
                     f"Table {table_name} not found in 2D OME-Zarr {zarr_url}."
                 )
             table = ome_zarr_container_2d.get_table(table_name)
-            print(table.type())
-            if table.type() == "roi_table" or table.type() == "masking_roi_table":
+            if (
+                table.table_type() == "roi_table"
+                or table.table_type() == "masking_roi_table"
+            ):
                 for roi in table.rois():
                     roi.z_length = z_extent
                 ome_zarr_container_3d.add_table(

src/fractal_helper_tasks/drop_t_dimension.py

@@ -52,11 +52,11 @@ def drop_t_dimension(
         )
     # 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")
+    t_index = old_ome_zarr_img.axes_handler.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
+    axes_names = list(old_ome_zarr_img.axes)
     del axes_names[t_index]
     chunk_sizes = old_ome_zarr_img.chunks
     new_chunk_sizes = chunk_sizes[:t_index] + chunk_sizes[t_index + 1 :]

src/fractal_helper_tasks/rechunk_zarr.py

@@ -10,9 +10,6 @@
 from typing import Any, Optional
 
 import ngio
-import ngio.images
-import ngio.images.label
-from ngio.ome_zarr_meta import AxesMapper
 from pydantic import validate_call
 
 from fractal_helper_tasks.utils import normalize_chunk_size_dict
@@ -22,7 +19,7 @@
 
 def change_chunks(
     initial_chunks: list[int],
-    axes_mapper: AxesMapper,
+    multiscale: ngio.Image | ngio.Label,
     chunk_sizes: dict[str, Optional[int]],
 ) -> list[int]:
     """Create a new chunk_size list with rechunking.
@@ -33,12 +30,12 @@ def change_chunks(
     """
     for axes_name, chunk_value in chunk_sizes.items():
         if chunk_value is not None:
-            axes_index = axes_mapper.get_index(axes_name)
+            axes_index = multiscale.axes_handler.get_index(axes_name)
             if axes_index is None:
                 raise ValueError(
                     f"Rechunking with {axes_name=} is specified, but the "
                     "OME-Zarr only has the following axes: "
-                    f"{axes_mapper.on_disk_axes_names}"
+                    f"{multiscale.axes}"
                 )
             initial_chunks[axes_index] = chunk_value
     return initial_chunks
@@ -89,7 +86,7 @@ def rechunk_zarr(
     chunks = highest_res_img.chunks
     new_chunksize = change_chunks(
         initial_chunks=list(chunks),
-        axes_mapper=highest_res_img.meta.axes_mapper,
+        multiscale=highest_res_img,
         chunk_sizes=chunk_sizes,
     )
 
@@ -123,7 +120,7 @@ def rechunk_zarr(
             old_label = ome_zarr_container.get_label(name=label)
             new_chunksize = change_chunks(
                 initial_chunks=list(old_label.chunks),
-                axes_mapper=old_label.meta.axes_mapper,
+                multiscale=old_label,
                 chunk_sizes=chunk_sizes,
             )
             new_label = new_ome_zarr_container.derive_label(
@@ -139,9 +136,7 @@ def rechunk_zarr(
                 label_pyramid_paths = old_label.meta.paths
                 for path in label_pyramid_paths:
                     old_label = ome_zarr_container.get_label(name=label, path=path)
-                    new_label.set_array(
-                        old_label.get_array(mode="dask")
-                    )
+                    new_label.set_array(old_label.get_array(mode="dask"))
 
     if overwrite_input:
         os.rename(zarr_url, f"{zarr_url}_tmp")

tests/test_add_z_singleton.py

@@ -60,7 +60,7 @@ def test_add_singleton(
 
     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
+        list(new_ome_zarr_container.image_meta.axes_handler.axes_names)
         == target_axes_names
     )
     assert new_ome_zarr_container.get_image().pixel_size.z == 1.0

tests/test_convert_2d_to_3d_segmentation.py

@@ -6,7 +6,6 @@
 import numpy as np
 import pandas as pd
 import pytest
-from ngio.images.label import build_masking_roi_table
 from ngio.tables import GenericTable
 
 from fractal_helper_tasks.convert_2D_segmentation_to_3D import (
@@ -134,10 +133,9 @@ def test_2d_to_3d_table_copying(tmp_path: Path):
 
     create_synthetic_data(zarr_url, zarr_url_3d, label_name, z_spacing=1.0)
     ome_zarr_2d = ngio.open_ome_zarr_container(zarr_url)
-    rois = ome_zarr_2d.get_table(
-        "masking_ROI_table", check_type="masking_roi_table"
-    ).rois()
-    assert rois[0].z_length == 1.0
+    roi_table = ome_zarr_2d.get_table("masking_ROI_table")
+    assert roi_table.table_type() == "masking_roi_table"
+    assert roi_table.rois()[0].z_length == 1.0
 
     convert_2D_segmentation_to_3D(
         zarr_url=zarr_url,
@@ -147,16 +145,18 @@ def test_2d_to_3d_table_copying(tmp_path: Path):
     ome_zarr_3d = ngio.open_ome_zarr_container(zarr_url_3d)
     # Validate correctness of tables
     assert ome_zarr_3d.list_tables() == ["masking_ROI_table", "image_ROI_table"]
-    rois = ome_zarr_3d.get_table(
-        "masking_ROI_table", check_type="masking_roi_table"
-    ).rois()
+    roi_table = ome_zarr_3d.get_table("masking_ROI_table")
+    assert roi_table.table_type() == "masking_roi_table"
+    rois = roi_table.rois()
     assert len(rois) == 1
     assert rois[0].name == "1"
     assert rois[0].x_length == 10.0
     # z_length goes from 1 to 10 because we have 10 z planes
     assert rois[0].z_length == 10.0
 
-    rois = ome_zarr_3d.get_table("image_ROI_table", check_type="roi_table").rois()
+    roi_table = ome_zarr_3d.get_table("image_ROI_table")
+    assert roi_table.table_type() == "roi_table"
+    rois = roi_table.rois()
     assert len(rois) == 1
     assert rois[0].name == "image_ROI_table"
     assert rois[0].x_length == 50.0
@@ -193,7 +193,7 @@ def test_2d_to_3d_real_data(tmp_zenodo_zarr: list[str]):
 
     # Create a masking roi table in the 2D image
     ome_zarr_2d = ngio.open_ome_zarr_container(zarr_url)
-    masking_roi_table = build_masking_roi_table(ome_zarr_2d.get_label(name=label_name))
+    masking_roi_table = ome_zarr_2d.build_masking_roi_table(label_name)
 
     ome_zarr_2d.add_table(
         name=tables_to_copy[0],
@@ -227,7 +227,7 @@ def test_2d_to_3d_real_data_no_label_copy(tmp_zenodo_zarr: list[str]):
 
     # Add a generic table to be copied over
     generic_table = GenericTable(
-        dataframe=pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
+        table_data=pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
     )
 
     ome_zarr_2d.add_table(

tests/test_drop_t_dimension.py

@@ -30,9 +30,10 @@ def test_drop_t_dimension(
     )
     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 == [
+
+    assert new_ome_zarr_container.image_meta.axes_handler.axes_names == (
         "c",
         "z",
         "y",
         "x",
-    ]
+    )

tests/test_rechunk_zarr.py

@@ -8,17 +8,16 @@
 )
 
 
-# FIXME: Reactive these tests with ngio 0.2.3
 @pytest.mark.parametrize(
     "chunk_sizes, output_chunk_sizes",
     [
         ({"x": 1000, "y": 1000}, (1, 1, 1000, 1000)),
         ({"X": 1000, "Y": 1000}, (1, 1, 1000, 1000)),
-        # ({"x": 6000, "y": 6000}, (1, 1, 2160, 5120)),
+        ({"x": 6000, "y": 6000}, (1, 1, 2160, 5120)),
         ({}, (1, 1, 2160, 2560)),
         ({"x": None, "y": None}, (1, 1, 2160, 2560)),
-        # ({"z": 10}, (1, 1, 2160, 2560)),
-        # ({"Z": 10}, (1, 1, 2160, 2560)),
+        ({"z": 10}, (1, 1, 2160, 2560)),
+        ({"Z": 10}, (1, 1, 2160, 2560)),
     ],
 )
 def test_rechunk_2d(tmp_zenodo_zarr: list[str], chunk_sizes, output_chunk_sizes):
@@ -33,12 +32,11 @@ def test_rechunk_2d(tmp_zenodo_zarr: list[str], chunk_sizes, output_chunk_sizes)
     assert chunk_sizes == output_chunk_sizes
 
 
-# FIXME: Reactive these tests with ngio 0.2.3
 @pytest.mark.parametrize(
     "chunk_sizes, output_chunk_sizes",
     [
         ({"x": None, "y": None}, (1, 1, 2160, 2560)),
-        # ({"z": 10}, (1, 2, 2160, 2560)),
+        ({"z": 10}, (1, 2, 2160, 2560)),
     ],
 )
 def test_rechunk_3d(tmp_zenodo_zarr: list[str], chunk_sizes, output_chunk_sizes):