Merge pull request #799 from fractal-analytics-platform/797-improve-t…

…ype-hints-and-docstring-for-calculate_registration_image_based Minor fixes for `calculate_registration_image_based`
fractal-analytics-platform · Jul 16, 2024 · dbf6d75 · dbf6d75
2 parents cd45cd3 + e756599
commit dbf6d75
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 8 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,8 +4,10 @@
 
 * 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 masking ROI label types in `masked_loading_wrapper` for Cellpose task (\#785).
     * 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).
 
 # 1.1.0
 
@@ -21,7 +23,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).

diff --git a/fractal_tasks_core/__FRACTAL_MANIFEST__.json b/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"

diff --git a/fractal_tasks_core/tasks/calculate_registration_image_based.py b/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