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

nabu/stitching/stitcher/post_processing.py

@@ -0,0 +1,555 @@
+import logging
+import numpy
+import os
+import h5py
+from typing import Union
+from nabu.stitching.config import PostProcessedSingleAxisStitchingConfiguration
+from nabu.stitching.alignment import AlignmentAxis1
+from nabu.stitching.alignment import PaddedRawData
+from math import ceil
+from tomoscan.io import HDF5File
+from tomoscan.esrf.scan.utils import cwd_context
+from tomoscan.esrf import NXtomoScan
+from tomoscan.series import Series
+from tomoscan.volumebase import VolumeBase
+from tomoscan.esrf.volume import HDF5Volume
+from typing import Iterable
+from contextlib import AbstractContextManager
+from pyunitsystem.metricsystem import MetricSystem
+from nabu.stitching.config import (
+    PostProcessedSingleAxisStitchingConfiguration,
+    KEY_IMG_REG_METHOD,
+)
+from nabu.stitching.utils.utils import find_volumes_relative_shifts
+from nabu.io.utils import DatasetReader
+from .single_axis import SingleAxisStitcher
+
+_logger = logging.getLogger(__name__)
+
+
+class FlippingValueError(ValueError):
+    pass
+
+
+class PostProcessingStitching(SingleAxisStitcher):
+    """
+    Loader to be used when load data during post-processing stitching (on recosntructed volume). Output is expected to be an NXtomo
+    """
+
+    def __init__(self, configuration, progress=None) -> None:
+        if not isinstance(configuration, PostProcessedSingleAxisStitchingConfiguration):
+            raise TypeError(
+                f"configuration is expected to be an instance of {PostProcessedSingleAxisStitchingConfiguration}. Get {type(configuration)} instead"
+            )
+        self._input_volumes = configuration.input_volumes
+        self.__output_data_type = None
+
+        self._series = Series("series", iterable=self._input_volumes, use_identifiers=False)
+
+        super().__init__(configuration, progress=progress)
+
+    @property
+    def stitching_axis_in_frame_space(self):
+        if self.axis == 0:
+            return 0
+        elif self.axis in (1, 2):
+            raise NotImplementedError(f"post-processing stitching along axis {self.axis} is not handled.")
+        else:
+            raise NotImplementedError(f"stitching axis must be in (0, 1, 2). Get {self.axis}")
+
+    def settle_flips(self):
+        super().settle_flips()
+        if not self.configuration.duplicate_data:
+            if len(numpy.unique(self.configuration.flip_lr)) > 1:
+                raise FlippingValueError(
+                    "Stitching without data duplication cannot handle volume with different flip. Please run the stitching with data duplication"
+                )
+            if True in self.configuration.flip_ud:
+                raise FlippingValueError(
+                    "Stitching without data duplication cannot handle with up / down flips. Please run the stitching with data duplication"
+                )
+
+    def order_input_tomo_objects(self):
+
+        def get_min_bound(volume):
+            try:
+                bb = volume.get_bounding_box(axis=self.axis)
+            except ValueError:  #  if missing information
+                bb = None
+            if bb is not None:
+                return bb.min
+            else:
+                # if can't find bounding box (missing metadata to the volume
+                # try to get it from the scan
+                metadata = volume.metadata or volume.load_metadata()
+                scan_location = metadata.get("nabu_config", {}).get("dataset", {}).get("location", None)
+                scan_entry = metadata.get("nabu_config", {}).get("dataset", {}).get("hdf5_entry", None)
+                if scan_location is not None:
+                    # this work around (until most volume have position metadata) works only for Hdf5volume
+                    with cwd_context(os.path.dirname(volume.file_path)):
+                        o_scan = NXtomoScan(scan_location, scan_entry)
+                        bb_acqui = o_scan.get_bounding_box(axis=None)
+                        # for next step volume position will be required.
+                        # if you can find it set it directly
+                        volume.position = (numpy.array(bb_acqui.max) - numpy.array(bb_acqui.min)) / 2.0 + numpy.array(
+                            bb_acqui.min
+                        )
+                        # for now translation are stored in pixel size ref instead of real_pixel_size
+                        volume.pixel_size = o_scan.x_real_pixel_size
+                        if bb_acqui is not None:
+                            return bb_acqui.min[0]
+                raise ValueError("Unable to find volume position. Unable to deduce z position")
+
+        try:
+            # order volumes from higher z to lower z
+            # if axis 0 position is provided then use directly it
+            if self.configuration.axis_0_pos_px is not None and len(self.configuration.axis_0_pos_px) > 0:
+                order = numpy.argsort(self.configuration.axis_0_pos_px)
+                sorted_series = Series(
+                    self.series.name,
+                    numpy.take_along_axis(numpy.array(self.series[:]), order, axis=0)[::-1],
+                    use_identifiers=False,
+                )
+            else:
+                # else use bounding box
+                sorted_series = Series(
+                    self.series.name,
+                    sorted(self.series[:], key=get_min_bound, reverse=True),
+                    use_identifiers=False,
+                )
+        except ValueError:
+            _logger.warning(
+                "Unable to find volume positions in metadata. Expect the volume to be ordered already (decreasing along axis 0.)"
+            )
+        else:
+            if sorted_series == self.series:
+                pass
+            elif sorted_series != self.series:
+                if sorted_series[:] != self.series[::-1]:
+                    raise ValueError(
+                        "Unable to get comprehensive input. ordering along axis 0 is not respected (decreasing)."
+                    )
+                else:
+                    _logger.warning(
+                        f"decreasing order haven't been respected. Need to reorder {self.serie_label} ({[str(scan) for scan in sorted_series[:]]}). Will also reorder positions"
+                    )
+                    if self.configuration.axis_0_pos_mm is not None:
+                        self.configuration.axis_0_pos_mm = self.configuration.axis_0_pos_mm[::-1]
+                    if self.configuration.axis_0_pos_px is not None:
+                        self.configuration.axis_0_pos_px = self.configuration.axis_0_pos_px[::-1]
+                    if self.configuration.axis_1_pos_mm is not None:
+                        self.configuration.axis_1_pos_mm = self.configuration.axis_1_pos_mm[::-1]
+                    if self.configuration.axis_1_pos_px is not None:
+                        self.configuration.axis_1_pos_px = self.configuration.axis_1_pos_px[::-1]
+                    if self.configuration.axis_2_pos_mm is not None:
+                        self.configuration.axis_2_pos_mm = self.configuration.axis_2_pos_mm[::-1]
+                    if self.configuration.axis_2_pos_px is not None:
+                        self.configuration.axis_2_pos_px = self.configuration.axis_2_pos_px[::-1]
+                    if not numpy.isscalar(self._configuration.flip_ud):
+                        self._configuration.flip_ud = self._configuration.flip_ud[::-1]
+                    if not numpy.isscalar(self._configuration.flip_lr):
+                        self._configuration.flip_ud = self._configuration.flip_lr[::-1]
+
+                    self._series = sorted_series
+
+    def check_inputs(self):
+        """
+        insure input data is coherent
+        """
+        # check input volume
+        if self.configuration.output_volume is None:
+            raise ValueError("input volume should be provided")
+
+        n_volumes = len(self.series)
+        if n_volumes == 0:
+            raise ValueError("no scan to stich together")
+
+        if not isinstance(self.configuration.output_volume, VolumeBase):
+            raise TypeError(f"make sure we return a volume identifier not {(type(self.configuration.output_volume))}")
+
+        # check axis 0 position
+        if isinstance(self.configuration.axis_0_pos_px, Iterable) and len(self.configuration.axis_0_pos_px) != (
+            n_volumes
+        ):
+            raise ValueError(f"expect {n_volumes} overlap defined. Get {len(self.configuration.axis_0_pos_px)}")
+        if isinstance(self.configuration.axis_0_pos_mm, Iterable) and len(self.configuration.axis_0_pos_mm) != (
+            n_volumes
+        ):
+            raise ValueError(f"expect {n_volumes} overlap defined. Get {len(self.configuration.axis_0_pos_mm)}")
+
+        # check axis 1 position
+        if isinstance(self.configuration.axis_1_pos_px, Iterable) and len(self.configuration.axis_1_pos_px) != (
+            n_volumes
+        ):
+            raise ValueError(f"expect {n_volumes} overlap defined. Get {len(self.configuration.axis_1_pos_px)}")
+        if isinstance(self.configuration.axis_1_pos_mm, Iterable) and len(self.configuration.axis_1_pos_mm) != (
+            n_volumes
+        ):
+            raise ValueError(f"expect {n_volumes} overlap defined. Get {len(self.configuration.axis_1_pos_mm)}")
+
+        # check axis 2 position
+        if isinstance(self.configuration.axis_1_pos_px, Iterable) and len(self.configuration.axis_1_pos_px) != (
+            n_volumes
+        ):
+            raise ValueError(f"expect {n_volumes} overlap defined. Get {len(self.configuration.axis_1_pos_px)}")
+        if isinstance(self.configuration.axis_2_pos_mm, Iterable) and len(self.configuration.axis_2_pos_mm) != (
+            n_volumes
+        ):
+            raise ValueError(f"expect {n_volumes} overlap defined. Get {len(self.configuration.axis_2_pos_mm)}")
+
+        self._reading_orders = []
+        # the first scan will define the expected reading orderd, and expected flip.
+        # if all scan are flipped then we will keep it this way
+        self._reading_orders.append(1)
+
+    @staticmethod
+    def _get_bunch_of_data(
+        bunch_start: int,
+        bunch_end: int,
+        step: int,
+        volumes: tuple,
+        flip_lr_arr: bool,
+        flip_ud_arr: bool,
+    ):
+        """
+        goal is to load contiguous frames as much as possible...
+        return for each volume the bunch of slice along axis 1
+        warning: they can have different shapes
+        """
+
+        def get_sub_volume(volume, flip_lr, flip_ud):
+            sub_volume = volume[:, bunch_start:bunch_end:step, :]
+            if flip_lr:
+                sub_volume = numpy.fliplr(sub_volume)
+            if flip_ud:
+                sub_volume = numpy.flipud(sub_volume)
+            return sub_volume
+
+        sub_volumes = [
+            get_sub_volume(volume, flip_lr, flip_ud)
+            for volume, flip_lr, flip_ud in zip(volumes, flip_lr_arr, flip_ud_arr)
+        ]
+        # generator on it self: we want to iterate over the y axis
+        n_slices_in_bunch = ceil((bunch_end - bunch_start) / step)
+        assert isinstance(n_slices_in_bunch, int)
+        for i in range(n_slices_in_bunch):
+            yield [sub_volume[:, i, :] for sub_volume in sub_volumes]
+
+    def compute_estimated_shifts(self):
+        axis_0_pos_px = self.configuration.axis_0_pos_px
+        self._axis_0_rel_ini_shifts = []
+        # compute overlap along axis 0
+        for upper_volume, lower_volume, upper_volume_axis_0_pos, lower_volume_axis_0_pos in zip(
+            self.series[:-1], self.series[1:], axis_0_pos_px[:-1], axis_0_pos_px[1:]
+        ):
+            upper_volume_low_pos = upper_volume_axis_0_pos - upper_volume.get_volume_shape()[0] / 2
+            lower_volume_high_pos = lower_volume_axis_0_pos + lower_volume.get_volume_shape()[0] / 2
+            self._axis_0_rel_ini_shifts.append(
+                int(lower_volume_high_pos - upper_volume_low_pos)  # overlap are expected to be int for now
+            )
+        self._axis_1_rel_ini_shifts = self.from_abs_pos_to_rel_pos(self.configuration.axis_1_pos_px)
+        self._axis_2_rel_ini_shifts = [0.0] * (len(self.series) - 1)
+
+    def _compute_positions_as_px(self):
+        """compute if necessary position other axis 0 from volume metadata"""
+
+        def get_position_as_px_on_axis(axis, pos_as_px, pos_as_mm):
+            if pos_as_px is not None:
+                if pos_as_mm is not None:
+                    raise ValueError(
+                        f"position of axis {axis} is provided twice: as mm and as px. Please provide one only ({pos_as_mm} vs {pos_as_px})"
+                    )
+                else:
+                    return pos_as_px
+
+            elif pos_as_mm is not None:
+                # deduce from position given in configuration and pixel size
+                axis_N_pos_px = []
+                for volume, pos_in_mm in zip(self.series, pos_as_mm):
+                    voxel_size_m = self.configuration.voxel_size or volume.voxel_size
+                    axis_N_pos_px.append((pos_in_mm / MetricSystem.MILLIMETER.value) / voxel_size_m[0])
+                return axis_N_pos_px
+            else:
+                # deduce from motor position and pixel size
+                axis_N_pos_px = []
+                base_position_m = self.series[0].get_bounding_box(axis=axis).min
+                for volume in self.series:
+                    voxel_size_m = self.configuration.voxel_size or volume.voxel_size
+                    volume_axis_bb = volume.get_bounding_box(axis=axis)
+                    axis_N_mean_pos_m = (volume_axis_bb.max - volume_axis_bb.min) / 2 + volume_axis_bb.min
+                    axis_N_mean_rel_pos_m = axis_N_mean_pos_m - base_position_m
+                    axis_N_pos_px.append(int(axis_N_mean_rel_pos_m / voxel_size_m[0]))
+                return axis_N_pos_px
+
+        self.configuration.axis_0_pos_px = get_position_as_px_on_axis(
+            axis=0,
+            pos_as_px=self.configuration.axis_0_pos_px,
+            pos_as_mm=self.configuration.axis_0_pos_mm,
+        )
+        self.configuration.axis_0_pos_mm = None
+
+        self.configuration.axis_1_pos_px = get_position_as_px_on_axis(
+            axis=1,
+            pos_as_px=self.configuration.axis_1_pos_px,
+            pos_as_mm=self.configuration.axis_1_pos_mm,
+        )
+
+        self.configuration.axis_2_pos_px = get_position_as_px_on_axis(
+            axis=2,
+            pos_as_px=self.configuration.axis_2_pos_px,
+            pos_as_mm=self.configuration.axis_2_pos_mm,
+        )
+        self.configuration.axis_2_pos_mm = None
+
+    def _compute_shifts(self):
+        n_volumes = len(self.configuration.input_volumes)
+        if n_volumes == 0:
+            raise ValueError("no scan to stich provided")
+
+        slice_for_shift = self.configuration.slice_for_cross_correlation or "middle"
+        y_rel_shifts = self._axis_0_rel_ini_shifts
+        x_rel_shifts = self._axis_1_rel_ini_shifts
+        dim_axis_1 = max([volume.get_volume_shape()[1] for volume in self.series])
+
+        final_rel_shifts = []
+        for (
+            upper_volume,
+            lower_volume,
+            x_rel_shift,
+            y_rel_shift,
+            flip_ud_upper,
+            flip_ud_lower,
+        ) in zip(
+            self.series[:-1],
+            self.series[1:],
+            x_rel_shifts,
+            y_rel_shifts,
+            self.configuration.flip_ud[:-1],
+            self.configuration.flip_ud[1:],
+        ):
+            x_cross_algo = self.configuration.axis_2_params.get(KEY_IMG_REG_METHOD, None)
+            y_cross_algo = self.configuration.axis_0_params.get(KEY_IMG_REG_METHOD, None)
+
+            # compute relative shift
+            found_shift_y, found_shift_x = find_volumes_relative_shifts(
+                upper_volume=upper_volume,
+                lower_volume=lower_volume,
+                dtype=self.get_output_data_type(),
+                dim_axis_1=dim_axis_1,
+                slice_for_shift=slice_for_shift,
+                x_cross_correlation_function=x_cross_algo,
+                y_cross_correlation_function=y_cross_algo,
+                x_shifts_params=self.configuration.axis_2_params,
+                y_shifts_params=self.configuration.axis_0_params,
+                estimated_shifts=(y_rel_shift, x_rel_shift),
+                flip_ud_lower_frame=flip_ud_lower,
+                flip_ud_upper_frame=flip_ud_upper,
+                alignment_axis_1=self.configuration.alignment_axis_1,
+                alignment_axis_2=self.configuration.alignment_axis_2,
+                overlap_axis=self.axis,
+            )
+            final_rel_shifts.append(
+                (found_shift_y, found_shift_x),
+            )
+
+        # set back values. Now position should start at 0
+        self._axis_0_rel_final_shifts = [final_shift[0] for final_shift in final_rel_shifts]
+        self._axis_1_rel_final_shifts = [final_shift[1] for final_shift in final_rel_shifts]
+        self._axis_2_rel_final_shifts = [0.0] * len(final_rel_shifts)
+        _logger.info(f"axis 2 relative shifts (x in radio ref) to be used will be {self._axis_1_rel_final_shifts}")
+        print(f"axis 2 relative shifts (x in radio ref) to be used will be {self._axis_1_rel_final_shifts}")
+        _logger.info(f"axis 0 relative shifts (y in radio ref) y to be used will be {self._axis_0_rel_final_shifts}")
+        print(f"axis 0 relative shifts (y in radio ref) y to be used will be {self._axis_0_rel_final_shifts}")
+
+    def get_output_data_type(self):
+        if self.__output_data_type is None:
+
+            def find_output_data_type():
+                first_vol = self._input_volumes[0]
+                if first_vol.data is not None:
+                    return first_vol.data.dtype
+                elif isinstance(first_vol, HDF5Volume):
+                    with DatasetReader(first_vol.data_url) as vol_dataset:
+                        return vol_dataset.dtype
+                else:
+                    return first_vol.load_data(store=False).dtype
+
+            self.__output_data_type = find_output_data_type()
+        return self.__output_data_type
+
+    def _create_stitched_volume(self, store_composition: bool):
+        overlap_kernels = self._overlap_kernels
+        self._slices_to_stitch, n_slices = self.configuration.settle_slices()
+
+        # sync overwrite_results with volume overwrite parameter
+        self.configuration.output_volume.overwrite = self.configuration.overwrite_results
+
+        # init final volume
+        final_volume = self.configuration.output_volume
+        final_volume_shape = (
+            int(
+                numpy.asarray([volume.get_volume_shape()[0] for volume in self._input_volumes]).sum()
+                - numpy.asarray([abs(overlap) for overlap in self._axis_0_rel_final_shifts]).sum(),
+            ),
+            n_slices,
+            self._stitching_constant_length,
+        )
+
+        data_type = self.get_output_data_type()
+
+        if self.progress:
+            self.progress.total = final_volume_shape[1]
+
+        y_index = 0
+        if isinstance(self._slices_to_stitch, slice):
+            step = self._slices_to_stitch.step or 1
+        else:
+            step = 1
+
+        output_dataset_args = {
+            "volume": final_volume,
+            "volume_shape": final_volume_shape,
+            "dtype": data_type,
+            "dumper": self.dumper,
+        }
+        from .dumper.postprocessing import PostProcessingStitchingDumperNoDD
+
+        # TODO: FIXME: for now not very elegant but in the case of avoiding data duplication
+        # we need to provide the the information about the stitched part shape.
+        # this should be move to the dumper in the future
+        if isinstance(self.dumper, PostProcessingStitchingDumperNoDD):
+            output_dataset_args["stitching_sources_arr_shapes"] = tuple(
+                [(abs(overlap), n_slices, self._stitching_constant_length) for overlap in self._axis_0_rel_final_shifts]
+            )
+
+        with self.dumper.OutputDatasetContext(**output_dataset_args):
+            # note: output_dataset is a HDF5 dataset if final volume is an HDF5 volume else is a numpy array
+            with _RawDatasetsContext(
+                self._input_volumes,
+                alignment_axis_1=self.configuration.alignment_axis_1,
+            ) as raw_datasets:
+                # note: raw_datasets can be numpy arrays or HDF5 dataset (in the case of HDF5Volume)
+                # to speed up we read by bunch of dataset. For numpy array this doesn't change anything
+                # but for HDF5 dataset this can speed up a lot the processing (depending on HDF5 dataset chuncks)
+                # note: we read trhough axis 1
+                if isinstance(self.dumper, PostProcessingStitchingDumperNoDD):
+                    self.dumper.raw_regions_hdf5_dataset = raw_datasets
+                for bunch_start, bunch_end in PostProcessingStitching._data_bunch_iterator(
+                    slices=self._slices_to_stitch, bunch_size=50
+                ):
+                    for data_frames in PostProcessingStitching._get_bunch_of_data(
+                        bunch_start,
+                        bunch_end,
+                        step=step,
+                        volumes=raw_datasets,
+                        flip_lr_arr=self.configuration.flip_lr,
+                        flip_ud_arr=self.configuration.flip_ud,
+                    ):
+                        if self.configuration.rescale_frames:
+                            data_frames = self.rescale_frames(data_frames)
+                        if self.configuration.normalization_by_sample.is_active():
+                            data_frames = self.normalize_frame_by_sample(data_frames)
+
+                        sf = PostProcessingStitching.stitch_frames(
+                            frames=data_frames,
+                            axis=self.axis,
+                            output_dtype=data_type,
+                            x_relative_shifts=self._axis_1_rel_final_shifts,
+                            y_relative_shifts=self._axis_0_rel_final_shifts,
+                            overlap_kernels=overlap_kernels,
+                            dumper=self.dumper,
+                            i_frame=y_index,
+                            return_composition_cls=store_composition if y_index == 0 else False,
+                            stitching_axis=self.axis,
+                            check_inputs=y_index == 0,  # on process check on the first iteration
+                        )
+                        if y_index == 0 and store_composition:
+                            _, self._frame_composition = sf
+
+                        if self.progress is not None:
+                            self.progress.update()
+                        y_index += 1
+
+    # alias to general API
+    def _create_stitching(self, store_composition):
+        self._create_stitched_volume(store_composition=store_composition)
+
+
+class _RawDatasetsContext(AbstractContextManager):
+    """
+    return volume data for all input volume (target: used for volume stitching).
+    If the volume is an HDF5Volume then the HDF5 dataset will be used (on disk)
+    If the volume is of another type then it will be loaded in memory then used (more memory consuming)
+    """
+
+    def __init__(self, volumes: tuple, alignment_axis_1) -> None:
+        super().__init__()
+        for volume in volumes:
+            if not isinstance(volume, VolumeBase):
+                raise TypeError(
+                    f"Volumes are expected to be an instance of {VolumeBase}. {type(volume)} provided instead"
+                )
+
+        self._volumes = volumes
+        self.__file_handlers = []
+        self._alignment_axis_1 = alignment_axis_1
+
+    @property
+    def alignment_axis_1(self):
+        return self._alignment_axis_1
+
+    def __enter__(self):
+        # handle the specific case of HDF5. Goal: avoid getting the full stitched volume in memory
+        datasets = []
+        shapes = {volume.get_volume_shape()[1] for volume in self._volumes}
+        axis_1_dim = max(shapes)
+        axis_1_need_padding = len(shapes) > 1
+
+        try:
+            for volume in self._volumes:
+                if volume.data is not None:
+                    data = volume.data
+                elif isinstance(volume, HDF5Volume):
+                    file_handler = HDF5File(volume.data_url.file_path(), mode="r")
+                    dataset = file_handler[volume.data_url.data_path()]
+                    data = dataset
+                    self.__file_handlers.append(file_handler)
+                # for other file format: load the full dataset in memory
+                else:
+                    data = volume.load_data(store=False)
+                    if data is None:
+                        raise ValueError(f"No data found for volume {volume.get_identifier()}")
+                if axis_1_need_padding:
+                    data = self.add_padding(data=data, axis_1_dim=axis_1_dim, alignment=self.alignment_axis_1)
+                datasets.append(data)
+        except Exception as e:
+            # if some errors happen during loading HDF5
+            for file_handled in self.__file_handlers:
+                file_handled.close()
+            raise e
+
+        return datasets
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        success = True
+        for file_handler in self.__file_handlers:
+            success = success and file_handler.close()
+        if exc_type is None:
+            return success
+
+    def add_padding(self, data: Union[h5py.Dataset, numpy.ndarray], axis_1_dim, alignment: AlignmentAxis1):
+        alignment = AlignmentAxis1.from_value(alignment)
+        if alignment is AlignmentAxis1.BACK:
+            axis_1_pad_width = (axis_1_dim - data.shape[1], 0)
+        elif alignment is AlignmentAxis1.CENTER:
+            half_width = int((axis_1_dim - data.shape[1]) / 2)
+            axis_1_pad_width = (half_width, axis_1_dim - data.shape[1] - half_width)
+        elif alignment is AlignmentAxis1.FRONT:
+            axis_1_pad_width = (0, axis_1_dim - data.shape[1])
+        else:
+            raise ValueError(f"alignment {alignment} is not handled")
+
+        return PaddedRawData(
+            data=data,
+            axis_1_pad_width=axis_1_pad_width,
+        )