Add upper and lower quantile rescaling to calculate_registration_image_based.py for more robust registration

CHANGELOG.md

@@ -1,5 +1,8 @@
 **Note**: Numbers like (\#123) point to closed Pull Requests on the fractal-tasks-core repository.
 
+# 1.3.2
+* Tasks:
+    * Add percentile-based rescaling to calculate registration task to make it more robust (\#848)
 * Dependencies:
   * Relax pandas constraint to `<2`.
   * Relax torch constraint to `<=3.0.0`.

fractal_tasks_core/__FRACTAL_MANIFEST__.json

@@ -1123,6 +1123,18 @@
             "title": "Method",
             "description": "Method to use for image registration. The available methods are `phase_cross_correlation` (scikit-image package, works for 2D & 3D) and \"chi2_shift\" (image_registration package, only works for 2D images)."
           },
+          "lower_rescale_quantile": {
+            "default": 0.0,
+            "title": "Lower Rescale Quantile",
+            "type": "number",
+            "description": "Lower quantile for rescaling the image intensities before applying registration. Can be helpful to deal with image artifacts. Default is 0."
+          },
+          "upper_rescale_quantile": {
+            "default": 0.99,
+            "title": "Upper Rescale Quantile",
+            "type": "number",
+            "description": "Upper quantile for rescaling the image intensities before applying registration. Can be helpful to deal with image artifacts. Default is 0.99."
+          },
           "roi_table": {
             "default": "FOV_ROI_table",
             "title": "Roi Table",

fractal_tasks_core/tasks/calculate_registration_image_based.py

@@ -20,6 +20,7 @@
 import numpy as np
 import zarr
 from pydantic import validate_call
+from skimage.exposure import rescale_intensity
 from skimage.registration import phase_cross_correlation
 
 from fractal_tasks_core.channels import get_channel_from_image_zarr
@@ -77,6 +78,8 @@ def calculate_registration_image_based(
     # Core parameters
     wavelength_id: str,
     method: RegistrationMethod = RegistrationMethod.PHASE_CROSS_CORRELATION,
+    lower_rescale_quantile: float = 0.0,
+    upper_rescale_quantile: float = 0.99,
     roi_table: str = "FOV_ROI_table",
     level: int = 2,
 ) -> None:
@@ -102,6 +105,12 @@ def calculate_registration_image_based(
             are `phase_cross_correlation` (scikit-image package, works for 2D
             & 3D) and "chi2_shift" (image_registration package, only works for
             2D images).
+        lower_rescale_quantile: Lower quantile for rescaling the image
+            intensities before applying registration. Can be helpful
+             to deal with image artifacts. Default is 0.
+        upper_rescale_quantile: Upper quantile for rescaling the image
+            intensities before applying registration. Can be helpful
+            to deal with image artifacts. Default is 0.99.
         roi_table: Name of the ROI table over which the task loops to
             calculate the registration. Examples: `FOV_ROI_table` => loop over
             the field of views, `well_ROI_table` => process the whole well as
@@ -247,6 +256,22 @@ def calculate_registration_image_based(
             compute=compute,
         )
 
+        # Rescale the images
+        img_ref = rescale_intensity(
+            img_ref,
+            in_range=(
+                np.quantile(img_ref, lower_rescale_quantile),
+                np.quantile(img_ref, upper_rescale_quantile),
+            ),
+        )
+        img_acq_x = rescale_intensity(
+            img_acq_x,
+            in_range=(
+                np.quantile(img_acq_x, lower_rescale_quantile),
+                np.quantile(img_acq_x, upper_rescale_quantile),
+            ),
+        )
+
         ##############
         #  Calculate the transformation
         ##############

tests/tasks/test_registration.py

@@ -237,6 +237,8 @@ def test_multiplexing_registration(
     tmp_path,
     monkeypatch: MonkeyPatch,
     method: str,
+    lower_rescale_quantile: float = 0.0,
+    upper_rescale_quantile: float = 0.99,
     # Given the test data, only implemented per FOV
     roi_table="FOV_ROI_table",
 ):
@@ -348,6 +350,8 @@ def test_multiplexing_registration(
             init_args=image["init_args"],
             wavelength_id="A01_C01",
             method=method,
+            lower_rescale_quantile=lower_rescale_quantile,
+            upper_rescale_quantile=upper_rescale_quantile,
             roi_table=roi_table,
         )
 
@@ -438,6 +442,8 @@ def test_multiplexing_registration_3d(
     tmp_path,
     monkeypatch: MonkeyPatch,
     method: str,
+    lower_rescale_quantile: float = 0.0,
+    upper_rescale_quantile: float = 0.99,
     # Given the test data, only implemented per FOV
     roi_table="FOV_ROI_table",
 ):
@@ -526,6 +532,8 @@ def test_multiplexing_registration_3d(
                     init_args=image["init_args"],
                     wavelength_id="A01_C01",
                     method=method,
+                    lower_rescale_quantile=lower_rescale_quantile,
+                    upper_rescale_quantile=upper_rescale_quantile,
                     roi_table=roi_table,
                 )
             pytest.skip(