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

nabu/stitching/utils/post_processing.py

@@ -0,0 +1,281 @@
+import os
+import logging
+from typing import Optional, Union
+from nabu import version as nabu_version
+from nabu.stitching.config import (
+    PreProcessedSingleAxisStitchingConfiguration,
+    PostProcessedSingleAxisStitchingConfiguration,
+    SingleAxisStitchingConfiguration,
+)
+from nabu.stitching.stitcher.single_axis import PROGRESS_BAR_STITCH_VOL_DESC
+from nabu.io.writer import get_datetime
+from tomoscan.factory import Factory as TomoscanFactory
+from silx.io.dictdump import dicttonx
+from nxtomo.application.nxtomo import NXtomo
+from tomoscan.utils.volume import concatenate as concatenate_volumes
+from tomoscan.esrf.volume import HDF5Volume
+from contextlib import AbstractContextManager
+from threading import Thread
+from time import sleep
+
+_logger = logging.getLogger(__name__)
+
+
+class StitchingPostProcAggregation:
+    """
+    for remote stitching each process will stitch a part of the volume or projections.
+    Then once all are finished we want to aggregate them all to a final volume or NXtomo.
+
+    This is the goal of this class.
+    Please be careful with API. This is already inheriting from a tomwer class
+
+    :param stitching_config: configuration of the stitching configuration
+    :param futures: futures that just run
+    :param existing_objs: futures that just run
+    :param progress_bars: tqdm progress bars for each jobs
+    """
+
+    def __init__(
+        self,
+        stitching_config: SingleAxisStitchingConfiguration,
+        futures: Optional[tuple] = None,
+        existing_objs_ids: Optional[tuple] = None,
+        progress_bars: Optional[dict] = None,
+    ) -> None:
+        if not isinstance(stitching_config, (SingleAxisStitchingConfiguration)):
+            raise TypeError(f"stitching_config should be an instance of {SingleAxisStitchingConfiguration}")
+        if not ((existing_objs_ids is None) ^ (futures is None)):
+            raise ValueError("Either existing_objs or futures should be provided (can't provide both)")
+        if progress_bars is not None and not isinstance(progress_bars, dict):
+            raise TypeError(f"'progress_bars' should be None or an instance of a dict. Got {type(progress_bars)}")
+        self._futures = futures
+        self._stitching_config = stitching_config
+        self._existing_objs_ids = existing_objs_ids
+        self._progress_bars = progress_bars or {}
+
+    @property
+    def futures(self):
+        return self._futures
+
+    @property
+    def progress_bars(self) -> dict:
+        return self._progress_bars
+
+    def retrieve_tomo_objects(self) -> tuple:
+        """
+        Return tomo objects to be stitched together. Either from future or from existing_objs
+        """
+        if self._existing_objs_ids is not None:
+            scan_ids = self._existing_objs_ids
+        else:
+            results = {}
+            _logger.info(
+                f"wait for slurm job to be completed. Advancement will be created once slurm job output file will be available"
+            )
+            for obj_id, future in self.futures.items():
+                results[obj_id] = future.result()
+
+            failed = tuple(
+                filter(
+                    lambda x: x.exception() is not None,
+                    self.futures.values(),
+                )
+            )
+            if len(failed) > 0:
+                # if some job failed: useless to do the concatenation
+                exceptions = " ; ".join([f"{job} : {job.exception()}" for job in failed])
+                raise RuntimeError(f"some job failed. Won't do the concatenation. Exceptiosn are {exceptions}")
+
+            canceled = tuple(
+                filter(
+                    lambda x: x.cancelled(),
+                    self.futures.values(),
+                )
+            )
+            if len(canceled) > 0:
+                # if some job canceled: useless to do the concatenation
+                raise RuntimeError(f"some job failed. Won't do the concatenation. Jobs are {' ; '.join(canceled)}")
+            scan_ids = results.keys()
+        return [TomoscanFactory.create_tomo_object_from_identifier(scan_id) for scan_id in scan_ids]
+
+    def dump_stitching_config_as_nx_process(self, file_path: str, data_path: str, overwrite: bool, process_name: str):
+        dict_to_dump = {
+            process_name: {
+                "config": self._stitching_config.to_dict(),
+                "program": "nabu-stitching",
+                "version": nabu_version,
+                "date": get_datetime(),
+            },
+            f"{process_name}@NX_class": "NXprocess",
+        }
+
+        dicttonx(
+            dict_to_dump,
+            h5file=file_path,
+            h5path=data_path,
+            update_mode="replace" if overwrite else "add",
+            mode="a",
+        )
+
+    @property
+    def stitching_config(self) -> SingleAxisStitchingConfiguration:
+        return self._stitching_config
+
+    def process(self) -> None:
+        """
+        main function
+        """
+
+        # concatenate result
+        _logger.info("all job succeeded. Concatenate results")
+        if isinstance(self._stitching_config, PreProcessedSingleAxisStitchingConfiguration):
+            # 1: case of a pre-processing stitching
+            with self.follow_progress():
+                scans = self.retrieve_tomo_objects()
+            nx_tomos = []
+            for scan in scans:
+                if not os.path.exists(scan.master_file):
+                    raise RuntimeError(
+                        f"output file not created ({scan.master_file}). Stitching failed. "
+                        "Please check slurm .out files to have more information. Most likely the slurm configuration is invalid. "
+                        "(partition name not existing...)"
+                    )
+                nx_tomos.append(
+                    NXtomo().load(
+                        file_path=scan.master_file,
+                        data_path=scan.entry,
+                    )
+                )
+            final_nx_tomo = NXtomo.concatenate(nx_tomos)
+            final_nx_tomo.save(
+                file_path=self.stitching_config.output_file_path,
+                data_path=self.stitching_config.output_data_path,
+                overwrite=self.stitching_config.overwrite_results,
+            )
+
+            # dump NXprocess if possible
+            parts = self.stitching_config.output_data_path.split("/")
+            process_name = parts[-1] + "_stitching"
+            if len(parts) < 2:
+                data_path = "/"
+            else:
+                data_path = "/".join(parts[:-1])
+
+            self.dump_stitching_config_as_nx_process(
+                file_path=self.stitching_config.output_file_path,
+                data_path=data_path,
+                process_name=process_name,
+                overwrite=self.stitching_config.overwrite_results,
+            )
+
+        elif isinstance(self.stitching_config, PostProcessedSingleAxisStitchingConfiguration):
+            # 2: case of a post-processing stitching
+            with self.follow_progress():
+                outputs_sub_volumes = self.retrieve_tomo_objects()
+            concatenate_volumes(
+                output_volume=self.stitching_config.output_volume,
+                volumes=tuple(outputs_sub_volumes),
+                axis=1,
+            )
+
+            if isinstance(self.stitching_config.output_volume, HDF5Volume):
+                parts = self.stitching_config.output_volume.metadata_url.data_path().split("/")
+                process_name = parts[-1] + "_stitching"
+                if len(parts) < 2:
+                    data_path = "/"
+                else:
+                    data_path = "/".join(parts[:-1])
+
+                self.dump_stitching_config_as_nx_process(
+                    file_path=self.stitching_config.output_volume.metadata_url.file_path(),
+                    data_path=data_path,
+                    process_name=process_name,
+                    overwrite=self.stitching_config.overwrite_results,
+                )
+        else:
+            raise TypeError(f"stitching_config type ({type(self.stitching_config)}) not handled")
+
+    def follow_progress(self) -> AbstractContextManager:
+        return SlurmStitchingFollowerContext(
+            output_files_to_progress_bars={
+                job._get_output_file_path(): progress_bar for (job, progress_bar) in self.progress_bars.items()
+            }
+        )
+
+
+class SlurmStitchingFollowerContext(AbstractContextManager):
+    """Util class to provide user feedback from stitching done on slurm"""
+
+    def __init__(self, output_files_to_progress_bars: dict):
+        self._update_thread = SlurmStitchingFollowerThread(file_to_progress_bar=output_files_to_progress_bars)
+
+    def __enter__(self) -> None:
+        self._update_thread.start()
+
+    def __exit__(self, *args, **kwargs):
+        self._update_thread.join(timeout=1.5)
+        for progress_bar in self._update_thread.file_to_progress_bar.values():
+            progress_bar.close()  # close to clean display as leave == False
+
+
+class SlurmStitchingFollowerThread(Thread):
+    """
+    Thread to check progression of stitching slurm job(s)
+    Read slurm jobs .out file each 'delay time' and look for a tqdm line at the end.
+    If it exists then deduce progress from it.
+
+    file_to_progress_bar provide for each slurm .out file the progress bar to update
+    """
+
+    def __init__(self, file_to_progress_bar: dict, delay_time: float = 0.5) -> None:
+        super().__init__()
+        self._stop_run = False
+        self._wait_time = delay_time
+        self._file_to_progress_bar = file_to_progress_bar
+        self._first_run = True
+
+    @property
+    def file_to_progress_bar(self) -> dict:
+        return self._file_to_progress_bar
+
+    def run(self) -> None:
+        while not self._stop_run:
+            for file_path, progress_bar in self._file_to_progress_bar.items():
+                if self._first_run:
+                    # make sure each progress bar have been refreshed at least one
+                    progress_bar.refresh()
+
+                if not os.path.exists(file_path):
+                    continue
+                with open(file_path, "r") as f:
+                    try:
+                        last_line = f.readlines()[-1]
+                    except IndexError:
+                        continue
+                    advancement = self.cast_progress_line_from_log(line=last_line)
+                    if advancement is not None:
+                        progress_bar.n = advancement
+                        progress_bar.refresh()
+
+            self._first_run = False
+
+            sleep(self._wait_time)
+
+    def join(self, timeout: Union[float, None] = None) -> None:
+        self._stop_run = True
+        return super().join(timeout)
+
+    @staticmethod
+    def cast_progress_line_from_log(line: str) -> Optional[float]:
+        """Try to retrieve from a line from log the advancement (in percentage)"""
+        if PROGRESS_BAR_STITCH_VOL_DESC not in line or "%" not in line:
+            return None
+
+        str_before_percentage = line.split("%")[0].split(" ")[-1]
+        try:
+            advancement = float(str_before_percentage)
+        except ValueError:
+            _logger.debug(f"Failed to retrieve advancement from log file. Value got is {str_before_percentage}")
+            return None
+        else:
+            return advancement

nabu/stitching/utils/tests/test_post-processing.py

@@ -0,0 +1,21 @@
+import pytest
+
+from nabu.stitching.stitcher.single_axis import PROGRESS_BAR_STITCH_VOL_DESC
+from nabu.stitching.utils.post_processing import SlurmStitchingFollowerThread
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    {
+        "dump configuration: 100%|": None,
+        f"stitching : 100%|": None,
+        f"{PROGRESS_BAR_STITCH_VOL_DESC}: 42%": 42.0,
+        f"{PROGRESS_BAR_STITCH_VOL_DESC}: 56% toto: 23%": 56.0,
+        "": None,
+        "my%": None,
+    }.items(),
+)
+def test_SlurmStitchingFollowerContext(test_case):
+    """Test that the conversion from log lines created by tqdm can be read back"""
+    str_to_test, expected_result = test_case
+    assert SlurmStitchingFollowerThread.cast_progress_line_from_log(str_to_test) == expected_result

nabu/stitching/{utils.py → utils/utils.py}

@@ -3,22 +3,18 @@ from typing import Optional, Union
 import logging
 import functools
 import numpy
-from scipy.ndimage import affine_transform
 from tomoscan.scanbase import TomoScanBase
 from tomoscan.volumebase import VolumeBase
-from nxtomo.utils.transformation import build_matrix, UDDetTransformation
+from nxtomo.utils.transformation import build_matrix, DetYFlipTransformation
 from silx.utils.enum import Enum as _Enum
 from scipy.fft import rfftn as local_fftn
 from scipy.fft import irfftn as local_ifftn
-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
+from ..overlap import OverlapStitchingStrategy, ImageStichOverlapKernel
+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
@@ -66,39 +62,22 @@ class ShiftAlgorithm(_Enum):
             return super().from_value(value=value)
 
 
-def test_overlap_stitching_strategy(overlap_1, overlap_2, stitching_strategies):
-    """
-    stitch the two ovrelap with all the requested strategies.
-    Return a dictionary with stitching strategy as key and a result dict as value.
-    result dict keys are: 'weights_overlap_1', 'weights_overlap_2', 'stiching'
-    """
-    res = {}
-    for strategy in stitching_strategies:
-        s = OverlapStitchingStrategy.from_value(strategy)
-        stitcher = ZStichOverlapKernel(
-            stitching_strategy=s,
-            frame_width=overlap_1.shape[1],
-        )
-        stiched_overlap, w1, w2 = stitcher.stitch(overlap_1, overlap_2, check_input=True)
-        res[s.value] = {
-            "stitching": stiched_overlap,
-            "weights_overlap_1": w1,
-            "weights_overlap_2": w2,
-        }
-    return res
-
-
 def find_frame_relative_shifts(
     overlap_upper_frame: numpy.ndarray,
     overlap_lower_frame: numpy.ndarray,
-    estimated_shifts,
+    estimated_shifts: tuple,
+    overlap_axis: int,
     x_cross_correlation_function=None,
     y_cross_correlation_function=None,
     x_shifts_params: Optional[dict] = None,
     y_shifts_params: Optional[dict] = None,
 ):
+    """
+    :param overlap_axis: axis in [0, 1] on which the overlap exists. In image space. So 0 is aka y and 1 as x
+    """
+    if not overlap_axis in (0, 1):
+        raise ValueError(f"overlap_axis should be in (0, 1). Get {overlap_axis}")
     from nabu.stitching.config import (
-        KEY_WINDOW_SIZE,
         KEY_LOW_PASS_FILTER,
         KEY_HIGH_PASS_FILTER,
     )  # avoid cyclic import
@@ -188,6 +167,7 @@ def find_frame_relative_shifts(
 def find_volumes_relative_shifts(
     upper_volume: VolumeBase,
     lower_volume: VolumeBase,
+    overlap_axis: int,
     estimated_shifts,
     dim_axis_1: int,
     dtype,
@@ -210,6 +190,15 @@ def find_volumes_relative_shifts(
 
     if x_shifts_params is None:
         x_shifts_params = {}
+    # convert from overlap_axis (3D acquisition space) to overlap_axis_proj_space.
+    if overlap_axis == 1:
+        raise NotImplementedError("finding projection shift along axis 1 is not handled for projections")
+    elif overlap_axis == 0:
+        overlap_axis_proj_space = 0
+    elif overlap_axis == 2:
+        overlap_axis_proj_space = 1
+    else:
+        raise ValueError(f"Stitching is done in 3D space. Expect axis to be in [0,2]. Get {overlap_axis}")
 
     alignment_axis_2 = AlignmentAxis2.from_value(alignment_axis_2)
     alignment_axis_1 = AlignmentAxis1.from_value(alignment_axis_1)
@@ -299,6 +288,7 @@ def find_volumes_relative_shifts(
         y_cross_correlation_function=y_cross_correlation_function,
         x_shifts_params=x_shifts_params,
         y_shifts_params=y_shifts_params,
+        overlap_axis=overlap_axis_proj_space,
     )
 
 
@@ -308,7 +298,8 @@ from nabu.pipeline.estimators import estimate_cor
 def find_projections_relative_shifts(
     upper_scan: TomoScanBase,
     lower_scan: TomoScanBase,
-    estimated_shifts,
+    estimated_shifts: tuple,
+    axis: int,
     flip_ud_upper_frame: bool = False,
     flip_ud_lower_frame: bool = False,
     projection_for_shift: Union[int, str] = "middle",
@@ -326,13 +317,16 @@ def find_projections_relative_shifts(
 
     :param TomoScanBase scan_0:
     :param TomoScanBase scan_1:
-    :param int axis_0_overlap_px: overlap between the two scans in pixel
+    :param tuple estimated_shifts: 'a priori' shift estimation
+    :param int axis: axis on which the overlap / stitching is happening. In the 3D space (sample, detector referential)
+    :param bool flip_ud_upper_frame: is the upper frame flipped
+    :param bool flip_ud_lower_frame: is the lower frame flipped
     :param Union[int,str] projection_for_shift: index fo the projection to use (in projection space or in scan space ?. For now in projection) or str. If str must be in (`middle`, `first`, `last`)
+    :param bool invert_order: are projections inverted between the two scans (case if rotation angle are inverted)
     :param str x_cross_correlation_function: optional method to refine x shift from computing cross correlation. For now valid values are: ("skimage", "nabu-fft")
     :param str y_cross_correlation_function: optional method to refine y shift from computing cross correlation. For now valid values are: ("skimage", "nabu-fft")
-    :param int minimal_overlap_area_for_cross_correlation: if first approximated overlap shift found from z_translation is lower than this value will fall back on taking the full image for the cross correlation and log a warning
-    :param bool invert_order: are projections inverted between the two scans (case if rotation angle are inverted)
-    :param tuple estimated_shifts: 'a priori' shift estimation
+    :param x_shifts_params: parameters to find the shift over x
+    :param y_shifts_params: parameters to find the shift over y
     :return: relative shift of scan_1 with scan_0 as reference: (y_shift, x_shift)
     :rtype: tuple
 
@@ -342,8 +336,18 @@ def find_projections_relative_shifts(
         x_shifts_params = {}
     if y_shifts_params is None:
         y_shifts_params = {}
-    if estimated_shifts[0] < 0:
-        raise ValueError("y_overlap_px is expected to be stricktly positive")
+
+    # convert from overlap_axis (3D acquisition space) to overlap_axis_proj_space.
+    if axis == 1:
+        axis_proj_space = 1
+    elif axis == 0:
+        axis_proj_space = 0
+    elif axis == 2:
+        raise NotImplementedError(
+            "finding projection shift along axis 1 (x-ray direction) is not handled for projections"
+        )
+    else:
+        raise ValueError(f"Stitching is done in 3D space. Expect axis to be in [0,2]. Get {axis}")
 
     x_cross_correlation_function = ShiftAlgorithm.from_value(x_cross_correlation_function)
     y_cross_correlation_function = ShiftAlgorithm.from_value(y_cross_correlation_function)
@@ -432,11 +436,11 @@ def find_projections_relative_shifts(
 
     upper_scan_transformations = list(upper_scan.get_detector_transformations(tuple()))
     if flip_ud_upper_frame:
-        upper_scan_transformations.append(UDDetTransformation())
+        upper_scan_transformations.append(DetYFlipTransformation(flip=True))
     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_transformations.append(DetYFlipTransformation(flip=True))
     lower_scan_trans_matrix = build_matrix(lower_scan_transformations)
     upper_proj = get_flat_fielded_proj(
         upper_scan,
@@ -453,16 +457,30 @@ def find_projections_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] // 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:]
+    if axis_proj_space == 0:
+        w_window_size = int(y_shifts_params.get(KEY_WINDOW_SIZE, 400))
     else:
-        overlap_upper_frame = upper_proj[-end_overlap:-start_overlap]
-    overlap_lower_frame = lower_proj[start_overlap:end_overlap]
+        w_window_size = int(x_shifts_params.get(KEY_WINDOW_SIZE, 400))
+    start_overlap = max(estimated_shifts[axis_proj_space] // 2 - w_window_size // 2, 0)
+    end_overlap = min(
+        estimated_shifts[axis_proj_space] // 2 + w_window_size // 2,
+        min(upper_proj.shape[axis_proj_space], lower_proj.shape[axis_proj_space]),
+    )
+    o_upper_sel = numpy.array(range(-end_overlap, -start_overlap))
+    overlap_upper_frame = numpy.take_along_axis(
+        upper_proj,
+        o_upper_sel[:, None] if axis_proj_space == 0 else o_upper_sel[None, :],
+        axis=axis_proj_space,
+    )
+    o_lower_sel = numpy.array(range(start_overlap, end_overlap))
+    overlap_lower_frame = numpy.take_along_axis(
+        lower_proj,
+        o_lower_sel[:, None] if axis_proj_space == 0 else o_upper_sel[None, :],
+        axis=axis_proj_space,
+    )
+
     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})")
+        raise ValueError(f"Fail to get consistent overlap ({overlap_upper_frame.shape} vs {overlap_lower_frame.shape})")
 
     return find_frame_relative_shifts(
         overlap_upper_frame=overlap_upper_frame,
@@ -472,6 +490,7 @@ def find_projections_relative_shifts(
         y_cross_correlation_function=y_cross_correlation_function,
         x_shifts_params=x_shifts_params,
         y_shifts_params=y_shifts_params,
+        overlap_axis=axis_proj_space,
     )
 
 
@@ -561,3 +580,11 @@ def find_shift_with_itk(img1: numpy.ndarray, img2: numpy.ndarray) -> tuple:
     translation_along_y = final_parameters.GetElement(1)
 
     return numpy.round(translation_along_y), numpy.round(translation_along_x)
+
+
+def from_slice_to_n_elements(slice_: Union[slice, tuple]):
+    """Return the number of element in a slice or in a tuple"""
+    if isinstance(slice_, slice):
+        return (slice_.stop - slice_.start) / (slice_.step or 1)
+    else:
+        return len(slice_)

nabu/stitching/y_stitching.py

@@ -0,0 +1,27 @@
+from tomoscan.identifier import BaseIdentifier
+from nabu.stitching.stitcher.y_stitcher import PreProcessingYStitcher as PreProcessYStitcher
+from nabu.stitching.config import PreProcessedYStitchingConfiguration
+
+
+def y_stitching(configuration: PreProcessedYStitchingConfiguration, progress=None) -> BaseIdentifier:
+    """
+    Apply stitching from provided configuration.
+    Stitching will be applied along the first axis - 1 (aka y).
+
+    like:
+                    axis 0
+                      ^
+                      |
+    x-ray             |
+    -------->          ------> axis 2
+                     /
+                    /
+                     axis 1
+    """
+    if isinstance(configuration, PreProcessedYStitchingConfiguration):
+        stitcher = PreProcessYStitcher(configuration=configuration, progress=progress)
+    else:
+        raise TypeError(
+            f"configuration is expected to be in {(PreProcessedYStitchingConfiguration, )}. {type(configuration)} provided"
+        )
+    return stitcher.stitch()