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

nabu/stitching/stitcher_2D.py

@@ -0,0 +1,278 @@
+import numpy
+from math import ceil
+from typing import Union, Optional
+from nabu.stitching.overlap import ImageStichOverlapKernel
+from nabu.stitching.frame_composition import FrameComposition
+from nabu.stitching.alignment import align_frame, _Alignment
+
+
+def stitch_raw_frames(
+    frames: tuple,
+    key_lines: tuple,
+    overlap_kernels: Union[ImageStichOverlapKernel, tuple],
+    output_dtype: numpy.dtype = numpy.float32,
+    check_inputs=True,
+    raw_frames_compositions: Optional[FrameComposition] = None,
+    overlap_frames_compositions: Optional[FrameComposition] = None,
+    return_composition_cls=False,
+    alignment: _Alignment = "center",
+    pad_mode="constant",
+    new_unstitched_axis_size: Optional[int] = None,
+) -> numpy.ndarray:
+    """
+    stitches raw frames (already shifted and flat fielded !!!) together using
+    raw stitching (no pixel interpolation, y_overlap_in_px is expected to be a int).
+    Sttiching depends on the kernel used.
+
+    It can be done:
+
+    * vertically:
+
+                                    X
+        ------------------------------------------------------------------>
+        |    --------------
+        |    |            |
+        |    |  Frame 1   |                           --------------
+        |    |            |                           |  Frame 1    |
+        |    --------------                           |             |
+        Y |                         --> stitching  -->  |~ stitching ~|
+        |    --------------                           |             |
+        |    |            |                           |  Frame 2    |
+        |    |  Frame 2   |                            --------------
+        |    |            |
+        |    --------------
+        \/
+
+    * horizontally:
+
+        ------------------------------------------------------------------>
+        |    --------------    --------------                     -----------------------
+        |    |            |    |            |                     |         ~ ~         |
+        Y |    |  Frame 1   |    |  Frame 2   | --> stitching  -->  | Frame 1 ~ ~ Frame 2 |
+        |    |            |    |            |                     |         ~ ~         |
+        |    --------------    --------------                     -----------------------
+        |
+        \/
+
+    returns stitched_projection, raw_img_1, raw_img_2, computed_overlap
+    proj_0 and pro_1 are already expected to be in a row. Having stitching_height_in_px in common. At top of proj_0
+    and at bottom of proj_1
+
+    :param tuple frames: tuple of 2D numpy array. Expected to be Z up oriented at this stage
+    :param tuple key_lines: for each jonction define the two lines to overlaid (from the upper and the lower frames). In the reference where 0 is the bottom line of the image.
+    :param overlap_kernels: ZStichOverlapKernel overlap kernel to be used or a list of kernel (one per overlap). Define startegy and overlap heights
+    :param numpy.dtype output_dtype: dataset dtype. For now must be provided because flat field corrcetion change data type (numpy.float32 for now)
+    :param bool check_inputs: if True will do more test on inputs parameters like checking frame shapes, coherence of the request.. As it can be time consuming it is optional
+    :param raw_frames_compositions: pre computed raw frame composition. If not provided will compute them. allow providing it to speed up calculation
+    :param overlap_frames_compositions: pre computed stitched frame composition. If not provided will compute them. allow providing it to speed up calculation
+    :param bool return_frame_compositions: if False return simply the stitched frames. Else return a tuple with stitching frame and the dictionnary with the composition frames...
+    :param alignment: how to align frame if two frames have different size along the unstitched axis
+    :param pad_mode: how to pad data for alignment (provided to numpy.pad function)
+    :param new_unstitched_axis_size: size of the image along the axis not stitched. So it will be the frame width if the stitching axis is 0 and the frame height if the stitching axis is 1
+    """
+    if overlap_kernels is None or len(overlap_kernels) == 0:
+        raise ValueError("overlap kernels must be provided")
+
+    stitched_axis = overlap_kernels[0].stitched_axis
+    unstitched_axis = overlap_kernels[0].unstitched_axis
+
+    if check_inputs:
+        # check frames are 2D numpy arrays
+        def check_frame(proj):
+            if not isinstance(proj, numpy.ndarray) and proj.ndim == 2:
+                raise ValueError(f"frames are expected to be 2D numpy array")
+
+        [check_frame(frame) for frame in frames]
+        for frame_0, frame_1 in zip(frames[:-1], frames[1:]):
+            if not (frame_0.ndim == frame_1.ndim == 2):
+                raise ValueError("Frames are expected to be 2D")
+
+        # check there is coherence between overlap kernels and frames
+        for frame_0, frame_1, kernel in zip(frames[:-1], frames[1:], overlap_kernels):
+            if frame_0.shape[stitched_axis] < kernel.overlap_size:
+                raise ValueError(
+                    f"frame_0 height ({frame_0.shape[stitched_axis]}) is less than kernel overlap ({kernel.overlap_size})"
+                )
+            if frame_1.shape[stitched_axis] < kernel.overlap_size:
+                raise ValueError(
+                    f"frame_1 height ({frame_1.shape[stitched_axis]}) is less than kernel overlap ({kernel.overlap_size})"
+                )
+        # check key lines are coherent with overlp kernels
+        if not len(key_lines) == len(overlap_kernels):
+            raise ValueError("we expect to have the same number of key_lines then the number of kernel")
+        else:
+            for key_line in key_lines:
+                for value in key_line:
+                    if not isinstance(value, (int, numpy.integer)):
+                        raise TypeError(f"key_line is expected to be an integer. {type(key_line)} provided")
+                    elif value < 0:
+                        raise ValueError(f"key lines are expected to be positive values. Get {value} as key line value")
+
+        # check overlap kernel stitching axis are coherent (for now make sure they are all along the same axis)
+        if len(overlap_kernels) > 1:
+            for previous_kernel, next_kernel in zip(overlap_kernels[:-1], overlap_kernels[1:]):
+                if not isinstance(previous_kernel, ImageStichOverlapKernel):
+                    raise TypeError(
+                        f"overlap kernels must be instances of {ImageStichOverlapKernel}. Get {type(previous_kernel)}"
+                    )
+                if not isinstance(next_kernel, ImageStichOverlapKernel):
+                    raise TypeError(
+                        f"overlap kernels must be instances of {ImageStichOverlapKernel}. Get {type(next_kernel)}"
+                    )
+                if previous_kernel.stitched_axis != next_kernel.stitched_axis:
+                    raise ValueError(
+                        "kernels with different stitching axis provided. For now all kernels must have the same stitchign axis"
+                    )
+
+    if new_unstitched_axis_size is None:
+        new_unstitched_axis_size = max([frame.shape[unstitched_axis] for frame in frames])
+
+    frames = tuple(
+        [
+            align_frame(
+                data=frame,
+                alignment=alignment,
+                new_aligned_axis_size=new_unstitched_axis_size,
+                pad_mode=pad_mode,
+                alignment_axis=unstitched_axis,
+            )
+            for frame in frames
+        ]
+    )
+
+    # step 1: create numpy array that will contain stitching
+    # if raw composition doesn't exists create it
+    if raw_frames_compositions is None:
+        raw_frames_compositions = FrameComposition.compute_raw_frame_compositions(
+            frames=frames,
+            overlap_kernels=overlap_kernels,
+            key_lines=key_lines,
+            stitching_axis=stitched_axis,
+        )
+
+    new_stitched_axis_size = raw_frames_compositions.global_end[-1] - raw_frames_compositions.global_start[0]
+    if stitched_axis == 0:
+        stitched_projection_shape = (
+            int(new_stitched_axis_size),
+            new_unstitched_axis_size,
+        )
+    else:
+        stitched_projection_shape = (
+            new_unstitched_axis_size,
+            int(new_stitched_axis_size),
+        )
+
+    stitch_array = numpy.empty(stitched_projection_shape, dtype=output_dtype)
+
+    # step 2: set raw data
+    # fill stitch array with raw data raw data
+    raw_frames_compositions.compose(
+        output_frame=stitch_array,
+        input_frames=frames,
+    )
+
+    # step 3 set stitched data
+
+    # 3.1 create stitched overlaps
+    stitched_overlap = []
+    for frame_0, frame_1, kernel, key_line in zip(frames[:-1], frames[1:], overlap_kernels, key_lines):
+        assert kernel.overlap_size >= 0
+        frame_0_overlap, frame_1_overlap = get_overlap_areas(
+            upper_frame=frame_0,
+            lower_frame=frame_1,
+            upper_frame_key_line=key_line[0],
+            lower_frame_key_line=key_line[1],
+            overlap_size=kernel.overlap_size,
+            stitching_axis=stitched_axis,
+        )
+
+        assert (
+            frame_0_overlap.shape[stitched_axis] == frame_1_overlap.shape[stitched_axis] == kernel.overlap_size
+        ), f"{frame_0_overlap.shape[stitched_axis]} == {frame_1_overlap.shape[stitched_axis]} == {kernel.overlap_size}"
+
+        stitched_overlap.append(
+            kernel.stitch(
+                frame_0_overlap,
+                frame_1_overlap,
+            )[0]
+        )
+    # 3.2 fill stitched overlap on output array
+    if overlap_frames_compositions is None:
+        overlap_frames_compositions = FrameComposition.compute_stitch_frame_composition(
+            frames=frames,
+            overlap_kernels=overlap_kernels,
+            key_lines=key_lines,
+            stitching_axis=stitched_axis,
+        )
+    overlap_frames_compositions.compose(
+        output_frame=stitch_array,
+        input_frames=stitched_overlap,
+    )
+    if return_composition_cls:
+        return (
+            stitch_array,
+            {
+                "raw_composition": raw_frames_compositions,
+                "overlap_composition": overlap_frames_compositions,
+            },
+        )
+
+    return stitch_array
+
+
+def get_overlap_areas(
+    upper_frame: numpy.ndarray,
+    lower_frame: numpy.ndarray,
+    upper_frame_key_line: int,
+    lower_frame_key_line: int,
+    overlap_size: int,
+    stitching_axis: int,
+):
+    """
+    return the requested area from lower_frame and upper_frame.
+
+    Lower_frame contains at the end of it the 'real overlap' with the upper_frame.
+    Upper_frame contains the 'real overlap' at the end of it.
+
+    For some reason the user can ask the stitching height to be smaller than the `real overlap`.
+
+    Here are some drawing to have a better of view of those regions:
+
+    .. image:: images/stitching/z_stitch_real_overlap.png
+        :width: 600
+
+    .. image:: z_stitch_stitch_height.png
+        :width: 600
+    """
+    assert stitching_axis in (0, 1)
+    for pf, pn in zip((lower_frame_key_line, upper_frame_key_line), ("lower_frame", "upper_frame")):
+        if not isinstance(pf, (int, numpy.number)):
+            raise TypeError(f"{pn} is expected to be a number. {type(pf)} provided")
+    assert overlap_size >= 0
+
+    lf_start = ceil(lower_frame_key_line - overlap_size / 2)
+    lf_end = ceil(lower_frame_key_line + overlap_size / 2)
+    uf_start = ceil(upper_frame_key_line - overlap_size / 2)
+    uf_end = ceil(upper_frame_key_line + overlap_size / 2)
+
+    lf_start, lf_end = min(lf_start, lf_end), max(lf_start, lf_end)
+    uf_start, uf_end = min(uf_start, uf_end), max(uf_start, uf_end)
+    if lf_start < 0 or uf_start < 0:
+        raise ValueError(
+            f"requested overlap ({overlap_size}) is incoherent with key line positions ({lower_frame_key_line}, {upper_frame_key_line}) - expected to be smaller."
+        )
+
+    if stitching_axis == 0:
+        overlap_upper = upper_frame[uf_start:uf_end]
+        overlap_lower = lower_frame[lf_start:lf_end]
+    elif stitching_axis == 1:
+        overlap_upper = upper_frame[:, uf_start:uf_end]
+        overlap_lower = lower_frame[:, lf_start:lf_end]
+    else:
+        raise NotImplementedError
+    if not overlap_upper.shape == overlap_lower.shape:
+        # maybe in the future: try to reduce one according to the other ????
+        raise RuntimeError(
+            f"lower and upper frame have different overlap size ({overlap_upper.shape} vs {overlap_lower.shape})"
+        )
+    return overlap_upper, overlap_lower

nabu/stitching/tests/test_config.py

@@ -1,33 +1,3 @@
-# coding: utf-8
-# /*##########################################################################
-#
-# Copyright (c) 2016-2017 European Synchrotron Radiation Facility
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in
-# all copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-# THE SOFTWARE.
-#
-# ###########################################################################*/
-
-__authors__ = ["H. Payno"]
-__license__ = "MIT"
-__date__ = "10/05/2022"
-
-
 import os
 from tempfile import TemporaryDirectory
 import pytest
@@ -71,12 +41,17 @@ def test_stitching_config(stitching_type, option_level):
         if stitching_type is stiching_config.StitchingType.Z_POSTPROC:
             assert isinstance(
                 stiching_config.dict_to_config_obj(config),
-                stiching_config.PostProcessedZStitchingConfiguration,
+                stiching_config.PostProcessedSingleAxisStitchingConfiguration,
             )
         elif stitching_type is stiching_config.StitchingType.Z_PREPROC:
             assert isinstance(
                 stiching_config.dict_to_config_obj(config),
-                stiching_config.PreProcessedZStitchingConfiguration,
+                stiching_config.PreProcessedSingleAxisStitchingConfiguration,
+            )
+        elif stitching_type is stiching_config.StitchingType.Y_PREPROC:
+            assert isinstance(
+                stiching_config.dict_to_config_obj(config),
+                stiching_config.PreProcessedSingleAxisStitchingConfiguration,
             )
         else:
             raise ValueError("not handled")
@@ -98,12 +73,12 @@ def test_stitching_config(stitching_type, option_level):
         if stitching_type is stiching_config.StitchingType.Z_POSTPROC:
             assert isinstance(
                 config_class_instance,
-                stiching_config.PostProcessedZStitchingConfiguration,
+                stiching_config.PostProcessedSingleAxisStitchingConfiguration,
             )
         elif stitching_type is stiching_config.StitchingType.Z_PREPROC:
             assert isinstance(
                 config_class_instance,
-                stiching_config.PreProcessedZStitchingConfiguration,
+                stiching_config.PreProcessedSingleAxisStitchingConfiguration,
             )
 
         assert isinstance(config_class_instance.to_dict(), dict)
@@ -183,7 +158,7 @@ def test_PreProcessedZStitchingConfiguration(
             stiching_config.OUTPUT_SECTION: {
                 stiching_config.OVERWRITE_RESULTS_FIELD: overwrite_results,
             },
-            stiching_config.Z_PRE_PROC_SECTION: {
+            stiching_config.PRE_PROC_SECTION: {
                 stiching_config.DATA_FILE_FIELD: "my_file.nx",
                 stiching_config.DATA_PATH_FIELD: "entry",
                 stiching_config.NEXUS_VERSION_FIELD: None,
@@ -219,8 +194,8 @@ def test_description_dict():
     """
     make sure the description dict (used for generating the file) is working and generates a dict
     """
-    assert isinstance(stiching_config.PreProcessedZStitchingConfiguration.get_description_dict(), dict)
+    assert isinstance(stiching_config.PreProcessedSingleAxisStitchingConfiguration.get_description_dict(), dict)
     assert isinstance(
-        stiching_config.PostProcessedZStitchingConfiguration.get_description_dict(),
+        stiching_config.PostProcessedSingleAxisStitchingConfiguration.get_description_dict(),
         dict,
     )

nabu/stitching/tests/test_frame_composition.py

@@ -1,43 +1,13 @@
-# coding: utf-8
-# /*##########################################################################
-#
-# Copyright (c) 2016-2017 European Synchrotron Radiation Facility
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in
-# all copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-# THE SOFTWARE.
-#
-# ###########################################################################*/
-
-__authors__ = ["H. Payno"]
-__license__ = "MIT"
-__date__ = "17/05/2022"
-
-
 import pytest
-from nabu.stitching.frame_composition import ZFrameComposition
+from nabu.stitching.frame_composition import FrameComposition
 import numpy
 
-from nabu.stitching.overlap import OverlapStitchingStrategy, ZStichOverlapKernel
+from nabu.stitching.overlap import OverlapStitchingStrategy, ImageStichOverlapKernel
 
 
 def test_frame_composition():
     """
-    Test _ZFrameComposition
+    Test FrameComposition
     """
     frame_0 = numpy.zeros((100, 1))
     frame_1 = numpy.ones((98, 1))
@@ -48,19 +18,21 @@ def test_frame_composition():
     y_shifts = -20, -10
 
     kernels = [
-        ZStichOverlapKernel(
-            frame_width=1,
+        ImageStichOverlapKernel(
+            stitching_axis=0,
+            frame_unstitched_axis_size=1,
             stitching_strategy=OverlapStitchingStrategy.LINEAR_WEIGHTS,
             overlap_size=4,
         ),
-        ZStichOverlapKernel(
-            frame_width=1,
+        ImageStichOverlapKernel(
+            stitching_axis=0,
+            frame_unstitched_axis_size=1,
             stitching_strategy=OverlapStitchingStrategy.MEAN,
             overlap_size=8,
         ),
     ]
     # check raw composition
-    raw_composition = ZFrameComposition.compute_raw_frame_compositions(
+    raw_composition = FrameComposition.compute_raw_frame_compositions(
         frames=frames,
         key_lines=(
             (90, 10),
@@ -69,12 +41,12 @@ def test_frame_composition():
         overlap_kernels=kernels,
         stitching_axis=0,
     )
-    assert isinstance(raw_composition, ZFrameComposition)
+    assert isinstance(raw_composition, FrameComposition)
 
-    assert raw_composition.local_start_y == (0, 12, 9)
-    assert raw_composition.local_end_y == (88, 89, 205)
-    assert raw_composition.global_start_y == (0, 92, 177)
-    assert raw_composition.global_end_y == (88, 169, 373)
+    assert raw_composition.local_start == (0, 12, 9)
+    assert raw_composition.local_end == (88, 89, 205)
+    assert raw_composition.global_start == (0, 92, 177)
+    assert raw_composition.global_end == (88, 169, 373)
 
     stitched_data = numpy.empty((100 + 98 + 205 - 30, 1))
     raw_composition.compose(output_frame=stitched_data, input_frames=frames)
@@ -83,7 +55,7 @@ def test_frame_composition():
     assert stitched_data[-1, 0] == 2.0
 
     # check stitch composition
-    stitch_composition = ZFrameComposition.compute_stitch_frame_composition(
+    stitch_composition = FrameComposition.compute_stitch_frame_composition(
         frames=frames,
         key_lines=(
             (90, 10),
@@ -93,12 +65,12 @@ def test_frame_composition():
         stitching_axis=0,
     )
 
-    ZFrameComposition.pprint_z_stitching(raw_composition, stitch_composition)
+    FrameComposition.pprint_composition(raw_composition, stitch_composition)
 
-    assert stitch_composition.local_start_y == (0, 0)
-    assert stitch_composition.local_end_y == (4, 8)
-    assert stitch_composition.global_start_y == (88, 169)
-    assert stitch_composition.global_end_y == (92, 177)
+    assert stitch_composition.local_start == (0, 0)
+    assert stitch_composition.local_end == (4, 8)
+    assert stitch_composition.global_start == (88, 169)
+    assert stitch_composition.global_end == (92, 177)
 
     stitched_frames = []
     for frame_0, frame_1, kernel, y_shift in zip(frames[:-1], frames[1:], kernels, y_shifts):
@@ -130,13 +102,15 @@ _raw_comp_config = (
         "raw_local_start": (0, 4, 5),
         "raw_local_end": (16, 34, 20),
         "kernels": (
-            ZStichOverlapKernel(
-                frame_width=1,
+            ImageStichOverlapKernel(
+                stitching_axis=0,
+                frame_unstitched_axis_size=1,
                 stitching_strategy=OverlapStitchingStrategy.LINEAR_WEIGHTS,
                 overlap_size=3,
             ),
-            ZStichOverlapKernel(
-                frame_width=1,
+            ImageStichOverlapKernel(
+                stitching_axis=0,
+                frame_unstitched_axis_size=1,
                 stitching_strategy=OverlapStitchingStrategy.MEAN,
                 overlap_size=4,
             ),
@@ -159,21 +133,21 @@ def test_raw_frame_composition_exotic_config(configuration):
     key_lines = configuration.get("key_lines")
     kernels = configuration.get("kernels")
     # check raw composition
-    raw_composition = ZFrameComposition.compute_raw_frame_compositions(
+    raw_composition = FrameComposition.compute_raw_frame_compositions(
         frames=frames,
         overlap_kernels=kernels,
         key_lines=key_lines,
         stitching_axis=0,
     )
 
-    assert raw_composition.global_start_y == configuration.get("raw_global_start")
-    assert raw_composition.global_end_y == configuration.get("raw_global_end")
-    assert raw_composition.local_start_y == configuration.get("raw_local_start")
-    assert raw_composition.local_end_y == configuration.get("raw_local_end")
+    assert raw_composition.global_start == configuration.get("raw_global_start")
+    assert raw_composition.global_end == configuration.get("raw_global_end")
+    assert raw_composition.local_start == configuration.get("raw_local_start")
+    assert raw_composition.local_end == configuration.get("raw_local_end")
 
     stitched_data = numpy.empty(
         (
-            (raw_composition.global_end_y[-1] - raw_composition.global_start_y[0]),
+            (raw_composition.global_end[-1] - raw_composition.global_start[0]),
             1,
         )
     )