diff --git a/httomolibgpu/prep/alignment.py b/httomolibgpu/prep/alignment.py index 96dc14c6..8f282bba 100644 --- a/httomolibgpu/prep/alignment.py +++ b/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 diff --git a/tests/test_prep/test_alignment.py b/tests/test_prep/test_alignment.py index afbe7af5..40a11be7 100644 --- a/tests/test_prep/test_alignment.py +++ b/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)