Merge pull request #866 from fractal-analytics-platform/ngio-projection

Ngio projection
fractal-analytics-platform · Dec 19, 2024 · b122c5b · b122c5b
2 parents f7779b9 + 4c12f2b
commit b122c5b
Show file tree

Hide file tree

Showing 11 changed files with 1,663 additions and 1,150 deletions.
diff --git a/.github/workflows/ci_pip.yml b/.github/workflows/ci_pip.yml
@@ -15,10 +15,8 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-22.04, macos-latest]
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12"]
         exclude:
-          - os: macos-latest
-            python-version: '3.9'
           - os: macos-latest
             python-version: '3.10'
     name: "Core, Python ${{ matrix.python-version }}, ${{ matrix.os }}"
@@ -50,10 +48,8 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest, macos-latest]
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12"]
         exclude:
-          - os: macos-latest
-            python-version: '3.9'
           - os: macos-latest
             python-version: '3.10'
     name: "Tasks, Python ${{ matrix.python-version }}, ${{ matrix.os }}"

diff --git a/.github/workflows/ci_poetry.yml b/.github/workflows/ci_poetry.yml
@@ -15,7 +15,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12"]
 
     steps:
       - uses: actions/checkout@v4
@@ -60,7 +60,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12"]
 
     steps:
       - uses: actions/checkout@v4
@@ -131,7 +131,7 @@ jobs:
           python-version: "3.10"
 
       - name: Install dependencies
-        run: poetry install --only dev
+        run: poetry install --with dev -E fractal-tasks
 
       - name: Download data
         uses: actions/download-artifact@v4

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,7 +1,15 @@
 **Note**: Numbers like (\#123) point to closed Pull Requests on the fractal-tasks-core repository.
 
+
 # Unreleased
 
+* Tasks:
+    * Refactor projection task to use ngio
+* Dependencies:
+    * Add `ngio==0.1.4` to the dependencies
+    * Require `python >=3.10,<3.13`
+* CI:
+    * Remove Python 3.9 from the CI matrix
 * Tests:
     * Use locked version of `coverage` in GitHub action (\#882).
     * Bump `coverage` version from 6.5 to 7.6 (\#882).
@@ -13,6 +21,7 @@
     * Support providing `docs_info=file:task_info/description.md` (\#876).
     * Deprecate `check_manifest.py` module, in favor of additional GitHub action steps (\#876).
 
+
 # 1.3.3
 
 * Add new metadata (authors, category, modality, tags) to manifest models and to tasks (\#855).

diff --git a/docs/development.md b/docs/development.md
@@ -36,7 +36,7 @@ poetry run pytest --ignore tests/tasks
 
 The tests files are in the `tests` folder of the repository. Its structure reflects the `fractal_tasks_core` structure, with tests for the core library in the main folder and tests for `tasks` and `dev` subpckages in their own subfolders.
 
-Tests are also run through GitHub Actions, with Python 3.9, 3.10 and 3.11. Note that within GitHub actions we run tests for both the `poetry`-installed and `pip`-installed versions of the code, which may e.g. have different versions of some dependencies (since `pip install` does not rely on the `poetry.lock` lockfile).
+Tests are also run through GitHub Actions, with Python 3.10, 3.11 and 3.12. Note that within GitHub actions we run tests for both the `poetry`-installed and `pip`-installed versions of the code, which may e.g. have different versions of some dependencies (since `pip install` does not rely on the `poetry.lock` lockfile).
 
 ## Documentation
 

diff --git a/fractal_tasks_core/__FRACTAL_MANIFEST__.json b/fractal_tasks_core/__FRACTAL_MANIFEST__.json
@@ -653,12 +653,17 @@
               "overwrite": {
                 "title": "Overwrite",
                 "type": "boolean"
+              },
+              "new_plate_name": {
+                "title": "New Plate Name",
+                "type": "string"
               }
             },
             "required": [
               "origin_url",
               "method",
-              "overwrite"
+              "overwrite",
+              "new_plate_name"
             ],
             "title": "InitArgsMIP",
             "type": "object"

diff --git a/fractal_tasks_core/tasks/copy_ome_zarr_hcs_plate.py b/fractal_tasks_core/tasks/copy_ome_zarr_hcs_plate.py
@@ -244,7 +244,10 @@ def copy_ome_zarr_hcs_plate(
         parallelization_item = dict(
             zarr_url=new_zarr_url,
             init_args=dict(
-                origin_url=zarr_url, method=method.value, overwrite=overwrite
+                origin_url=zarr_url,
+                method=method.value,
+                overwrite=overwrite,
+                new_plate_name=f"{new_plate_name}.zarr",
             ),
         )
         InitArgsMIP(**parallelization_item["init_args"])

diff --git a/fractal_tasks_core/tasks/io_models.py b/fractal_tasks_core/tasks/io_models.py
@@ -86,11 +86,13 @@ class InitArgsMIP(BaseModel):
         origin_url: Path to the zarr_url with the 3D data
         method: Projection method to be used. See `DaskProjectionMethod`
         overwrite: If `True`, overwrite the task output.
+        new_plate_name: Name of the new OME-Zarr HCS plate
     """
 
     origin_url: str
     method: str
     overwrite: bool
+    new_plate_name: str
 
 
 class MultiplexingAcquisition(BaseModel):

diff --git a/fractal_tasks_core/tasks/projection.py b/fractal_tasks_core/tasks/projection.py
@@ -12,30 +12,39 @@
 """
 Task for 3D->2D maximum-intensity projection.
 """
+from __future__ import annotations
+
 import logging
 from typing import Any
 
-import anndata as ad
 import dask.array as da
-import zarr
+from ngio import NgffImage
+from ngio.core import Image
 from pydantic import validate_call
-from zarr.errors import ContainsArrayError
-
-from fractal_tasks_core.ngff import load_NgffImageMeta
-from fractal_tasks_core.pyramids import build_pyramid
-from fractal_tasks_core.roi import (
-    convert_ROIs_from_3D_to_2D,
-)
-from fractal_tasks_core.tables import write_table
-from fractal_tasks_core.tables.v1 import get_tables_list_v1
+
 from fractal_tasks_core.tasks.io_models import InitArgsMIP
 from fractal_tasks_core.tasks.projection_utils import DaskProjectionMethod
-from fractal_tasks_core.zarr_utils import OverwriteNotAllowedError
-
 
 logger = logging.getLogger(__name__)
 
 
+def _compute_new_shape(source_image: Image) -> tuple[int]:
+    """Compute the new shape of the image after the projection.
+
+    The new shape is the same as the original one,
+    except for the z-axis, which is set to 1.
+    """
+    on_disk_shape = source_image.on_disk_shape
+    logger.info(f"Source {on_disk_shape=}")
+
+    on_disk_z_index = source_image.dataset.on_disk_axes_names.index("z")
+
+    dest_on_disk_shape = list(on_disk_shape)
+    dest_on_disk_shape[on_disk_z_index] = 1
+    logger.info(f"Destination {dest_on_disk_shape=}")
+    return tuple(dest_on_disk_shape)
+
+
 @validate_call
 def projection(
     *,
@@ -60,123 +69,68 @@ def projection(
     logger.info(f"{method=}")
 
     # Read image metadata
-    ngff_image = load_NgffImageMeta(init_args.origin_url)
-    # Currently not using the validation models due to wavelength_id issue
-    # See #681 for discussion
-    # new_attrs = ngff_image.model_dump(exclude_none=True)
-    # Current way to get the necessary metadata for MIP
-    group = zarr.open_group(init_args.origin_url, mode="r")
-    new_attrs = group.attrs.asdict()
-
-    # Create the zarr image with correct
-    new_image_group = zarr.group(zarr_url)
-    new_image_group.attrs.put(new_attrs)
-
-    # Load 0-th level
-    data_czyx = da.from_zarr(init_args.origin_url + "/0")
-    num_channels = data_czyx.shape[0]
-    chunksize_y = data_czyx.chunksize[-2]
-    chunksize_x = data_czyx.chunksize[-1]
-    logger.info(f"{num_channels=}")
-    logger.info(f"{chunksize_y=}")
-    logger.info(f"{chunksize_x=}")
-
-    # Loop over channels
-    accumulate_chl = []
-    for ind_ch in range(num_channels):
-        # Perform MIP for each channel of level 0
-        project_yx = da.stack(
-            [method.apply(data_czyx[ind_ch], axis=0)], axis=0
-        )
-        accumulate_chl.append(project_yx)
-    accumulated_array = da.stack(accumulate_chl, axis=0)
-
-    # Write to disk (triggering execution)
-    try:
-        accumulated_array.to_zarr(
-            f"{zarr_url}/0",
-            overwrite=init_args.overwrite,
-            dimension_separator="/",
-            write_empty_chunks=False,
-        )
-    except ContainsArrayError as e:
-        error_msg = (
-            f"Cannot write array to zarr group at '{zarr_url}/0', "
-            f"with {init_args.overwrite=} (original error: {str(e)}).\n"
-            "Hint: try setting overwrite=True."
+    original_ngff_image = NgffImage(init_args.origin_url)
+    orginal_image = original_ngff_image.get_image()
+
+    if orginal_image.is_2d or orginal_image.is_2d_time_series:
+        raise ValueError(
+            "The input image is 2D, "
+            "projection is only supported for 3D images."
         )
-        logger.error(error_msg)
-        raise OverwriteNotAllowedError(error_msg)
 
-    # Starting from on-disk highest-resolution data, build and write to disk a
-    # pyramid of coarser levels
-    build_pyramid(
-        zarrurl=zarr_url,
+    # Compute the new shape and pixel size
+    dest_on_disk_shape = _compute_new_shape(orginal_image)
+
+    dest_pixel_size = orginal_image.pixel_size
+    dest_pixel_size.z = 1.0
+    logger.info(f"New shape: {dest_on_disk_shape=}")
+
+    # Create the new empty image
+    new_ngff_image = original_ngff_image.derive_new_image(
+        store=zarr_url,
+        name="MIP",
+        on_disk_shape=dest_on_disk_shape,
+        pixel_sizes=dest_pixel_size,
         overwrite=init_args.overwrite,
-        num_levels=ngff_image.num_levels,
-        coarsening_xy=ngff_image.coarsening_xy,
-        chunksize=(1, 1, chunksize_y, chunksize_x),
+        copy_labels=False,
+        copy_tables=True,
     )
+    logger.info(f"New Projection image created - {new_ngff_image=}")
+    new_image = new_ngff_image.get_image()
 
-    # Copy over any tables from the original zarr
-    # Generate the list of tables:
-    tables = get_tables_list_v1(init_args.origin_url)
-    roi_tables = get_tables_list_v1(init_args.origin_url, table_type="ROIs")
-    non_roi_tables = [table for table in tables if table not in roi_tables]
-
-    for table in roi_tables:
-        logger.info(
-            f"Reading {table} from "
-            f"{init_args.origin_url=}, convert it to 2D, and "
-            "write it back to the new zarr file."
-        )
-        new_ROI_table = ad.read_zarr(f"{init_args.origin_url}/tables/{table}")
-        old_ROI_table_attrs = zarr.open_group(
-            f"{init_args.origin_url}/tables/{table}"
-        ).attrs.asdict()
-
-        # Convert 3D ROIs to 2D
-        pxl_sizes_zyx = ngff_image.get_pixel_sizes_zyx(level=0)
-        new_ROI_table = convert_ROIs_from_3D_to_2D(
-            new_ROI_table, pixel_size_z=pxl_sizes_zyx[0]
-        )
-        # Write new table
-        write_table(
-            new_image_group,
-            table,
-            new_ROI_table,
-            table_attrs=old_ROI_table_attrs,
-            overwrite=init_args.overwrite,
-        )
+    # Process the image
+    z_axis_index = orginal_image.find_axis("z")
+    source_dask = orginal_image.get_array(
+        mode="dask", preserve_dimensions=True
+    )
 
-    for table in non_roi_tables:
-        logger.info(
-            f"Reading {table} from "
-            f"{init_args.origin_url=}, and "
-            "write it back to the new zarr file."
-        )
-        new_non_ROI_table = ad.read_zarr(
-            f"{init_args.origin_url}/tables/{table}"
-        )
-        old_non_ROI_table_attrs = zarr.open_group(
-            f"{init_args.origin_url}/tables/{table}"
-        ).attrs.asdict()
-
-        # Write new table
-        write_table(
-            new_image_group,
-            table,
-            new_non_ROI_table,
-            table_attrs=old_non_ROI_table_attrs,
-            overwrite=init_args.overwrite,
-        )
+    dest_dask = method.apply(dask_array=source_dask, axis=z_axis_index)
+    dest_dask = da.expand_dims(dest_dask, axis=z_axis_index)
+    new_image.set_array(dest_dask)
+    new_image.consolidate()
+    # Ends
+
+    # Copy over the tables
+    for roi_table_name in new_ngff_image.tables.list(table_type="roi_table"):
+        table = new_ngff_image.tables.get_table(roi_table_name)
+
+        roi_list = []
+        for roi in table.rois:
+            roi.z = 0.0
+            roi.z_length = 1.0
+            roi_list.append(roi)
+
+        table.set_rois(roi_list, overwrite=True)
+        table.consolidate()
+        logger.info(f"Table {roi_table_name} Projection done")
 
     # Generate image_list_updates
     image_list_update_dict = dict(
         image_list_updates=[
             dict(
                 zarr_url=zarr_url,
                 origin=init_args.origin_url,
+                attributes=dict(plate=init_args.new_plate_name),
                 types=dict(is_3D=False),
             )
         ]