dataset_analyzer: prepare force_flatfield

nabu/resources/dataset_analyzer.py

@@ -24,8 +24,34 @@ class DatasetAnalyzer(object):
     """
     Base class for datasets analyzers.
     """
-    def __init__(self, location):
+    def __init__(self, location, processes_file=None, extra_options=None):
+        """
+        Initialize a Dataset analyzer.
+
+        Parameters
+        ----------
+        location: str
+            Dataset location (directory or file name)
+        processes_file: str, optional
+            Processes file providing supplementary information (ex. pre-processed data)
+        extra_options: dict, optional
+            Extra options on how to interpret the dataset.
+            Available options are the following:
+              - force_flatfield
+        """
         self.location = location
+        self.processes_file = processes_file
+        self._set_extra_options(extra_options)
+
+
+    def _set_extra_options(self, extra_options):
+        if extra_options is None:
+            extra_options = {}
+        advanced_options = {
+            "force_flatfield": False,
+        }
+        advanced_options.update(extra_options)
+        self.extra_options = advanced_options
 
 
     def _init_dataset_scan(self, filetype, **kwargs):
@@ -147,6 +173,17 @@ class HDF5DatasetAnalyzer(DatasetAnalyzer):
         self._rot_angles = None
 
 
+    def _get_flats_darks(self):
+        if len(self.flats) == 0:
+            # No flats at all in the dataset. Do nothing.
+            return
+        if self._load_flats_from_tomwer():
+            # Loaded from tomwer_processes.h5
+            return
+        # Otherwise load or compute flats/darks with nabu
+        self._compute_or_load_flats()
+
+
     def _load_flats_from_tomwer(self, tomwer_processes_fname=None):
         tomwer_processes_fname = tomwer_processes_fname or "tomwer_processes.h5"
         tomwer_processes_file = os.path.join(self.dataset_scanner.path, "tomwer_processes.h5")
@@ -165,7 +202,9 @@ class HDF5DatasetAnalyzer(DatasetAnalyzer):
 
 
     def _compute_or_load_flats(self):
-        processes_file = os.path.join(self.dataset_scanner.path, "nabu_processes.h5")
+        processes_file = self.processes_file
+        if processes_file is None:
+            processes_file = os.path.join(self.dataset_scanner.path, "nabu_processes.h5")
         lookup_files = [
             DataUrl(file_path=processes_file, data_path="/entry/flat_field")
         ]
@@ -180,17 +219,6 @@ class HDF5DatasetAnalyzer(DatasetAnalyzer):
         self._nxflatfield.get_final_images()
 
 
-    def _get_flats_darks(self):
-        if len(self.flats) == 0:
-            # No flats at all in the dataset. Do nothing.
-            return
-        if self._load_flats_from_tomwer():
-            # Loaded from tomwer_processes.h5
-            return
-        # Otherwise load or compute flats/darks with nabu
-        self._compute_or_load_flats()
-
-
     def _get_rotation_angles(self):
         if self._rot_angles is None:
             angles = np.array(self.dataset_scanner.rotation_angle)
@@ -215,15 +243,15 @@ class HDF5DatasetAnalyzer(DatasetAnalyzer):
         return self._get_dataset_hdf5_url()
 
 
-def analyze_dataset(dataset_path):
+def analyze_dataset(dataset_path, extra_options=None):
     if not(os.path.isdir(dataset_path)):
         if not(os.path.isfile(dataset_path)):
             raise ValueError("Error: %s no such file or directory" % dataset_path)
-        if os.path.splitext(dataset_path)[-1] not in [".hdf5", ".h5", ".nx"]:
+        if not(is_hdf5_extension(os.path.splitext(dataset_path)[-1])):
             raise ValueError("Error: expected a HDF5 file")
         dataset_analyzer_class = HDF5DatasetAnalyzer
     else: # directory -> assuming EDF
         dataset_analyzer_class = EDFDatasetAnalyzer
-    dataset_structure = dataset_analyzer_class(dataset_path)
-    return dataset_structure
+    dataset_struct = dataset_analyzer_class(dataset_path, extra_options=extra_options)
+    return dataset_struct