nabu 2023.2.1__py3-none-any.whl → 2024.1.0rc3__py3-none-any.whl

nabu/stitching/utils.py

@@ -1,18 +1,24 @@
 from distutils.version import StrictVersion
-
 from typing import Optional, Union
 import logging
+import functools
 import numpy
+from scipy.ndimage import affine_transform
 from tomoscan.scanbase import TomoScanBase
-from nabu.misc import fourier_filters
-from nabu.stitching.overlap import OverlapStitchingStrategy, ZStichOverlapKernel
-from nabu.estimation.alignment import AlignmentBase
-from nabu.resources.dataset_analyzer import HDF5DatasetAnalyzer
+from tomoscan.volumebase import VolumeBase
+from nxtomo.utils.transformation import build_matrix, UDDetTransformation
 from silx.utils.enum import Enum as _Enum
-from scipy.ndimage import shift as scipy_shift
 from scipy.fft import rfftn as local_fftn
 from scipy.fft import irfftn as local_ifftn
-from nabu.resources.nxflatfield import update_dataset_info_flats_darks
+from silx.utils.enum import Enum as _Enum
+from nxtomo.utils.transformation import build_matrix, UDDetTransformation
+from tomoscan.scanbase import TomoScanBase
+from .overlap import OverlapStitchingStrategy, ZStichOverlapKernel
+from .alignment import AlignmentAxis1, AlignmentAxis2, PaddedRawData
+from ..misc import fourier_filters
+from ..estimation.alignment import AlignmentBase
+from ..resources.dataset_analyzer import HDF5DatasetAnalyzer
+from ..resources.nxflatfield import update_dataset_info_flats_darks
 
 try:
     import itk
@@ -40,7 +46,6 @@ class ShiftAlgorithm(_Enum):
 
     NABU_FFT = "nabu-fft"
     SKIMAGE = "skimage"
-    SHIFT_GRID = "shift-grid"
     ITK_IMG_REG_V4 = "itk-img-reg-v4"
     NONE = "None"
 
@@ -94,7 +99,6 @@ def find_frame_relative_shifts(
 ):
     from nabu.stitching.config import (
         KEY_WINDOW_SIZE,
-        KEY_SCORE_METHOD,
         KEY_LOW_PASS_FILTER,
         KEY_HIGH_PASS_FILTER,
     )  # avoid cyclic import
@@ -144,52 +148,49 @@ def find_frame_relative_shifts(
     initial_shifts = numpy.array(estimated_shifts).copy()
     extra_shifts = numpy.array([0.0, 0.0])
 
-    # 2.0 call cross correlation function from the estimated cor from motors
-    for axis, method, params in zip(
-        (0, 1),
-        (y_cross_correlation_function, x_cross_correlation_function),
-        (y_shifts_params, x_shifts_params),
-    ):
-        if method is ShiftAlgorithm.NABU_FFT:
-            extra_shifts[axis] = find_shift_correlate(img1=overlap_upper_frame, img2=overlap_lower_frame)[axis]
-        elif method is ShiftAlgorithm.SKIMAGE:
-            if not __has_sk_phase_correlation__:
-                raise ValueError("scikit-image not installed. Cannot do phase correlation from it")
-            else:
-                found_shift, _, _ = phase_cross_correlation(
-                    reference_image=overlap_upper_frame, moving_image=overlap_lower_frame, space="real"
-                )
-                extra_shifts[axis] = found_shift[axis]
-        elif method is ShiftAlgorithm.NONE:  # None as a string in case some uers give this value
-            # in the case we don't want to apply algorithm keep the initial 'guessed' shifts
-            continue
-        elif method is ShiftAlgorithm.SHIFT_GRID:
-            if axis == 0:
-                window_size = (int(y_shifts_params.get(KEY_WINDOW_SIZE, 200)), 0)
-            elif axis == 1:
-                window_size = (0, int(x_shifts_params.get(KEY_WINDOW_SIZE, 200)))
-            score_method = params.get(KEY_SCORE_METHOD, ScoreMethod.STD)
-            extra_shifts[axis] = -shift_grid_search(
-                img_1=overlap_upper_frame,
-                img_2=overlap_lower_frame,
-                window_sizes=window_size,
-                step_size=1,
-                axis=(axis,),
-                score_method=score_method,
-            )[axis]
-        elif method is ShiftAlgorithm.ITK_IMG_REG_V4:
-            extra_shifts[axis] = find_shift_with_itk(img1=overlap_upper_frame, img2=overlap_lower_frame)[axis]
+    def skimage_proxy(img1, img2):
+        if not __has_sk_phase_correlation__:
+            raise ValueError("scikit-image not installed. Cannot do phase correlation from it")
         else:
-            raise ValueError(f"requested cross correlation function not handled ({method})")
-    final_rel_shifts = numpy.array(extra_shifts) + initial_shifts
+            found_shift, _, _ = phase_cross_correlation(reference_image=img1, moving_image=img2, space="real")
+            return -found_shift
+
+    shift_methods = {
+        ShiftAlgorithm.NABU_FFT: functools.partial(
+            find_shift_correlate, img1=overlap_upper_frame, img2=overlap_lower_frame
+        ),
+        ShiftAlgorithm.SKIMAGE: functools.partial(skimage_proxy, img1=overlap_upper_frame, img2=overlap_lower_frame),
+        ShiftAlgorithm.ITK_IMG_REG_V4: functools.partial(
+            find_shift_with_itk, img1=overlap_upper_frame, img2=overlap_lower_frame
+        ),
+        ShiftAlgorithm.NONE: functools.partial(lambda: (0.0, 0.0)),
+    }
+
+    res_algo = {}
+    for shift_alg in set((x_cross_correlation_function, y_cross_correlation_function)):
+        if shift_alg not in shift_methods:
+            raise ValueError(f"requested image alignment function not handled ({shift_alg})")
+        try:
+            res_algo[shift_alg] = shift_methods[shift_alg]()
+        except Exception as e:
+            _logger.error(f"Failed to find shift from {shift_alg.value}. Error is {e}")
+            res_algo[shift_alg] = (0, 0)
+
+    extra_shifts = (
+        res_algo[y_cross_correlation_function][0],
+        res_algo[x_cross_correlation_function][1],
+    )
 
+    final_rel_shifts = numpy.array(extra_shifts) + initial_shifts
     return tuple([int(shift) for shift in final_rel_shifts])
 
 
 def find_volumes_relative_shifts(
-    upper_volume: numpy.ndarray,
-    lower_volume: numpy.ndarray,
+    upper_volume: VolumeBase,
+    lower_volume: VolumeBase,
     estimated_shifts,
+    dim_axis_1: int,
+    dtype,
     flip_ud_upper_frame: bool = False,
     flip_ud_lower_frame: bool = False,
     slice_for_shift: Union[int, str] = "middle",
@@ -197,15 +198,59 @@ def find_volumes_relative_shifts(
     y_cross_correlation_function=None,
     x_shifts_params: Optional[dict] = None,
     y_shifts_params: Optional[dict] = None,
+    alignment_axis_2="center",
+    alignment_axis_1="center",
 ):
+    """
+
+    :param int dim_axis_1: axis 1 dimension (to handle axis 1 alignment)
+    """
     if y_shifts_params is None:
         y_shifts_params = {}
 
     if x_shifts_params is None:
         x_shifts_params = {}
 
-    upper_frame = upper_volume.get_slice(slice_for_shift, axis=1)
-    lower_frame = lower_volume.get_slice(slice_for_shift, axis=1)
+    alignment_axis_2 = AlignmentAxis2.from_value(alignment_axis_2)
+    alignment_axis_1 = AlignmentAxis1.from_value(alignment_axis_1)
+    assert dim_axis_1 > 0, "dim_axis_1 <= 0"
+
+    if isinstance(slice_for_shift, str):
+        if slice_for_shift == "first":
+            slice_for_shift = 0
+        elif slice_for_shift == "last":
+            slice_for_shift = dim_axis_1
+        elif slice_for_shift == "middle":
+            slice_for_shift = dim_axis_1 // 2
+        else:
+            raise ValueError("invalid slice provided to search shift", slice_for_shift)
+
+    def get_slice_along_axis_1(volume: VolumeBase, index: int):
+        assert isinstance(index, int), f"index should be an int, {type(index)} provided"
+        volume_shape = volume.get_volume_shape()
+        if alignment_axis_1 is AlignmentAxis1.BACK:
+            front_empty_width = dim_axis_1 - volume_shape[1]
+            if index < front_empty_width:
+                return PaddedRawData.get_empty_frame(shape=(volume_shape[0], volume_shape[2]), dtype=dtype)
+            else:
+                return volume.get_slice(index=index - front_empty_width, axis=1)
+        elif alignment_axis_1 is AlignmentAxis1.FRONT:
+            if index >= volume_shape[1]:
+                return PaddedRawData.get_empty_frame(shape=(volume_shape[0], volume_shape[2]), dtype=dtype)
+            else:
+                return volume.get_slice(index=index, axis=1)
+        elif alignment_axis_1 is AlignmentAxis1.CENTER:
+            front_empty_width = (dim_axis_1 - volume_shape[1]) // 2
+            back_empty_width = dim_axis_1 - front_empty_width
+            if index < front_empty_width or index > back_empty_width:
+                return PaddedRawData.get_empty_frame(shape=(volume_shape[0], volume_shape[2]), dtype=dtype)
+            else:
+                return volume.get_slice(index=index - front_empty_width, axis=1)
+        else:
+            raise TypeError(f"unmanaged alignment mode {alignment_axis_1.value}")
+
+    upper_frame = get_slice_along_axis_1(upper_volume, index=slice_for_shift)
+    lower_frame = get_slice_along_axis_1(lower_volume, index=slice_for_shift)
     if flip_ud_upper_frame:
         upper_frame = numpy.flipud(upper_frame.copy())
     if flip_ud_lower_frame:
@@ -214,14 +259,35 @@ def find_volumes_relative_shifts(
     from nabu.stitching.config import KEY_WINDOW_SIZE  # avoid cyclic import
 
     w_window_size = int(y_shifts_params.get(KEY_WINDOW_SIZE, 400))
-    start_overlap = max(estimated_shifts[0] - w_window_size // 2, 0)
-    end_overlap = min(estimated_shifts[0] + w_window_size // 2, min(upper_frame.shape[0], lower_frame.shape[0]))
+    start_overlap = max(estimated_shifts[0] // 2 - w_window_size // 2, 0)
+    end_overlap = min(estimated_shifts[0] // 2 + w_window_size // 2, min(upper_frame.shape[0], lower_frame.shape[0]))
 
     if start_overlap == 0:
         overlap_upper_frame = upper_frame[-end_overlap:]
     else:
         overlap_upper_frame = upper_frame[-end_overlap:-start_overlap]
     overlap_lower_frame = lower_frame[start_overlap:end_overlap]
+
+    # align if necessary
+    if overlap_upper_frame.shape[1] != overlap_lower_frame.shape[1]:
+        overlap_frame_width = min(overlap_upper_frame.shape[1], overlap_lower_frame.shape[1])
+        if alignment_axis_2 is AlignmentAxis2.CENTER:
+            upper_frame_left_pos = overlap_upper_frame.shape[1] // 2 - overlap_frame_width // 2
+            upper_frame_right_pos = upper_frame_left_pos + overlap_frame_width
+            overlap_upper_frame = overlap_upper_frame[:, upper_frame_left_pos:upper_frame_right_pos]
+
+            lower_frame_left_pos = overlap_lower_frame.shape[1] // 2 - overlap_frame_width // 2
+            lower_frame_right_pos = lower_frame_left_pos + overlap_frame_width
+            overlap_lower_frame = overlap_lower_frame[:, lower_frame_left_pos:lower_frame_right_pos]
+        elif alignment_axis_2 is AlignmentAxis2.LEFT:
+            overlap_upper_frame = overlap_upper_frame[:, :overlap_frame_width]
+            overlap_lower_frame = overlap_lower_frame[:, :overlap_frame_width]
+        elif alignment_axis_2 is AlignmentAxis2.RIGTH:
+            overlap_upper_frame = overlap_upper_frame[:, -overlap_frame_width:]
+            overlap_lower_frame = overlap_lower_frame[:, -overlap_frame_width:]
+        else:
+            raise ValueError(f"Alignement {alignment_axis_2.value} is not handled")
+
     if not overlap_upper_frame.shape == overlap_lower_frame.shape:
         raise ValueError(f"Fail to get consistant overlap ({overlap_upper_frame.shape} vs {overlap_lower_frame.shape})")
 
@@ -293,7 +359,6 @@ def find_projections_relative_shifts(
     ):
         cor_options = x_shifts_params.copy()
         cor_options.pop("img_reg_method", None)
-        cor_options.pop("score_method", None)
         # remove all none numeric options because estimate_cor will call 'literal_eval' on them
 
         upper_scan_dataset_info = HDF5DatasetAnalyzer(
@@ -326,16 +391,24 @@ def find_projections_relative_shifts(
 
     # } else we will compute shift from the flat projections
 
-    def get_flat_fielded_proj(scan: TomoScanBase, proj_index: int, reverse: bool, revert_x: bool, revert_y):
+    def get_flat_fielded_proj(
+        scan: TomoScanBase, proj_index: int, reverse: bool, transformation_matrix: Optional[numpy.ndarray]
+    ):
         first_proj_idx = sorted(lower_scan.projections.keys(), reverse=reverse)[proj_index]
         ff = scan.flat_field_correction(
             (scan.projections[first_proj_idx],),
             (first_proj_idx,),
         )[0]
-        if revert_x:
-            ff = numpy.fliplr(ff)
-        if revert_y:
-            ff = numpy.flipud(ff)
+        assert ff.ndim == 2, f"expects a single 2D frame. Get something with {ff.ndim} dimensions"
+        if transformation_matrix is not None:
+            assert (
+                transformation_matrix.ndim == 2
+            ), f"expects a 2D transformation matrix. Get a {transformation_matrix.ndim} D"
+            if numpy.isclose(transformation_matrix[2, 2], -1):
+                transformation_matrix[2, :] = 0
+                transformation_matrix[0, 2] = 0
+                transformation_matrix[2, 2] = 1
+                ff = numpy.flipud(ff)
         return ff
 
     if isinstance(projection_for_shift, str):
@@ -357,26 +430,32 @@ def find_projections_relative_shifts(
             f"projection_for_shift is expected to be an int. Not {type(projection_for_shift)} - {projection_for_shift}"
         )
 
+    upper_scan_transformations = list(upper_scan.get_detector_transformations(tuple()))
+    if flip_ud_upper_frame:
+        upper_scan_transformations.append(UDDetTransformation())
+    upper_scan_trans_matrix = build_matrix(upper_scan_transformations)
+    lower_scan_transformations = list(lower_scan.get_detector_transformations(tuple()))
+    if flip_ud_lower_frame:
+        lower_scan_transformations.append(UDDetTransformation())
+    lower_scan_trans_matrix = build_matrix(lower_scan_transformations)
     upper_proj = get_flat_fielded_proj(
         upper_scan,
         projection_for_shift,
         reverse=False,
-        revert_x=upper_scan.get_x_flipped(default=False),
-        revert_y=upper_scan.get_y_flipped(default=False) ^ flip_ud_upper_frame,
+        transformation_matrix=upper_scan_trans_matrix,
     )
     lower_proj = get_flat_fielded_proj(
         lower_scan,
         projection_for_shift,
         reverse=invert_order,
-        revert_x=lower_scan.get_x_flipped(default=False),
-        revert_y=lower_scan.get_y_flipped(default=False) ^ flip_ud_lower_frame,
+        transformation_matrix=lower_scan_trans_matrix,
     )
 
     from nabu.stitching.config import KEY_WINDOW_SIZE  # avoid cyclic import
 
     w_window_size = int(y_shifts_params.get(KEY_WINDOW_SIZE, 400))
-    start_overlap = max(estimated_shifts[0] - w_window_size // 2, 0)
-    end_overlap = min(estimated_shifts[0] + w_window_size // 2, min(upper_proj.shape[0], lower_proj.shape[0]))
+    start_overlap = max(estimated_shifts[0] // 2 - w_window_size // 2, 0)
+    end_overlap = min(estimated_shifts[0] // 2 + w_window_size // 2, min(upper_proj.shape[0], lower_proj.shape[0]))
     if start_overlap == 0:
         overlap_upper_frame = upper_proj[-end_overlap:]
     else:
@@ -410,144 +489,7 @@ def find_shift_correlate(img1, img2, padding_mode="reflect"):
 
     (f_vals, fv, fh) = alignment.extract_peak_region_2d(cc, cc_vs=cc_vs, cc_hs=cc_hs)
     shifts_vh = alignment.refine_max_position_2d(f_vals, fv, fh)
-    return shifts_vh
-
-
-class ScoreMethod(_Enum):
-    STD = "standard deviation"
-    TV = "total variation"
-    TV_INVERSE = "1 / (total variation)"
-    STD_INVERSE = "1 / std"
-
-    @classmethod
-    def from_value(cls, value):
-        if isinstance(value, str):
-            # for string handle the case where value as been provided as 'value'. As there is spaces this can happen
-            value = value.lstrip("'").rstrip("'")
-        if value in ("tv", "TV"):
-            return ScoreMethod.TV
-        elif value in ("std", "STD"):
-            return ScoreMethod.STD
-        else:
-            return super().from_value(value=value)
-
-
-def compute_score_contrast_std(data: numpy.ndarray):
-    """
-    Compute a contrast score by simply computing the standard deviation of
-    the frame
-    :param numpy.ndarray data: frame for which we should compute the score
-    :return: score of the frame
-    :rtype: float
-    """
-    if data is None:
-        return None
-    else:
-        return data.std()
-
-
-def compute_tv_score(data: numpy.ndarray):
-    """
-    Compute the data score as image total variation
-
-    :param numpy.ndarray data: frame for which we should compute the score
-    :return: score of the frame
-    :rtype: float
-    """
-    tv = numpy.sum(numpy.sqrt(numpy.gradient(data, axis=0) ** 2 + numpy.gradient(data, axis=1) ** 2))
-    return tv
-
-
-def compute_score(img_1, img_2, shift, score_method, score_region, return_img=False):
-    score_method = ScoreMethod.from_value(score_method)
-
-    img_2 = scipy_shift(img_2, shift=shift)
-
-    img_2_reduce = img_2[
-        score_region[0].start : score_region[0].stop,
-        score_region[1].start : score_region[1].stop,
-    ]
-    img_1_reduce = img_1[
-        score_region[0].start : score_region[0].stop,
-        score_region[1].start : score_region[1].stop,
-    ]
-
-    img_sum = img_1_reduce * 0.5 + img_2_reduce * 0.5
-
-    if score_method is ScoreMethod.TV:
-        result = compute_tv_score(img_sum)
-    elif score_method is ScoreMethod.STD:
-        result = compute_score_contrast_std(img_sum)
-    elif score_method is ScoreMethod.TV_INVERSE:
-        result = 1 / compute_tv_score(img_sum)
-    elif score_method is ScoreMethod.STD_INVERSE:
-        result = 1 / compute_score_contrast_std(img_sum)
-    else:
-        raise ValueError(f"{score_method} is not handled")
-    if return_img:
-        return result, img_sum
-    else:
-        return result
-
-
-def shift_grid_search(img_1, img_2, window_sizes: tuple, axis, step_size, score_method=ScoreMethod.STD):
-    """
-    we could consider adding weights to do the exact same operation that will be done for the stitching... ? overkilled ?
-    :param tuple window_sizes: as y_size, x_size
-    """
-    if not isinstance(window_sizes, tuple) and len(window_sizes) != 2:
-        raise TypeError(f"window_sizes is expected to be a tuple of two ints. {window_sizes} provided")
-    if not img_1.ndim == img_2.ndim == 2:
-        raise ValueError("image dimension should be 2D")
-
-    for value in axis:
-        if not value in (0, 1):
-            raise ValueError(f"axis {value} is not handled")
-    half_window_x = min((img_1.shape[1], img_2.shape[1], abs(window_sizes[1]))) // 2
-    half_window_y = min((img_1.shape[0], img_2.shape[0], abs(window_sizes[0]))) // 2
-    if 1 in axis:
-        x_research = numpy.arange(-half_window_x, half_window_x, step_size)
-        x_score_region = slice(
-            int(half_window_x * 2),
-            int(min(img_1.shape[1], img_2.shape[1]) - (half_window_x * 2)),
-        )
-    else:
-        x_research = tuple((0,))
-        x_score_region = slice(0, int(min(img_1.shape[1], img_2.shape[1])))
-
-    if 0 in axis:
-        y_research = numpy.arange(-half_window_y, half_window_y, step_size)
-        y_score_region = slice(
-            int((half_window_y * 2)),
-            int(min(img_1.shape[0], img_2.shape[0]) - (half_window_y * 2)),
-        )
-    else:
-        y_research = tuple((0,))
-        y_score_region = slice(0, int(min(img_1.shape[0], img_2.shape[0])))
-    score_region = (
-        y_score_region,
-        x_score_region,
-    )
-
-    score_width = score_region[1].stop - score_region[1].start
-    if score_width < 10:
-        _logger.warning("score_width seems very low. Try reducing window_sizes")
-
-    best_score, best_shift = None, (0, 0)
-    for x_shift in x_research:
-        for y_shift in y_research:
-            res = compute_score(
-                img_1,
-                img_2.copy(),
-                shift=(y_shift, x_shift),
-                score_method=score_method,
-                score_region=score_region,
-            )
-            local_score = res
-            if best_score is None or (local_score is not None and local_score > best_score):
-                best_score = local_score
-                best_shift = (y_shift, x_shift)
-    return numpy.array(best_shift)
+    return -shifts_vh
 
 
 def find_shift_with_itk(img1: numpy.ndarray, img2: numpy.ndarray) -> tuple: