nabu 2024.1.10__py3-none-any.whl → 2024.2.0__py3-none-any.whl

nabu/stitching/tests/test_overlap.py

@@ -3,14 +3,52 @@ import pytest
 
 from nabu.stitching.overlap import compute_image_minimum_divergence, compute_image_higher_signal, check_overlaps
 from nabu.testutils import get_data
+from nabu.stitching.overlap import ImageStichOverlapKernel, OverlapStitchingStrategy
+from nabu.stitching.stitcher_2D import stitch_raw_frames
+from silx.image.phantomgenerator import PhantomGenerator
+
+
+strategies_to_test_weights = (
+    OverlapStitchingStrategy.CLOSEST,
+    OverlapStitchingStrategy.COSINUS_WEIGHTS,
+    OverlapStitchingStrategy.LINEAR_WEIGHTS,
+    OverlapStitchingStrategy.MEAN,
+)
+
+
+@pytest.mark.parametrize("strategy", strategies_to_test_weights)
+@pytest.mark.parametrize("stitching_axis", (0, 1))
+def test_overlap_stitcher(strategy, stitching_axis):
+    frame_width = 128
+    frame_height = frame_width
+    frame_1 = PhantomGenerator.get2DPhantomSheppLogan(n=frame_width)
+    stitcher = ImageStichOverlapKernel(
+        stitching_strategy=strategy,
+        overlap_size=frame_height,
+        frame_unstitched_axis_size=128,
+        stitching_axis=stitching_axis,
+    )
+    stitched_frame = stitcher.stitch(frame_1, frame_1)[0]
+    assert stitched_frame.shape == (frame_height, frame_width)
+    # check result is close to the expected one
+    numpy.testing.assert_allclose(frame_1, stitched_frame, atol=10e-10)
+
+    # check sum of weights ~ 1.0
+    numpy.testing.assert_allclose(
+        stitcher.weights_img_1 + stitcher.weights_img_2,
+        numpy.ones_like(stitcher.weights_img_1),
+    )
 
 
-def test_compute_image_minimum_divergence():
+@pytest.mark.parametrize("stitching_axis", (0, 1))
+def test_compute_image_minimum_divergence(stitching_axis):
     """make sure the compute_image_minimum_divergence function is processing"""
     raw_data_1 = get_data("brain_phantom.npz")["data"]
     raw_data_2 = numpy.random.rand(*raw_data_1.shape) * 255.0
 
-    stitching = compute_image_minimum_divergence(raw_data_1, raw_data_2, high_frequency_threshold=2)
+    stitching = compute_image_minimum_divergence(
+        raw_data_1, raw_data_2, high_frequency_threshold=2, stitching_axis=stitching_axis
+    )
     assert stitching.shape == raw_data_1.shape
 
 
@@ -43,7 +81,7 @@ def test_check_overlaps():
             numpy.ones(10),
             numpy.ones(20),
         ),
-        positions=((10,), (0,)),
+        positions=((10, 0, 0), (0, 0, 0)),
         axis=0,
         raise_error=True,
     )
@@ -55,7 +93,7 @@ def test_check_overlaps():
                 numpy.ones(10),
                 numpy.ones(20),
             ),
-            positions=((0,), (100,)),
+            positions=((0, 0, 0), (100, 0, 0)),
             axis=0,
             raise_error=True,
         )
@@ -67,7 +105,7 @@ def test_check_overlaps():
                 numpy.ones(20),
                 numpy.ones(10),
             ),
-            positions=((8,), (5,)),
+            positions=((8, 0, 0), (5, 0, 0)),
             axis=0,
             raise_error=True,
         )
@@ -79,7 +117,7 @@ def test_check_overlaps():
             numpy.ones(20),
             numpy.ones(10),
         ),
-        positions=((20,), (10,), (0,)),
+        positions=((20, 0, 0), (10, 0, 0), (0, 0, 0)),
         axis=0,
         raise_error=True,
     )
@@ -92,7 +130,111 @@ def test_check_overlaps():
                 numpy.ones(10),
                 numpy.ones(10),
             ),
-            positions=((20,), (15,), (11,)),
+            positions=((20, 0, 0), (15, 0, 0), (11, 0, 0)),
             axis=0,
             raise_error=True,
         )
+
+
+@pytest.mark.parametrize("dtype", (numpy.float16, numpy.float32))
+def test_stitch_vertically_raw_frames(dtype):
+    """
+    ensure a stitching with 3 frames and different overlap can be done
+    """
+    ref_frame_width = 256
+    frame_ref = PhantomGenerator.get2DPhantomSheppLogan(n=ref_frame_width).astype(dtype)
+
+    # split the frame into several part
+    frame_1 = frame_ref[0:100]
+    frame_2 = frame_ref[80:164]
+    frame_3 = frame_ref[154:]
+
+    kernel_1 = ImageStichOverlapKernel(frame_unstitched_axis_size=ref_frame_width, overlap_size=20, stitching_axis=0)
+    kernel_2 = ImageStichOverlapKernel(frame_unstitched_axis_size=ref_frame_width, overlap_size=10, stitching_axis=0)
+
+    stitched = stitch_raw_frames(
+        frames=(frame_1, frame_2, frame_3),
+        output_dtype=dtype,
+        overlap_kernels=(kernel_1, kernel_2),
+        raw_frames_compositions=None,
+        overlap_frames_compositions=None,
+        key_lines=(
+            (
+                90,  # frame_1 height - kernel_1 height / 2.0
+                10,  # kernel_1 height / 2.0
+            ),
+            (
+                79,  # frame_2 height - kernel_2 height / 2.0 ou 102-20 ?
+                5,  # kernel_2 height / 2.0
+            ),
+        ),
+    )
+
+    assert stitched.shape == frame_ref.shape
+    numpy.testing.assert_array_almost_equal(frame_ref, stitched)
+
+
+def test_stitch_vertically_raw_frames_2():
+    """
+    ensure a stitching with 3 frames and different overlap can be done
+    """
+    ref_frame_width = 256
+    frame_ref = PhantomGenerator.get2DPhantomSheppLogan(n=ref_frame_width).astype(numpy.float32)
+
+    # split the frame into several part
+    frame_1 = frame_ref.copy()
+    frame_2 = frame_ref.copy()
+    frame_3 = frame_ref.copy()
+
+    kernel_1 = ImageStichOverlapKernel(frame_unstitched_axis_size=ref_frame_width, overlap_size=10, stitching_axis=0)
+    kernel_2 = ImageStichOverlapKernel(frame_unstitched_axis_size=ref_frame_width, overlap_size=10, stitching_axis=0)
+
+    stitched = stitch_raw_frames(
+        frames=(frame_1, frame_2, frame_3),
+        output_dtype=numpy.float32,
+        overlap_kernels=(kernel_1, kernel_2),
+        raw_frames_compositions=None,
+        overlap_frames_compositions=None,
+        key_lines=((20, 20), (105, 105)),
+    )
+
+    assert stitched.shape == frame_ref.shape
+    numpy.testing.assert_array_almost_equal(frame_ref, stitched)
+
+
+@pytest.mark.parametrize("dtype", (numpy.float16, numpy.float32))
+def test_stitch_horizontally_raw_frames(dtype):
+    """
+    ensure a stitching with 3 frames and different overlap can be done along axis 1
+    """
+    ref_frame_width = 256
+    frame_ref = PhantomGenerator.get2DPhantomSheppLogan(n=ref_frame_width).astype(dtype)
+
+    # split the frame into several part
+    frame_1 = frame_ref[:, 0:100]
+    frame_2 = frame_ref[:, 80:164]
+    frame_3 = frame_ref[:, 154:]
+
+    kernel_1 = ImageStichOverlapKernel(frame_unstitched_axis_size=ref_frame_width, overlap_size=20, stitching_axis=1)
+    kernel_2 = ImageStichOverlapKernel(frame_unstitched_axis_size=ref_frame_width, overlap_size=10, stitching_axis=1)
+
+    stitched = stitch_raw_frames(
+        frames=(frame_1, frame_2, frame_3),
+        output_dtype=dtype,
+        overlap_kernels=(kernel_1, kernel_2),
+        raw_frames_compositions=None,
+        overlap_frames_compositions=None,
+        key_lines=(
+            (
+                90,  # frame_1 height - kernel_1 height / 2.0
+                10,  # kernel_1 height / 2.0
+            ),
+            (
+                79,  # frame_2 height - kernel_2 height / 2.0 ou 102-20 ?
+                5,  # kernel_2 height / 2.0
+            ),
+        ),
+    )
+
+    assert stitched.shape == frame_ref.shape
+    numpy.testing.assert_array_almost_equal(frame_ref, stitched)

nabu/stitching/tests/test_utils.py

@@ -1,4 +1,4 @@
-from nabu.stitching.utils import has_itk, find_shift_with_itk
+from nabu.stitching.utils.utils import has_itk, find_shift_with_itk
 from scipy.ndimage import shift as shift_scipy
 import numpy
 import pytest

nabu/stitching/tests/test_y_preprocessing_stitching.py

@@ -0,0 +1,132 @@
+import os
+import pytest
+import numpy
+from tqdm import tqdm
+
+from nabu.stitching.y_stitching import y_stitching
+from nabu.stitching.config import PreProcessedYStitchingConfiguration
+from nxtomo.application.nxtomo import NXtomo
+from nxtomo.nxobject.nxdetector import ImageKey
+from tomoscan.esrf.scan.nxtomoscan import NXtomoScan
+
+
+def build_nxtomos(output_dir, flip_lr, flip_ud) -> tuple:
+    r"""
+    build two nxtomos in output_dir and return the list of NXtomos ready to be stitched
+      /\
+      |        ______________       ______________
+      |       |~           ~~|      |~             |
+      |       |~ nxtomo 1  ~~|      |~ nxtomo 0    |
+    Z |       |~ frame     ~~|      |~ frame       |
+              |______________|      |______________|
+    <-----------------------------------------------
+                    90                     40      0
+                    y (in acquisition space)
+    * ~: represent the overlap area
+    """
+    dark_data = numpy.array([0] * 64 * 120, dtype=numpy.float32).reshape((64, 120))
+    flat_data = numpy.array([1] * 64 * 120, dtype=numpy.float32).reshape((64, 120))
+    normalized_data = numpy.linspace(128, 1024, num=64 * 120, dtype=numpy.float32).reshape((64, 120))
+    if flip_lr:
+        dark_data = numpy.fliplr(dark_data)
+        flat_data = numpy.fliplr(flat_data)
+        normalized_data = numpy.fliplr(normalized_data)
+    if flip_ud:
+        dark_data = numpy.flipud(dark_data)
+        flat_data = numpy.flipud(flat_data)
+        normalized_data = numpy.flipud(normalized_data)
+
+    raw_data = (normalized_data + dark_data) * (flat_data + dark_data)
+
+    # create raw data
+    scans = []
+    slices = (slice(0, 80), slice(60, -1))
+    frame_y_positions = (40, 90)
+    for i_nxtomo, (my_slice, frame_y_position) in enumerate(zip(slices, frame_y_positions)):
+        my_raw_data = raw_data[:, my_slice]
+        assert my_raw_data.ndim == 2
+        my_dark_data = dark_data[:, my_slice]
+        assert my_dark_data.ndim == 2
+        my_flat_data = flat_data[:, my_slice]
+        assert my_flat_data.ndim == 2
+
+        n_projs = 3
+        nx_tomo = NXtomo()
+        nx_tomo.sample.x_translation = [0] * (n_projs + 2)
+        nx_tomo.sample.y_translation = [frame_y_position] * (n_projs + 2)
+        nx_tomo.sample.z_translation = [0] * (n_projs + 2)
+        nx_tomo.sample.rotation_angle = numpy.linspace(0, 180, num=(n_projs + 2), endpoint=False)
+        nx_tomo.instrument.detector.image_key_control = (
+            ImageKey.DARK_FIELD,
+            ImageKey.FLAT_FIELD,
+            ImageKey.PROJECTION,
+            ImageKey.PROJECTION,
+            ImageKey.PROJECTION,
+        )
+        nx_tomo.instrument.detector.x_pixel_size = 1.0
+        nx_tomo.instrument.detector.y_pixel_size = 1.0
+        nx_tomo.instrument.detector.distance = 2.3
+        nx_tomo.energy = 19.2
+        nx_tomo.instrument.detector.data = numpy.stack(
+            (
+                my_dark_data,
+                my_flat_data,
+                my_raw_data,
+                my_raw_data,
+                my_raw_data,
+            )
+        )
+
+        file_path = os.path.join(output_dir, f"nxtomo_{i_nxtomo}.nx")
+        entry = f"entry000{i_nxtomo}"
+        nx_tomo.save(file_path=file_path, data_path=entry)
+        scans.append(NXtomoScan(scan=file_path, entry=entry))
+    return scans, frame_y_positions, normalized_data
+
+
+@pytest.mark.parametrize("flip_lr", (True, False))
+@pytest.mark.parametrize("flip_ud", (True, False))
+@pytest.mark.parametrize("progress", (None, "with_tqdm"))
+def test_preprocessing_stitching(tmp_path, flip_lr, flip_ud, progress):
+    if progress == "with_tqdm":
+        progress = tqdm(total=100)
+
+    nxtomo_dir = tmp_path / "nxtomos"
+    nxtomo_dir.mkdir()
+    output_dir = tmp_path / "output"
+    output_dir.mkdir()
+
+    output_file_path = os.path.join(output_dir, "nxtomo.nxs")
+
+    nxtomos, _, normalized_data = build_nxtomos(
+        output_dir=nxtomo_dir,
+        flip_lr=flip_lr,
+        flip_ud=flip_ud,
+    )
+
+    configuration = PreProcessedYStitchingConfiguration(
+        input_scans=nxtomos,
+        axis_0_pos_px=None,
+        axis_1_pos_px=None,
+        axis_2_pos_px=None,
+        output_file_path=output_file_path,
+        output_data_path="stitched_volume",
+    )
+
+    output_identifier = y_stitching(
+        configuration=configuration,
+        progress=progress,
+    )
+    created_nx_tomo = NXtomo().load(
+        file_path=output_identifier.file_path,
+        data_path=output_identifier.data_path,
+        detector_data_as="as_numpy_array",
+    )
+    assert created_nx_tomo.instrument.detector.data.shape == (
+        3,
+        64,
+        120,
+    )  # 3 == number of projections, dark and flat will not be exported when doing the stitching
+    # TODO: improve me: the relative tolerance is pretty high. This doesn't comes from the algorithm on itself
+    # but more on the numerical calculation and the flat field normalization
+    numpy.testing.assert_allclose(normalized_data, created_nx_tomo.instrument.detector.data[0], rtol=0.06)