distortion correction explicit coefficients as input

httomolibgpu/prep/alignment.py

@@ -32,7 +32,7 @@
 else:
     map_coordinates = Mock()
 
-from typing import Literal, List
+from typing import Literal, List, Optional
 
 __all__ = [
     "distortion_correction_proj_discorpy",
@@ -51,7 +51,10 @@
 # but with the additional params `order` and `mode`).
 def distortion_correction_proj_discorpy(
     data: cp.ndarray,
-    metadata_path: str,
+    metadata_path: Optional[str] = None,
+    xcenter: Optional[float] = None,
+    ycenter: Optional[float] = None,
+    list_fact: Optional[List[float]] = None,
     shift_xy: List[int] = [0, 0],
     step_xy: List[int] = [1, 1],
     order: int = 3,
@@ -73,9 +76,17 @@ def distortion_correction_proj_discorpy(
     data : cp.ndarray
         3D array.
 
-    metadata_path : str
-        The path to the file containing the distortion coefficients for the
-        data.
+    metadata_path : Optional, str
+        The path to the file containing the distortion coefficients for the data. When the file is provided the center values are ignored.
+
+    xcenter : Optional, float
+        Center of distortion in x-direction. From the left of the image.
+
+    ycenter : Optional, float
+        Center of distortion in y-direction. From the top of the image.
+
+    list_fact : Optional, list
+        Polynomial coefficients of the backward model.
 
     shift_xy: List[int]
          Centers of distortion in x (from the left of the image) and y directions (from the top of the image).
@@ -100,7 +111,17 @@ def distortion_correction_proj_discorpy(
     __check_if_data_correct_type(
         data, accepted_type=["float32", "uint16"], methods_name=methods_name
     )
-    __check_variable_type(metadata_path, [str], "metadata_path", [], methods_name)
+    __check_variable_type(
+        metadata_path, [str, type(None)], "metadata_path", [], methods_name
+    )
+    __check_variable_type(
+        xcenter, [float, int, type(None)], "xcenter", [], methods_name
+    )
+    __check_variable_type(
+        ycenter, [float, int, type(None)], "ycenter", [], methods_name
+    )
+    __check_variable_type(list_fact, [list, type(None)], "list_fact", [], methods_name)
+
     __check_variable_type(shift_xy, [list], "shift_xy", [], methods_name)
     __check_variable_type(step_xy, [list], "step_xy", [], methods_name)
     __check_variable_type(order, [int], "order", [], methods_name)
@@ -126,8 +147,14 @@ def distortion_correction_proj_discorpy(
     if len(data.shape) == 2:
         data = cp.expand_dims(data, axis=0)
 
-    # Get info from metadata txt file
-    xcenter, ycenter, list_fact = _load_metadata_txt(metadata_path)
+    if metadata_path is not None:
+        # Get info from metadata txt file
+        xcenter, ycenter, list_fact = _load_metadata_txt(metadata_path)
+    else:
+        if any(v is None for v in (xcenter, ycenter, list_fact)):
+            raise ValueError(
+                "When file with distortion coefficients is not provided, xcenter, ycenter and list_fact must be given."
+            )
 
     # Use preview information to offset the x and y coords of the center of
     # distortion
@@ -146,8 +173,8 @@ def distortion_correction_proj_discorpy(
     det_x_shift = shift_xy[0]
     det_y_shift = shift_xy[1]
 
-    xcenter = xcenter - det_x_shift
-    ycenter = ycenter - det_y_shift
+    xcenter -= det_x_shift
+    ycenter -= det_y_shift
 
     height, width = data.shape[1], data.shape[2]
     xu_list = cp.arange(width) - xcenter

tests/test_prep/test_alignment.py

@@ -24,7 +24,7 @@
     [distortion_correction_proj_discorpy],
     ids=["tomopy"],
 )
-def test_correct_distortion(
+def test_correct_distortion_file(
     distortion_correction_path,
     ensure_clean_memory,
     image,
@@ -43,7 +43,65 @@ def test_correct_distortion(
     shift_xy = [0, 0]
     step_xy = [1, 1]
     corrected_data = implementation(
-        im, distortion_coeffs_path, shift_xy, step_xy, order=1, mode="reflect"
+        im,
+        distortion_coeffs_path,
+        None,
+        None,
+        None,
+        shift_xy,
+        step_xy,
+        order=1,
+        mode="reflect",
+    ).get()
+
+    assert_allclose(np.mean(corrected_data), mean_value)
+    assert np.max(corrected_data) == max_value
+    assert corrected_data.dtype == cp.uint16
+    assert corrected_data.flags.c_contiguous
+
+
+@pytest.mark.parametrize(
+    "image, max_value, mean_value",
+    [
+        ("dot_pattern_03.tif", 255, 200.16733869461675),
+        ("peppers.tif", 228, 95.51871109008789),
+        ("cameraman.tif", 254, 122.2400016784668),
+    ],
+    ids=["dot_pattern_03", "peppers", "cameraman"],
+)
+@pytest.mark.parametrize(
+    "implementation",
+    [distortion_correction_proj_discorpy],
+    ids=["tomopy"],
+)
+def test_correct_distortion_parameters(
+    distortion_correction_path,
+    ensure_clean_memory,
+    image,
+    implementation,
+    max_value,
+    mean_value,
+):
+    distortion_coeffs_path = os.path.join(
+        distortion_correction_path, "distortion-coeffs.txt"
+    )
+
+    path = os.path.join(distortion_correction_path, image)
+    im_host = imread(path)
+    im = cp.asarray(im_host.astype(cp.uint16))
+
+    shift_xy = [0, 0]
+    step_xy = [1, 1]
+    corrected_data = implementation(
+        im,
+        None,
+        1907.252211174307,
+        1522.9526884119646,
+        [1.0003402629927707, 5.332383469534285e-10, 3.426741168508618e-09],
+        shift_xy,
+        step_xy,
+        order=1,
+        mode="reflect",
     ).get()
 
     assert_allclose(np.mean(corrected_data), mean_value)