Merge branch 'main' into pydantic-v2-second-attempt

Author: tcompa

CHANGELOG.md

@@ -3,8 +3,11 @@
 # Unreleased
 
 * Tasks:
-    * Fix issue with masked ROI & relabeling in Cellpose task (\#785).
-    * Fix issue with masking ROI label types in masked_loading_wrapper for Cellpose task (\#785).
+    * Fix issue with masked ROI & relabeling in Cellpose task (\#786).
+    * Fix issue with masking ROI label types in `masked_loading_wrapper` for Cellpose task (\#786).
+    * Enable workaround to support yx images in Cellpose task (\#789).
+    * Fix error handling in `calculate_registration_image_based` (\#799).
+    * Fix minor issues with call-signature and type hints in `calculate_registration_image_based` (\#799).
 * Development:
     * Switch all core models to pydantic v2 (\#790).
 
@@ -22,7 +25,7 @@
     * Rename `task.cellpose_transforms` into `tasks.cellpose_utils` (\#738).
     * Fix wrong repeated overlap checks for bounding-boxes in Cellpose task (\#778).
     * Fix minor MIP issues related to plate metadata and expecting acquisition metadata in all NGFF plates(\#781).
-    * Add chi2_shift option to Calculate Registration (image-based) task(\#741).
+    * Add `chi2_shift` option to Calculate Registration (image-based) task(\#741).
 * Development:
     * Switch to transitional pydantic.v1 imports, changes pydantic requirement to `==1.10.16` or `>=2.6.3` (\#760).
     * Support JSON-Schema generation for `Enum` task arguments (\#749).

fractal_tasks_core/__FRACTAL_MANIFEST__.json

@@ -1076,7 +1076,7 @@
           },
           "RegistrationMethod": {
             "title": "RegistrationMethod",
-            "description": "An enumeration.",
+            "description": "RegistrationMethod Enum class",
             "enum": [
               "phase_cross_correlation",
               "chi2_shift"

fractal_tasks_core/tasks/calculate_registration_image_based.py

@@ -48,6 +48,16 @@
 
 
 class RegistrationMethod(Enum):
+    """
+    RegistrationMethod Enum class
+
+    Attributes:
+        PHASE_CROSS_CORRELATION: phase cross correlation based on scikit-image
+            (works with 2D & 3D images).
+        CHI2_SHIFT: chi2 shift based on image-registration library
+            (only works with 2D images).
+    """
+
     PHASE_CROSS_CORRELATION = "phase_cross_correlation"
     CHI2_SHIFT = "chi2_shift"
 
@@ -66,7 +76,7 @@ def calculate_registration_image_based(
     init_args: InitArgsRegistration,
     # Core parameters
     wavelength_id: str,
-    method: RegistrationMethod = "phase_cross_correlation",
+    method: RegistrationMethod = RegistrationMethod.PHASE_CROSS_CORRELATION,
     roi_table: str = "FOV_ROI_table",
     level: int = 2,
 ) -> None:
@@ -137,11 +147,11 @@ def calculate_registration_image_based(
     # Check if data is 3D (as not all registration methods work in 3D)
     # TODO: Abstract this check into a higher-level Zarr loading class
     if is_3D(data_reference_zyx):
-        if method == "chi2_shift":
+        if method == RegistrationMethod(RegistrationMethod.CHI2_SHIFT):
             raise ValueError(
-                "The `chi2_shift` registration method has not been "
-                "implemented for 3D images and the input image had a shape of "
-                f"{data_reference_zyx.shape}."
+                f"The `{RegistrationMethod.CHI2_SHIFT}` registration method "
+                "has not been implemented for 3D images and the input image "
+                f"had a shape of {data_reference_zyx.shape}."
             )
 
     # Read ROIs

fractal_tasks_core/tasks/cellpose_segmentation.py

@@ -314,11 +314,21 @@ def cellpose_segmentation(
             output_label_name = f"label_{ind_channel}"
 
     # Load ZYX data
-    data_zyx = da.from_zarr(f"{zarr_url}/{level}")[ind_channel]
+    # Workaround for #788: Only load channel index when there is a channel
+    # dimension
+    if ngff_image_meta.axes_names[0] != "c":
+        data_zyx = da.from_zarr(f"{zarr_url}/{level}")
+        if channel2.is_set():
+            raise ValueError(
+                "Dual channel input was specified for an OME-Zarr image "
+                "without a channel axis"
+            )
+    else:
+        data_zyx = da.from_zarr(f"{zarr_url}/{level}")[ind_channel]
+        if channel2.is_set():
+            data_zyx_c2 = da.from_zarr(f"{zarr_url}/{level}")[ind_channel_c2]
+            logger.info(f"Second channel: {data_zyx_c2.shape=}")
     logger.info(f"{data_zyx.shape=}")
-    if channel2.is_set():
-        data_zyx_c2 = da.from_zarr(f"{zarr_url}/{level}")[ind_channel_c2]
-        logger.info(f"Second channel: {data_zyx_c2.shape=}")
 
     # Read ROI table
     ROI_table_path = f"{zarr_url}/tables/{input_ROI_table}"
@@ -365,18 +375,21 @@ def cellpose_segmentation(
         logger.info(f"Anisotropy: {advanced_cellpose_model_params.anisotropy}")
 
     # Rescale datasets (only relevant for level>0)
+    # Workaround for #788
     if ngff_image_meta.axes_names[0] != "c":
-        raise ValueError(
-            "Cannot set `remove_channel_axis=True` for multiscale "
-            f"metadata with axes={ngff_image_meta.axes_names}. "
-            'First axis should have name "c".'
+        new_datasets = rescale_datasets(
+            datasets=[ds.dict() for ds in ngff_image_meta.datasets],
+            coarsening_xy=coarsening_xy,
+            reference_level=level,
+            remove_channel_axis=False,
+        )
+    else:
+        new_datasets = rescale_datasets(
+            datasets=[ds.dict() for ds in ngff_image_meta.datasets],
+            coarsening_xy=coarsening_xy,
+            reference_level=level,
+            remove_channel_axis=True,
         )
-    new_datasets = rescale_datasets(
-        datasets=[ds.dict() for ds in ngff_image_meta.datasets],
-        coarsening_xy=coarsening_xy,
-        reference_level=level,
-        remove_channel_axis=True,
-    )
 
     label_attrs = {
         "image-label": {