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

nabu/pipeline/estimators.py

@@ -2,41 +2,35 @@
 nabu.pipeline.estimators: helper classes/functions to estimate parameters of a dataset
 (center of rotation, detector tilt, etc).
 """
+
 import inspect
 import numpy as np
 import scipy.fft  #  pylint: disable=E0611
 from silx.io import get_data
-from typing import Union, Optional
 import math
-from numbers import Real
 from scipy import ndimage as nd
-
-from ..preproc.flatfield import FlatFieldDataUrls
+from ..preproc.flatfield import FlatField
 from ..estimation.cor import (
     CenterOfRotation,
     CenterOfRotationAdaptiveSearch,
     CenterOfRotationSlidingWindow,
     CenterOfRotationGrowingWindow,
-    CenterOfRotationFourierAngles,
     CenterOfRotationOctaveAccurate,
 )
-from ..estimation.cor_sino import SinoCorInterface
+from ..estimation.cor_sino import SinoCorInterface, CenterOfRotationFourierAngles, CenterOfRotationVo
 from ..estimation.tilt import CameraTilt
 from ..estimation.utils import is_fullturn_scan
 from ..resources.logger import LoggerOrPrint
 from ..resources.utils import extract_parameters
-from ..utils import check_supported, is_int
-from .params import tilt_methods
+from ..utils import check_supported, deprecation_warning, get_num_threads, is_int, is_scalar
 from ..resources.dataset_analyzer import get_radio_pair
 from ..processing.rotation import Rotation
-from ..io.reader import ChunkReader
 from ..preproc.ccd import Log, CCDFilter
 from ..misc import fourier_filters
-from .params import cor_methods
-from ..io.reader import load_images_from_dataurl_dict
+from .params import cor_methods, tilt_methods
 
 
-def estimate_cor(method, dataset_info, do_flatfield=True, cor_options: Optional[Union[str, dict]] = None, logger=None):
+def estimate_cor(method, dataset_info, do_flatfield=True, cor_options=None, logger=None):
     logger = LoggerOrPrint(logger)
     cor_options = cor_options or {}
     check_supported(method, list(cor_methods.keys()), "COR estimation method")
@@ -108,116 +102,41 @@ class CORFinderBase:
             Dataset information structure
         """
         check_supported(method, self.search_methods, "CoR estimation method")
+        self.method = method
+        self.cor_options = cor_options or {}
         self.logger = LoggerOrPrint(logger)
         self.dataset_info = dataset_info
         self.do_flatfield = do_flatfield
         self.shape = dataset_info.radio_dims[::-1]
-        self._init_cor_finder(method, cor_options)
+        self._get_lookup_side()
+        self._init_cor_finder()
 
-    def _init_cor_finder(self, method, cor_options):
-        self.method = method
-        if not isinstance(cor_options, (type(None), dict)):
-            raise TypeError(
-                f"cor_options is expected to be an optional instance of dict. Get {cor_options} ({type(cor_options)}) instead"
-            )
-        self.cor_options = {}
-        if isinstance(cor_options, dict):
-            self.cor_options.update(cor_options)
-
-        # tomotools internal meeting 07 feb 2024: Merge of options 'near_pos' and 'side'.
-        # See [minutes](https://gitlab.esrf.fr/tomotools/minutes/-/blob/master/minutes-20240207.md?ref_type=heads)
+    def _get_lookup_side(self):
+        """
+        Get the "initial guess" where the center-of-rotation (CoR) should be estimated.
+        For example 'center' means that CoR search will be done near the middle of the detector, i.e center column.
+        """
+        lookup_side = self.cor_options.get("side", None)
+        self._lookup_side = lookup_side
+        # User-provided scalar
+        if not (isinstance(lookup_side, str)) and np.isscalar(lookup_side):
+            return
 
-        detector_width = self.dataset_info.radio_dims[0]
         default_lookup_side = "right" if self.dataset_info.is_halftomo else "center"
-        near_init = self.cor_options.get("side", None)
 
-        if near_init is None:
-            near_init = default_lookup_side
-
-        if near_init == "from_file":
-            try:
-                near_pos = self.dataset_info.dataset_scanner.estimated_cor_frm_motor  # relative pos in pixels
-                if isinstance(near_pos, Real):
-                    # near_pos += detector_width // 2  # Field in NX is relative.
-                    self.cor_options.update({"near_pos": int(near_pos)})
-                else:
-                    near_init = default_lookup_side
-            except:
-                self.logger.warning(
-                    "COR estimation from motor position absent from NX file. Global search is performed."
-                )
-                near_init = default_lookup_side
-        elif isinstance(near_init, Real):
-            self.cor_options.update({"near_pos": int(near_init)})
-            near_init = "near"  # ???
-        elif near_init == "near":  # Legacy
-            if not isinstance(self.cor_options["near_pos"], Real):
-                self.logger.warning("Side option set to 'near' but no 'near_pos' option set.")
-                self.logger.warning("Set side to right if HA, center otherwise.")
-                near_init = default_lookup_side
-        elif near_init in ("left", "right", "center", "all"):
-            pass
-        else:
-            self.logger.warning(
-                f"COR option 'side' received {near_init} and should be either 'from_file' (default), 'left', 'right', 'center', 'near' or a number."
-            )
-
-        if isinstance(self.cor_options.get("near_pos", None), Real):
-            # Check validity of near_pos
-            if np.abs(self.cor_options["near_pos"]) > detector_width / 2:
-                self.logger.warning(
-                    f"Relative COR passed is greater than half the size of the detector. Did you enter a absolute COR position?"
-                )
-                self.logger.warning("Instead, the center of the detector is used.")
-                self.cor_options["near_pos"] = 0
-
-            # Set side from near_pos if passed.
-            if self.cor_options["near_pos"] < 0.0:
-                self.cor_options.update({"side": "left"})
-                near_init = "left"
+        # By default in nabu config, side='from_file' meaning that we inspect the dataset information for CoR metadata
+        if lookup_side == "from_file":
+            initial_cor_pos = self.dataset_info.dataset_scanner.x_rotation_axis_pixel_position  # relative pos in pixels
+            if initial_cor_pos is None or initial_cor_pos == 0:
+                self.logger.warning("Could not get an initial estimate for center of rotation in data file")
+                lookup_side = default_lookup_side
             else:
-                self.cor_options.update({"side": "right"})
-                near_init = "right"
-
-        self.cor_options.update({"side": near_init})
-
-        # At this stage : side is set to one of left, right, center near.
-        # and near_pos to a numeric value.
-
-        # if isinstance(self.cor_options["near_pos"], Real):
-        #    # estimated_cor_frm_motor value is supposed to be relative. Since the config documentation expects the "near_pos" options
-        #    # to be given as an absolute COR estimate, a conversion is needed.
-        #    self.cor_options["near_pos"] += detector_width // 2  # converted in absolute nb of pixels.
-        # if not (isinstance(self.cor_options["near_pos"], Real) or self.cor_options["near_pos"] == "ignore"):
-        #    self.cor_options.update({"near_pos": "ignore"})
-
-        # At this stage, cor_options["near_pos"] is either
-        #   - 'ignore':
-        #   - an (absolute) integer value (either the user-provided one if present or the NX one).
-
-        cor_class = self.search_methods[method]["class"]
-        self.cor_finder = cor_class(logger=self.logger, cor_options=self.cor_options)
-
-        lookup_side = self.cor_options.get("side", default_lookup_side)
-
-        # OctaveAccurate
-        # if cor_class == CenterOfRotationOctaveAccurate:
-        #    lookup_side = "center"
-        angles = self.dataset_info.rotation_angles
-
-        self.cor_exec_args = []
-        self.cor_exec_args.extend(self.search_methods[method].get("default_args", []))
-
-        # CenterOfRotationSlidingWindow is the only class to have a mandatory argument ("side")
-        # TODO - it would be more elegant to have it as a kwarg...
-        if len(self.cor_exec_args) > 0:
-            if cor_class in (CenterOfRotationSlidingWindow, CenterOfRotationOctaveAccurate):
-                self.cor_exec_args[0] = lookup_side
-            elif cor_class in (CenterOfRotationFourierAngles,):
-                self.cor_exec_args[0] = angles
-                self.cor_exec_args[1] = lookup_side
-        #
-        self.cor_exec_kwargs = update_func_kwargs(self.cor_finder.find_shift, self.cor_options)
+                lookup_side = initial_cor_pos
+            self._lookup_side = initial_cor_pos
+
+    def _init_cor_finder(self):
+        cor_finder_cls = self.search_methods[self.method]["class"]
+        self.cor_finder = cor_finder_cls(verbose=False, logger=self.logger, extra_options=None)
 
 
 class CORFinder(CORFinderBase):
@@ -235,19 +154,17 @@ class CORFinder(CORFinderBase):
         },
         "sliding-window": {
             "class": CenterOfRotationSlidingWindow,
-            "default_args": ["center"],
         },
         "growing-window": {
             "class": CenterOfRotationGrowingWindow,
         },
         "octave-accurate": {
             "class": CenterOfRotationOctaveAccurate,
-            "default_args": ["center"],
         },
     }
 
     def __init__(
-        self, method, dataset_info, do_flatfield=True, cor_options=None, logger=None, radio_angles: tuple = (0.0, np.pi)
+        self, method, dataset_info, do_flatfield=True, cor_options=None, logger=None, radio_angles=(0.0, np.pi)
     ):
         """
         Initialize a CORFinder object.
@@ -261,7 +178,6 @@ class CORFinder(CORFinderBase):
         super().__init__(method, dataset_info, do_flatfield=do_flatfield, cor_options=cor_options, logger=logger)
         self._radio_angles = radio_angles
         self._init_radios()
-        self._init_flatfield()
         self._apply_flatfield()
         self._apply_tilt()
 
@@ -270,21 +186,16 @@ class CORFinder(CORFinderBase):
             self.dataset_info, radio_angles=self._radio_angles, return_indices=True
         )
 
-    def _init_flatfield(self):
+    def _apply_flatfield(self):
         if not (self.do_flatfield):
             return
-        self.flatfield = FlatFieldDataUrls(
+        self.flatfield = FlatField(
             self.radios.shape,
             flats=self.dataset_info.flats,
             darks=self.dataset_info.darks,
             radios_indices=self._radios_indices,
             interpolation="linear",
-            convert_float=True,
         )
-
-    def _apply_flatfield(self):
-        if not (self.do_flatfield):
-            return
         self.flatfield.normalize_radios(self.radios)
 
     def _apply_tilt(self):
@@ -306,11 +217,14 @@ class CORFinder(CORFinderBase):
             The estimated center of rotation for the current dataset.
         """
         self.logger.info("Estimating center of rotation")
-        self.logger.debug("%s.find_shift(%s)" % (self.cor_finder.__class__.__name__, str(self.cor_exec_kwargs)))
-        shift = self.cor_finder.find_shift(
-            self.radios[0], np.fliplr(self.radios[1]), *self.cor_exec_args, **self.cor_exec_kwargs
-        )
-        return self.shape[1] / 2 + shift
+        # All find_shift() methods in self.search_methods have the same API with "img_1" and "img_2"
+        cor_exec_kwargs = update_func_kwargs(self.cor_finder.find_shift, self.cor_options)
+        cor_exec_kwargs["return_relative_to_middle"] = False
+        if self._lookup_side is not None:
+            cor_exec_kwargs["side"] = self._lookup_side
+        self.logger.debug("%s.find_shift(%s)" % (self.cor_finder.__class__.__name__, str(cor_exec_kwargs)))
+        shift = self.cor_finder.find_shift(self.radios[0], np.fliplr(self.radios[1]), **cor_exec_kwargs)
+        return shift
 
 
 # alias
@@ -329,16 +243,26 @@ class SinoCORFinder(CORFinderBase):
         },
         "sino-sliding-window": {
             "class": CenterOfRotationSlidingWindow,
-            "default_args": ["right"],
         },
         "sino-growing-window": {
             "class": CenterOfRotationGrowingWindow,
         },
-        "fourier-angles": {"class": CenterOfRotationFourierAngles, "default_args": [None, "center"]},
+        "fourier-angles": {"class": CenterOfRotationFourierAngles},
+        "vo": {
+            "class": CenterOfRotationVo,
+        },
     }
 
     def __init__(
-        self, method, dataset_info, slice_idx="middle", subsampling=10, do_flatfield=True, cor_options=None, logger=None
+        self,
+        method,
+        dataset_info,
+        do_flatfield=True,
+        take_log=True,
+        cor_options=None,
+        logger=None,
+        slice_idx="middle",
+        subsampling=10,
     ):
         """
         Initialize a SinoCORFinder object.
@@ -355,20 +279,15 @@ class SinoCORFinder(CORFinderBase):
             subsampling strategy when building sinograms.
             As building the complete sinogram from raw projections might be tedious, the reading is done with subsampling.
             A positive integer value means the subsampling step (i.e `projections[::subsampling]`).
-            A negative integer value means we take -subsampling projections in total.
-            A float value indicates the angular step in DEGREES.
         """
         super().__init__(method, dataset_info, do_flatfield=do_flatfield, cor_options=cor_options, logger=logger)
-        self._check_360()
         self._set_slice_idx(slice_idx)
         self._set_subsampling(subsampling)
         self._load_raw_sinogram()
         self._flatfield(do_flatfield)
-        self._get_sinogram()
+        self._get_sinogram(take_log)
 
     def _check_360(self):
-        if self.dataset_info.dataset_scanner.scan_range == 360:
-            return
         if not is_fullturn_scan(self.dataset_info.rotation_angles):
             raise ValueError("Sinogram-based Center of Rotation estimation can only be used for 360 degrees scans")
 
@@ -382,50 +301,47 @@ class SinoCORFinder(CORFinderBase):
 
     def _set_subsampling(self, subsampling):
         projs_idx = sorted(self.dataset_info.projections.keys())
+        self.subsampling = None
         if is_int(subsampling):
             if subsampling < 0:  # Total number of angles
-                n_angles = -subsampling
-                indices_float = np.linspace(projs_idx[0], projs_idx[-1], n_angles, endpoint=True)
-                self.projs_indices = np.round(indices_float).astype(np.int32).tolist()
-            else:  # Subsampling step
+                raise NotImplementedError
+            else:
                 self.projs_indices = projs_idx[::subsampling]
                 self.angles = self.dataset_info.rotation_angles[::subsampling]
+                self.subsampling = subsampling
         else:  # Angular step
             raise NotImplementedError()
 
     def _load_raw_sinogram(self):
         if self.slice_idx is None:
             raise ValueError("Unknow slice index")
-        # Subsample projections
-        files = {}
-        for idx in self.projs_indices:
-            files[idx] = self.dataset_info.projections[idx]
-        self.files = files
-        self.data_reader = ChunkReader(
-            self.files,
-            sub_region=(None, None, self.slice_idx, self.slice_idx + 1),
-            convert_float=True,
-        )
-        self.data_reader.load_files()
-        self._radios = self.data_reader.files_data
+        reader_kwargs = {
+            "sub_region": (slice(None, None, self.subsampling), slice(self.slice_idx, self.slice_idx + 1), slice(None))
+        }
+        if self.dataset_info.kind == "edf":
+            reader_kwargs = {"n_reading_threads": get_num_threads()}
+        self.data_reader = self.dataset_info.get_reader(**reader_kwargs)
+        self._radios = self.data_reader.load_data()
 
     def _flatfield(self, do_flatfield):
         self.do_flatfield = bool(do_flatfield)
         if not self.do_flatfield:
             return
-        flatfield = FlatFieldDataUrls(
+        flats = {k: arr[self.slice_idx : self.slice_idx + 1, :] for k, arr in self.dataset_info.flats.items()}
+        darks = {k: arr[self.slice_idx : self.slice_idx + 1, :] for k, arr in self.dataset_info.darks.items()}
+        flatfield = FlatField(
             self._radios.shape,
-            self.dataset_info.flats,
-            self.dataset_info.darks,
+            flats,
+            darks,
             radios_indices=self.projs_indices,
-            sub_region=(None, None, self.slice_idx, self.slice_idx + 1),
         )
         flatfield.normalize_radios(self._radios)
 
-    def _get_sinogram(self):
-        log = Log(self._radios.shape, clip_min=1e-6, clip_max=10.0)
+    def _get_sinogram(self, take_log):
         sinogram = self._radios[:, 0, :].copy()
-        log.take_logarithm(sinogram)
+        if take_log:
+            log = Log(self._radios.shape, clip_min=1e-6, clip_max=10.0)
+            log.take_logarithm(sinogram)
         self.sinogram = sinogram
 
     @staticmethod
@@ -440,10 +356,29 @@ class SinoCORFinder(CORFinderBase):
 
     def find_cor(self):
         self.logger.info("Estimating center of rotation")
-        self.logger.debug("%s.find_shift(%s)" % (self.cor_finder.__class__.__name__, str(self.cor_exec_kwargs)))
-        img_1, img_2 = self._split_sinogram(self.sinogram)
-        shift = self.cor_finder.find_shift(img_1, np.fliplr(img_2), *self.cor_exec_args, **self.cor_exec_kwargs)
-        return self.shape[1] / 2 + shift
+
+        cor_exec_kwargs = update_func_kwargs(self.cor_finder.find_shift, self.cor_options)
+        cor_exec_kwargs["return_relative_to_middle"] = False
+
+        if self._lookup_side is not None:
+            cor_exec_kwargs["side"] = self._lookup_side
+
+        if self.method == "fourier-angles":
+            cor_exec_args = [self.sinogram]
+            cor_exec_kwargs["angles"] = self.dataset_info.rotation_angles
+        elif self.method == "vo":
+            cor_exec_args = [self.sinogram]
+            cor_exec_kwargs["halftomo"] = self.dataset_info.is_halftomo
+            cor_exec_kwargs["is_360"] = is_fullturn_scan(self.dataset_info.rotation_angles)
+        else:
+            # For these methods relying on find_shift() with two images, the sinogram needs to be split in two
+            img_1, img_2 = self._split_sinogram(self.sinogram)
+            cor_exec_args = [img_1, np.fliplr(img_2)]
+
+        self.logger.debug("%s.find_shift(%s)" % (self.cor_finder.__class__.__name__, str(cor_exec_kwargs)))
+        shift = self.cor_finder.find_shift(*cor_exec_args, **cor_exec_kwargs)
+
+        return shift
 
 
 # alias
@@ -472,14 +407,14 @@ class CompositeCORFinder(CORFinderBase):
             "class": CenterOfRotation,  # Hack. Not used. Everything is done in the find_cor() func.
         }
     }
-    _default_cor_options = {"low_pass": 0.4, "high_pass": 10, "side": "center", "near_pos": 0, "near_width": 20}
+    _default_cor_options = {"low_pass": 0.4, "high_pass": 10, "side": "near", "near_pos": 0, "near_width": 40}
 
     def __init__(
         self,
         dataset_info,
         oversampling=4,
         theta_interval=5,
-        n_subsampling_y=10,
+        n_subsampling_y=40,
         take_log=True,
         cor_options=None,
         spike_threshold=0.04,
@@ -530,8 +465,6 @@ class CompositeCORFinder(CORFinderBase):
             if useful_span < np.pi:
                 theta_interval = theta_interval * useful_span / np.pi
 
-        # self._get_cor_options(cor_options)
-
         self.take_log = take_log
         self.ovs = oversampling
         self.theta_interval = theta_interval
@@ -566,16 +499,15 @@ class CompositeCORFinder(CORFinderBase):
 
         self.absolute_indices = sorted(self.dataset_info.projections.keys())
 
-        my_flats = load_images_from_dataurl_dict(self.dataset_info.flats)
+        my_flats = self.dataset_info.flats
 
         if my_flats is not None and len(list(my_flats.keys())):
             self.use_flat = True
-            self.flatfield = FlatFieldDataUrls(
+            self.flatfield = FlatField(
                 (len(self.absolute_indices), self.sy, self.sx),
                 self.dataset_info.flats,
                 self.dataset_info.darks,
                 radios_indices=self.absolute_indices,
-                dtype=np.float64,
             )
         else:
             self.use_flat = False
@@ -678,7 +610,7 @@ class CompositeCORFinder(CORFinderBase):
                     other_i = sorted_angle_indexes[0]
                 elif insertion_point == len(sorted_all_angles):
                     other_i = sorted_angle_indexes[insertion_point - 1]
-                radio2 = self.get_radio(self.absolute_indices[other_i])
+                radio2 = self.get_radio(self.absolute_indices[other_i])  # pylint: disable=E0606
 
             self.sino[irad : irad + radio1.shape[0], :] = self._oversample(radio1)
             self.sino[
@@ -714,31 +646,38 @@ class CompositeCORFinder(CORFinderBase):
         tmp_sy, ovsd_sx = radio1.shape
         assert orig_sy == tmp_sy and orig_ovsd_sx == ovsd_sx, "this should not happen"
 
-        if self.cor_options["side"] == "center":
+        cor_side = self.cor_options["side"]
+        if cor_side == "center":
             overlap_min = max(round(ovsd_sx - ovsd_sx / 3), 4)
             overlap_max = min(round(ovsd_sx + ovsd_sx / 3), 2 * ovsd_sx - 4)
-        elif self.cor_options["side"] == "right":
+        elif cor_side == "right":
             overlap_min = max(4, self.ovs * self.high_pass * 3)
             overlap_max = ovsd_sx
-        elif self.cor_options["side"] == "left":
+        elif cor_side == "left":
             overlap_min = ovsd_sx
             overlap_max = min(2 * ovsd_sx - 4, 2 * ovsd_sx - self.ovs * self.ovs * self.high_pass * 3)
-        elif self.cor_options["side"] == "all":
+        elif cor_side == "all":
             overlap_min = max(4, self.ovs * self.high_pass * 3)
             overlap_max = min(2 * ovsd_sx - 4, 2 * ovsd_sx - self.ovs * self.ovs * self.high_pass * 3)
-
-        elif self.cor_options["side"] == "near":
+        elif is_scalar(cor_side):
+            near_pos = cor_side
+            near_width = self.cor_options["near_width"]
+            overlap_min = max(4, ovsd_sx - 2 * self.ovs * (near_pos + near_width))
+            overlap_max = min(2 * ovsd_sx - 4, ovsd_sx - 2 * self.ovs * (near_pos - near_width))
+        # COMPAT.
+        elif cor_side == "near":
+            deprecation_warning(
+                "using side='near' is deprecated, use side=<a scalar> instead",
+                do_print=True,
+                func_name="composite_near_pos",
+            )
             near_pos = self.cor_options["near_pos"]
             near_width = self.cor_options["near_width"]
-
             overlap_min = max(4, ovsd_sx - 2 * self.ovs * (near_pos + near_width))
             overlap_max = min(2 * ovsd_sx - 4, ovsd_sx - 2 * self.ovs * (near_pos - near_width))
-
+        # ---
         else:
-            message = f""" The cor options "side" can only have one of the three possible values ["","",""].
-            But it has the value "{self.cor_options["side"]}" instead
-            """
-            raise ValueError(message)
+            raise ValueError("Invalid option 'side=%s'" % self.cor_options["side"])
 
         if overlap_min > overlap_max:
             message = f""" There is no safe search range in find_cor once the margins corresponding to the high_pass filter are discarded.
@@ -938,13 +877,12 @@ class DetectorTiltEstimator:
     def _init_flatfield(self):
         if not (self.do_flatfield):
             return
-        self.flatfield = FlatFieldDataUrls(
+        self.flatfield = FlatField(
             self.radios.shape,
             flats=self.dataset_info.flats,
             darks=self.dataset_info.darks,
             radios_indices=self.radios_indices,
             interpolation="linear",
-            convert_float=True,
         )
 
     def _apply_flatfield(self):