Flat-field correction (without detector plane handling)

3 files changed, 33 insertions, 14 deletions

diff --git a/tests/roof/config.py b/tests/roof/config.py
index e085ed8..8f72014 100644
--- a/tests/roof/config.py
+++ b/tests/roof/config.py
@@ -13,10 +13,15 @@ class RoofConfig:
 
         self.path = self.get_opt('data', 'base_path', './')
         self.planes = self.get_opt('hardware', 'planes', 1)
-        self.modules = self.get_opt('hardware', 'modules', None)
-        self.streams = self.get_opt('network', 'streams', 1 if self.modules is None else self.modules)
+        self.modules = self.get_opt('hardware', 'modules', 0)
+        self.streams = self.get_opt('network', 'streams', self.modules if self.modules else 1)
         self.bit_depth = self.get_opt('hardware', 'bit_depth', 8)
 
+        self.sample_rate = self.get_opt('hardware', 'sample_rate', 0)
+        self.imaging_rate = self.get_opt('hardware', 'imaging_rate', 0)
+        self.fan_bins = self.modules * self.get_opt('hardware', 'channels_per_module', 16)
+        self.fan_projections = self.get_opt('hardware', 'samples_per_rotation', (self.sample_rate / self.imaging_rate) if (self.imaging_rate and self.sample_rate) else 1000)
+        
         if self.args.number is None: self.args.number = 0 if self.args.benchmark else self.planes
 
         # Consistency and default mode
@@ -46,6 +51,13 @@ class RoofConfig:
         if subpath is None: raise "Unknown data type %s is requested" % subpath
         return subpath if subpath.startswith('/') else  self.path + '/' + subpath
 
+    def get_roof_glob(self, data_type):
+        path = self.get_roof_path(data_type)
+        if path is None: return None
+
+        globre = re.compile('%.?(\d+)[luid]')
+        return globre.sub('*', path)
+
     def get_writer_type(self):
         return None if self.args.benchmark else self.args.write if self.args.write else 'raw_sinograms'
     
diff --git a/tests/roof/defaults.py b/tests/roof/defaults.py
index eed3fe5..f6b67f6 100644
--- a/tests/roof/defaults.py
+++ b/tests/roof/defaults.py
@@ -15,7 +15,7 @@ roof_default_paths = {
 #}
 
 roof_filters = {
- 'correction':                  [ "flat-field-correct" ],
+ 'correction':                  [ "roof-flat-field-correct" ],
  'reconstruction':              [ "roof-fan2par", "fft", "filter", "ifft", "backproject" ],
  'control':                     [ ],
  'visualization':               [ ]
@@ -23,9 +23,9 @@ roof_filters = {
 
 roof_data_types = {
  'correction': {
-    'flat_fields':              'flat-field-correct',
-    'dark_fields':              'flat-field-correct',
-    'raw_sinograms':            'flat-field-correct',
+    'flat_fields':              'roof-flat-field-correct',
+    'dark_fields':              'roof-flat-field-correct',
+    'raw_sinograms':            'roof-flat-field-correct',
     'fan_sinograms':            None    
   },
   'reconstruction': {
diff --git a/tests/roof/graph.py b/tests/roof/graph.py
index c34a3ed..495c629 100644
--- a/tests/roof/graph.py
+++ b/tests/roof/graph.py
@@ -39,11 +39,16 @@ class RoofGraph(RoofConfig):
         if (re.compile('roof').match(name)): kwargs.update(config = self.config_file)
         return self.save_task(stage, name, self.get_task(name, **kwargs))
     
+    def get_read_task(self, path):
+        params = { 'raw_width': self.fan_bins, 'raw_height': self.fan_projections, 'raw_bitdepth': self.bit_depth } if re.search('\.raw$', path) else { }
+        params['path'] = path
+        return self.get_task('read', **params)
+
     def get_reader(self):
         first = self.get_opt('data', 'first_file_number', 1)
         if self.args.read:
             # Reconstruction from standard UFO files
-            path = self.get_roof_path(self.args.read)
+            path = self.get_roof_glob(self.args.read)
             step = 1
             if (self.args.plane is not None) and (self.args.plane > 0): 
                 first += self.args.plane - 1;
@@ -54,6 +59,7 @@ class RoofGraph(RoofConfig):
                 params['number'] = self.args.number
             
             print ("Reading {} data from {}".format(self.args.read,path))
+            # FIXME: handle raw data parameters
             return self.get_task('read', **params)
         else:
             path = None
@@ -85,16 +91,16 @@ class RoofGraph(RoofConfig):
             write = self.get_task('write', filename=path)
         return write
 
-    def get_correction_flat_field_correct(self, head):
+    def get_correction_roof_flat_field_correct(self, head):
         # Standard UFO reconstruction stack distinguish flat/dark-fields recorded before and after experiment. We only do 'before experiment' part.
-        darks = self.get_roof_path('dark_fields')
+        darks = self.get_roof_glob('dark_fields')
         n_darks = len(get_filenames(darks))
         if n_darks == 0: raise FileNotFoundError("Dark fields are not found in {}".format(darks))
-        flats = self.get_roof_path('falt_fields')
+        flats = self.get_roof_glob('flat_fields')
         n_flats = len(get_filenames(flats))
         if n_flats == 0: raise FileNotFoundError("Flat fields are not found in {}".format(flats))
-        dark_reader = self.get_task('read', path = darks)
-        flat_reader = self.get_task('read', path = flats)
+        dark_reader = self.get_read_task(path = darks)
+        flat_reader = self.get_read_task(path = flats)
 
         # We are using standard get_task here because this is too generic plugin to allow config-based customization
         mode = self.get_opt('correction', 'aggregation', 'average')
@@ -116,7 +122,8 @@ class RoofGraph(RoofConfig):
         else:
             raise ValueError('Invalid reduction mode')
 
-        ffc = self.get_task('flat-field-correct') # dark_scale=args.dark_scale, absorption_correct=args.absorptivity, fix_nan_and_inf=args.fix_nan_and_inf)
+        #ffc = self.get_task('flat-field-correct', fix_nan_and_inf=False, absorption_correct=True)
+        ffc = self.get_task('roof-flat-field-correct', fix_nan_and_inf=False)  # absorption_correct=True)
         self.graph.connect_nodes_full(head, ffc, 0)
         self.graph.connect_nodes_full(dark_reduced, ffc, 1)
         self.graph.connect_nodes_full(flat_reduced, ffc, 2)
@@ -146,7 +153,7 @@ class RoofGraph(RoofConfig):
             if method in dir(self):
                 f = getattr(self, method)(head)
             else:
-                f = self.get_processor_task(stage, filters[pos])
+                f = self.get_processor_task(stage, filters[i])
                 graph.connect_nodes(head, f)
             head = f