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

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/pipeline/utils.py

@@ -0,0 +1,116 @@
+from ..utils import deprecated_class
+from .config_validators import str2bool
+from dataclasses import dataclass
+import os
+
+#
+# Decorators and callback mechanism
+#
+
+
+def use_options(step_name, step_attr):
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            self = args[0]
+            if step_name not in self.processing_steps:
+                self.__setattr__(step_attr, None)
+                return
+            self._steps_name2component[step_name] = step_attr
+            self._steps_component2name[step_attr] = step_name
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def pipeline_step(step_attr, step_desc):
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            self = args[0]
+            if self.__getattribute__(step_attr) is None:
+                return
+            self.logger.info(step_desc)
+            res = func(*args, **kwargs)
+            step_name = self._steps_component2name[step_attr]
+            callbacks = self._callbacks.get(step_name, None)
+            if callbacks is not None:
+                for callback in callbacks:
+                    callback(self)
+            if self.datadump_manager is not None and step_name in self.datadump_manager.data_dump:
+                self.datadump_manager.dump_data_to_file(step_name, self.radios)
+            return res
+
+        return wrapper
+
+    return decorator
+
+
+#
+# sub-region, shapes, etc
+#
+
+
+def get_subregion(sub_region, ndim=3):
+    """
+    Return a "normalized" sub-region in the form ((start_z, end_z), (start_y, end_y), (start_x, end_x)).
+
+    Parameters
+    ----------
+    sub_region: tuple
+        A tuple of tuples or tuple of integers.
+
+    Notes
+    -----
+    The parameter "sub_region" is normally a tuple of tuples of integers.
+    However it can be more convenient to use tuple of integers.
+    This function will attempt at catching the different cases, but will fail if
+    'sub_region' contains heterogeneous types (ex. tuples along with int)
+    """
+    if sub_region is None:
+        res = ((None, None),)
+    elif hasattr(sub_region[0], "__iter__"):
+        if set(map(len, sub_region)) != set([2]):
+            raise ValueError("Expected each tuple to be in the form (start, end)")
+        res = sub_region
+    else:
+        if len(sub_region) % 2:
+            raise ValueError("Expected even number of elements")
+        starts, ends = sub_region[::2], sub_region[1::2]
+        res = tuple([(s, e) for s, e in zip(starts, ends)])
+    if len(res) != ndim:
+        res += ((None, None),) * (ndim - len(res))
+    return res
+
+
+#
+# Writer - moved to pipeline.writer
+#
+
+from .writer import WriterManager
+
+WriterConfigurator = deprecated_class("WriterConfigurator moved to nabu.pipeline.writer.WriterManager", do_print=True)(
+    WriterManager
+)
+
+
+@dataclass
+class EnvSettings:
+    """This class centralises the definitions, possibly documentation, and access to environnt variable
+    driven settings.
+    It is meant to be used in the following way:
+          from nabu.utils import nabu_env_settings
+          if not nabu_env_settings.skip_tomoscan_checks:
+                do something
+    """
+
+    skip_tomoscan_checks: bool = False
+
+
+def _get_nabu_environment_variables():
+    nabu_env_settings = EnvSettings()
+    nabu_env_settings.skip_tomoscan_checks = str2bool(os.getenv("SKIP_TOMOSCAN_CHECK", "0"))
+    return nabu_env_settings
+
+
+nabu_env_settings = _get_nabu_environment_variables()

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/pipeline/writer.py

@@ -0,0 +1,193 @@
+from os import path
+from tomoscan.esrf import TIFFVolume, MultiTIFFVolume, EDFVolume, JP2KVolume
+from ..utils import check_supported, get_num_threads
+from ..resources.logger import LoggerOrPrint
+from ..io.writer import NXProcessWriter, HSTVolVolume, NXVolVolume
+from ..io.utils import convert_dict_values
+from .params import files_formats
+
+
+class WriterManager:
+
+    """
+    This class is a wrapper on top of all "writers".
+    It will create the right "writer" with all the necessary options, and the histogram writer.
+
+    The layout is the following.
+
+       * Single-big-file volume formats (big-tiff, .vol):
+           - no start index
+           - everything is increasingly written in one file
+       * Multiple-frames per file (HDF5 + master-file):
+          - needs a start index (change file_prefix) for each partial file
+          - needs a subdirectory for partial files
+       * One-file-per-frame (tiff, edf, jp2)
+          - start_index
+
+    When saving intermediate steps (eg. sinogram): HDF5 format is always used.
+    """
+
+    _overwrite_warned = False
+    _writer_classes = {
+        "hdf5": NXVolVolume,
+        "tiff": TIFFVolume,
+        "bigtiff": MultiTIFFVolume,
+        "jp2": JP2KVolume,
+        "edf": EDFVolume,
+        "vol": HSTVolVolume,
+    }
+
+    def __init__(
+        self,
+        output_dir,
+        file_prefix,
+        file_format="hdf5",
+        overwrite=False,
+        start_index=0,
+        logger=None,
+        metadata=None,
+        histogram=False,
+        extra_options=None,
+    ):
+        """
+        Create a Writer from a set of parameters.
+
+        Parameters
+        ----------
+        output_dir: str
+            Directory where the file(s) will be written.
+        file_prefix: str
+            File prefix (without leading path)
+        start_index: int, optional
+            Index to start the files numbering (filename_0123.ext).
+            Default is 0.
+            Ignored for HDF5 extension.
+        logger: nabu.resources.logger.Logger, optional
+            Logger object
+        metadata: dict, optional
+            Metadata, eg. information on various processing steps. For HDF5, it will go to "configuration"
+        histogram: bool, optional
+            Whether to also write a histogram of data. If set to True, it will configure
+            an additional "writer".
+        extra_options: dict, optional
+            Other advanced options to pass to Writer class.
+        """
+        self.extra_options = extra_options or {}
+        self._set_file_format(file_format)
+        self.overwrite = overwrite
+        self.start_index = start_index
+        self.logger = LoggerOrPrint(logger)
+        self.histogram = histogram
+        self.output_dir = output_dir
+        self.file_prefix = file_prefix
+        self.metadata = convert_dict_values(metadata or {}, {None: "None"})
+        self._init_writer()
+        self._init_histogram_writer()
+
+    def _set_file_format(self, file_format):
+        check_supported(file_format, files_formats, "file format")
+        self.file_format = files_formats[file_format]
+        self._is_bigtiff = file_format in ["tiff", "tif"] and any(
+            [self.extra_options.get(opt, False) for opt in ["tiff_single_file", "use_bigtiff"]]
+        )
+        if self._is_bigtiff:
+            self.file_format = "bigtiff"
+
+    @staticmethod
+    def get_first_fname(vol_writer):
+        if hasattr(vol_writer, "file_path"):
+            return path.dirname(vol_writer.file_path)
+        dirname = vol_writer.data_url.file_path()
+        fname = vol_writer.data_url.data_path().format(
+            volume_basename=vol_writer._volume_basename,
+            index_zfill6=vol_writer.start_index,
+            data_extension=vol_writer.extension or vol_writer.DEFAULT_DATA_EXTENSION,
+        )
+        return path.join(dirname, fname)
+
+    @staticmethod
+    def get_fname(vol_writer):
+        if hasattr(vol_writer, "file_path"):
+            # several frames per file - return the file itself
+            return vol_writer.file_path
+        # one file per frame - return the directory
+        return vol_writer.data_url.file_path()
+
+    def _init_writer(self):
+        if self.file_format in ["tiff", "edf", "jp2", "hdf5"]:
+            writer_kwargs = {
+                "folder": self.output_dir,
+                "volume_basename": self.file_prefix,
+                "start_index": self.start_index,
+                "overwrite": self.overwrite,
+            }
+            if self.file_format == "hdf5":
+                writer_kwargs["data_path"] = self.metadata.get("entry", "entry")
+                writer_kwargs["process_name"] = self.metadata.get("process_name", "reconstruction")
+                writer_kwargs["create_subfolder"] = self.extra_options.get("create_subfolder", True)
+            elif self.file_format == "jp2":
+                writer_kwargs["cratios"] = self.metadata.get("jpeg2000_compression_ratio", None)
+                writer_kwargs["clip_values"] = self.metadata.get("float_clip_values", None)
+                writer_kwargs["n_threads"] = get_num_threads()
+        elif self.file_format in ["vol", "bigtiff"]:
+            writer_kwargs = {
+                "file_path": path.join(
+                    self.output_dir, self.file_prefix + "." + self.file_format.replace("bigtiff", "tiff")
+                ),
+                "overwrite": self.overwrite,
+                "append": self.extra_options.get("single_output_file_initialized", False),
+                "hst_metadata": self.extra_options.get("raw_vol_metadata", {}),
+            }
+        else:
+            raise ValueError("Unsupported file format: %s" % self.file_format)
+        self._h5_entry = self.metadata.get("entry", "entry")
+        self.writer = self._writer_classes[self.file_format](**writer_kwargs)
+        self.fname = self.get_fname(self.writer)
+        if path.exists(self.fname):
+            err = "File already exists: %s" % self.fname
+            if self.overwrite:
+                if not (self.__class__._overwrite_warned):
+                    self.logger.warning(err + ". It will be overwritten as requested in configuration")
+                    self.__class__._overwrite_warned = True
+            else:
+                self.logger.fatal(err)
+                raise ValueError(err)
+
+    def _init_histogram_writer(self):
+        if not self.histogram:
+            return
+        separate_histogram_file = not (self.file_format == "hdf5")
+        if separate_histogram_file:
+            fmode = "w"
+            hist_fname = path.join(self.output_dir, "histogram_%05d.hdf5" % self.start_index)
+        else:
+            fmode = "a"
+            hist_fname = self.fname
+        # Nabu's original NXProcessWriter has to be used here, as histogram is not 3D
+        self.histogram_writer = NXProcessWriter(
+            hist_fname,
+            entry=self._h5_entry,
+            filemode=fmode,
+            overwrite=True,
+        )
+
+    def write_histogram(self, data, config=None, processing_index=1):
+        if not (self.histogram):
+            return
+        self.histogram_writer.write(
+            data,
+            "histogram",
+            processing_index=processing_index,
+            config=config,
+            is_frames_stack=False,
+            direct_access=False,
+        )
+
+    def _write_metadata(self):
+        self.writer.metadata = self.metadata
+        self.writer.save_metadata()
+
+    def write_data(self, data):
+        self.writer.data = data
+        self.writer.save()
+        # self._write_metadata()

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/preproc/__init__.py

@@ -0,0 +1,7 @@
+from .ccd import CCDFilter, Log
+from .ctf import CTFPhaseRetrieval
+from .distortion import DistortionCorrection
+from .double_flatfield import DoubleFlatField
+from .flatfield import FlatField, FlatFieldDataUrls
+from .phase import PaganinPhaseRetrieval
+from .shift import VerticalShift

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/preproc/alignment.py

@@ -0,0 +1,11 @@
+# Backward compat.
+from ..estimation.alignment import AlignmentBase
+from ..estimation.cor import (
+    CenterOfRotation,
+    CenterOfRotationAdaptiveSearch,
+    CenterOfRotationGrowingWindow,
+    CenterOfRotationSlidingWindow,
+)
+from ..estimation.translation import DetectorTranslationAlongBeam
+from ..estimation.focus import CameraFocus
+from ..estimation.tilt import CameraTilt

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/preproc/ccd.py

@@ -0,0 +1,161 @@
+import numpy as np
+from ..utils import check_supported
+from silx.math.medianfilter import medfilt2d
+
+
+class CCDFilter:
+    """
+    Filtering applied on radios.
+    """
+
+    _supported_ccd_corrections = ["median_clip"]
+
+    def __init__(
+        self,
+        radios_shape: tuple,
+        correction_type: str = "median_clip",
+        median_clip_thresh: float = 0.1,
+        abs_diff=False,
+        preserve_borders=False,
+    ):
+        """
+        Initialize a CCDCorrection instance.
+
+        Parameters
+        -----------
+        radios_shape: tuple
+            A tuple describing the shape of the radios stack, in the form
+            `(n_radios, n_z, n_x)`.
+        correction_type: str
+            Correction type for radios ("median_clip", "sigma_clip", ...)
+        median_clip_thresh: float, optional
+            Threshold for the median clipping method.
+        abs_diff: boolean
+            by default False:  the correction is triggered  when img - median > threshold.
+            If equals True:    correction is triggered for abs(img-media) > threshold
+        preserve borders: boolean
+            by default False:
+            If equals True:  the borders (width=1) are not modified.
+
+
+        Notes
+        ------
+        A CCD correction is a process (usually filtering) taking place in the
+        radios space.
+        Available filters:
+           - median_clip: if the value of the current pixel exceeds the median
+             of adjacent pixels (a 3x3 neighborhood) more than a threshold,
+             then this pixel value is set to the median value.
+        """
+        self._set_radios_shape(radios_shape)
+        check_supported(correction_type, self._supported_ccd_corrections, "CCD correction mode")
+        self.correction_type = correction_type
+        self.median_clip_thresh = median_clip_thresh
+        self.abs_diff = abs_diff
+        self.preserve_borders = preserve_borders
+
+    def _set_radios_shape(self, radios_shape):
+        if len(radios_shape) == 2:
+            self.radios_shape = (1,) + radios_shape
+        elif len(radios_shape) == 3:
+            self.radios_shape = radios_shape
+        else:
+            raise ValueError("Expected radios to have 2 or 3 dimensions")
+        n_radios, n_z, n_x = self.radios_shape
+        self.n_radios = n_radios
+        self.n_angles = n_radios
+        self.shape = (n_z, n_x)
+
+    @staticmethod
+    def median_filter(img):
+        """
+        Perform a median filtering on an image.
+        """
+        return medfilt2d(img, (3, 3), mode="reflect")
+
+    def median_clip_mask(self, img, return_medians=False):
+        """
+        Compute a mask indicating whether a pixel is valid or not, according to
+        the median-clip method.
+
+        Parameters
+        ----------
+        img: numpy.ndarray
+            Input image
+        return_medians: bool, optional
+            Whether to return the median values additionally to the mask.
+        """
+        median_values = self.median_filter(img)
+        if not self.abs_diff:
+            invalid_mask = img >= median_values + self.median_clip_thresh
+        else:
+            invalid_mask = abs(img - median_values) > self.median_clip_thresh
+
+        if return_medians:
+            return invalid_mask, median_values
+        else:
+            return invalid_mask
+
+    def median_clip_correction(self, radio, output=None):
+        """
+        Compute the median clip correction on one image.
+
+        Parameters
+        ----------
+        radios: numpy.ndarray, optional
+            A radio image.
+        output: numpy.ndarray, optional
+            Output array
+        """
+        assert radio.shape == self.shape
+        if output is None:
+            output = np.copy(radio)
+        else:
+            output[:] = radio[:]
+        invalid_mask, medians = self.median_clip_mask(radio, return_medians=True)
+
+        if self.preserve_borders:
+            fixed_border = np.array(radio[[0, 0, -1, -1], [0, -1, 0, -1]])
+
+        output[invalid_mask] = medians[invalid_mask]
+
+        if self.preserve_borders:
+            output[[0, 0, -1, -1], [0, -1, 0, -1]] = fixed_border
+
+        return output
+
+
+class Log:
+    """
+    Helper class to take -log(radios)
+
+    Parameters
+    -----------
+    clip_min: float, optional
+        Before taking the logarithm, the values are clipped to this minimum.
+    clip_max: float, optional
+        Before taking the logarithm, the values are clipped to this maximum.
+    """
+
+    def __init__(self, radios_shape, clip_min=None, clip_max=None):
+        self.radios_shape = radios_shape
+        self.clip_min = clip_min
+        self.clip_max = clip_max
+
+    def take_logarithm(self, radios):
+        """
+        Take the negative logarithm of a radios chunk.
+        Processing is done in-place !
+
+        Parameters
+        -----------
+        radios: array
+            Radios chunk.
+        """
+        if (self.clip_min is not None) or (self.clip_max is not None):
+            np.clip(radios, self.clip_min, self.clip_max, out=radios)
+            np.log(radios, out=radios)
+        else:
+            np.log(radios, out=radios)
+        radios[:] *= -1
+        return radios

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/preproc/ccd_cuda.py

@@ -0,0 +1,169 @@
+import numpy as np
+from ..preproc.ccd import CCDFilter, Log
+from ..processing.medfilt_cuda import MedianFilter
+from ..utils import get_cuda_srcfile, updiv, deprecated_class
+from ..cuda.utils import __has_pycuda__
+
+if __has_pycuda__:
+    from ..cuda.kernel import CudaKernel
+
+# COMPAT.
+from .flatfield_cuda import (
+    CudaFlatField as CudaFlatfield_,
+    CudaFlatFieldArrays as CudaFlatFieldArrays_,
+    CudaFlatFieldDataUrls as CudaFlatFieldDataUrls_,
+)
+
+FlatField = deprecated_class(
+    "preproc.ccd_cuda.CudaFlatField was moved to preproc.flatfield_cuda.CudaFlatField", do_print=True
+)(CudaFlatfield_)
+FlatFieldArrays = deprecated_class(
+    "preproc.ccd_cuda.CudaFlatFieldArrays was moved to preproc.flatfield_cuda.CudaFlatFieldArrays", do_print=True
+)(CudaFlatFieldArrays_)
+FlatFieldDataUrls = deprecated_class(
+    "preproc.ccd_cuda.CudaFlatFieldDataUrls was moved to preproc.flatfield_cuda.CudaFlatFieldDataUrls", do_print=True
+)(CudaFlatFieldDataUrls_)
+#
+
+
+class CudaCCDFilter(CCDFilter):
+    def __init__(
+        self,
+        radios_shape,
+        correction_type="median_clip",
+        median_clip_thresh=0.1,
+        abs_diff=False,
+        cuda_options=None,
+    ):
+        """
+        Initialize a CudaCCDCorrection instance.
+        Please refer to the documentation of CCDCorrection.
+        """
+        super().__init__(
+            radios_shape,
+            correction_type=correction_type,
+            median_clip_thresh=median_clip_thresh,
+        )
+        self._set_cuda_options(cuda_options)
+        self.cuda_median_filter = None
+        if correction_type == "median_clip":
+            self.cuda_median_filter = MedianFilter(
+                self.shape,
+                footprint=(3, 3),
+                mode="reflect",
+                threshold=median_clip_thresh,
+                abs_diff=abs_diff,
+                cuda_options={
+                    "device_id": self.cuda_options["device_id"],
+                    "ctx": self.cuda_options["ctx"],
+                    "cleanup_at_exit": self.cuda_options["cleanup_at_exit"],
+                },
+            )
+
+    def _set_cuda_options(self, user_cuda_options):
+        self.cuda_options = {"device_id": None, "ctx": None, "cleanup_at_exit": None}
+        if user_cuda_options is None:
+            user_cuda_options = {}
+        self.cuda_options.update(user_cuda_options)
+
+    def median_clip_correction(self, radio, output=None):
+        """
+        Compute the median clip correction on one image.
+
+        Parameters
+        ----------
+        radio: pycuda.gpuarray
+            A radio image
+        output: pycuda.gpuarray, optional
+            Output data.
+        """
+        assert radio.shape == self.shape
+        return self.cuda_median_filter.medfilt2(radio, output=output)
+
+
+CudaCCDCorrection = deprecated_class("CudaCCDCorrection is replaced with CudaCCDFilter", do_print=True)(CudaCCDFilter)
+
+
+class CudaLog(Log):
+    """
+    Helper class to take -log(radios)
+    """
+
+    def __init__(self, radios_shape, clip_min=None, clip_max=None):
+        """
+        Initialize a Log processing.
+
+        Parameters
+        -----------
+        radios_shape: tuple
+            The shape of 3D radios stack.
+        clip_min: float, optional
+            Data smaller than this value is replaced by this value.
+        clip_max: float, optional.
+            Data bigger than this value is replaced by this value.
+        """
+        super().__init__(radios_shape, clip_min=clip_min, clip_max=clip_max)
+        self._init_kernels()
+
+    def _init_kernels(self):
+        self._do_clip_min = int(self.clip_min is not None)
+        self._do_clip_max = int(self.clip_max is not None)
+        self.clip_min = np.float32(self.clip_min or 0)
+        self.clip_max = np.float32(self.clip_max or 1)
+        self._nlog_srcfile = get_cuda_srcfile("ElementOp.cu")
+        nz, ny, nx = self.radios_shape
+        self._nx = np.int32(nx)
+        self._ny = np.int32(ny)
+        self._nz = np.int32(nz)
+        self._nthreadsperblock = (16, 16, 4)  # TODO tune ?
+        self._nblocks = tuple([updiv(n, p) for n, p in zip([nx, ny, nz], self._nthreadsperblock)])
+
+        self.nlog_kernel = CudaKernel(
+            "nlog",
+            filename=self._nlog_srcfile,
+            signature="Piiiff",
+            options=[
+                "-DDO_CLIP_MIN=%d" % self._do_clip_min,
+                "-DDO_CLIP_MAX=%d" % self._do_clip_max,
+            ],
+        )
+
+    def take_logarithm(self, radios, clip_min=None, clip_max=None):
+        """
+        Take the negative logarithm of a radios chunk.
+
+        Parameters
+        -----------
+        radios: `pycuda.gpuarray.GPUArray`
+            Radios chunk
+            If not provided, a new GPU array is created.
+        clip_min: float, optional
+            Before taking the logarithm, the values are clipped to this minimum.
+        clip_max: float, optional
+            Before taking the logarithm, the values are clipped to this maximum.
+        """
+        clip_min = clip_min or self.clip_min
+        clip_max = clip_max or self.clip_max
+        if radios.flags.c_contiguous:
+            self.nlog_kernel(
+                radios,
+                self._nx,
+                self._ny,
+                self._nz,
+                clip_min,
+                clip_max,
+                grid=self._nblocks,
+                block=self._nthreadsperblock,
+            )
+        else:
+            # map-like operations cannot be directly applied on 3D arrays
+            # that are not C-contiguous. We have to process image per image.
+            # TODO it's even worse when each single frame is not C-contiguous. For now this case is not handled
+            nz = np.int32(1)
+            nthreadsperblock = (32, 32, 1)
+            nblocks = tuple([updiv(n, p) for n, p in zip([int(self._nx), int(self._ny), int(nz)], nthreadsperblock)])
+            for i in range(radios.shape[0]):
+                self.nlog_kernel(
+                    radios[i], self._nx, self._ny, nz, clip_min, clip_max, grid=nblocks, block=nthreadsperblock
+                )
+        return radios