diff options
Diffstat (limited to 'src/processors/FBP_Stacked.py')
-rw-r--r-- | src/processors/FBP_Stacked.py | 85 |
1 files changed, 85 insertions, 0 deletions
diff --git a/src/processors/FBP_Stacked.py b/src/processors/FBP_Stacked.py new file mode 100644 index 0000000..b6310a1 --- /dev/null +++ b/src/processors/FBP_Stacked.py @@ -0,0 +1,85 @@ +from cil.framework import DataProcessor, ImageData +import numpy as np + +import gi +gi.require_version('Ufo','0.0') +from gi.repository import Ufo + +class FBP_Stacked(DataProcessor): + def __init__(self, volume_geometry, + sinogram_geometry, + precision_mode='single', + stack_num=2): + + kwargs = { + 'volume_geometry': volume_geometry, + 'sinogram_geometry': sinogram_geometry, + 'precision_mode': precision_mode, + 'stack_num': stack_num} + + super(FBP_Stacked, self).__init__(**kwargs) + + self.precision_mode = precision_mode + self.stack_num = stack_num + self.set_ImageGeometry(volume_geometry) + self.set_AcquisitionGeometry(sinogram_geometry) + + def check_input(self, dataset): + if self.sinogram_geometry.dimension == '2D': + raise ValueError("Expected input dimensions is 3, got {0}" \ + .format(dataset.number_of_dimensions)) + return True + + def set_ImageGeometry(self, volume_geometry): + self.volume_geometry = volume_geometry + + def set_AcquisitionGeometry(self, sinogram_geometry): + self.sinogram_geometry = sinogram_geometry + + def process(self, **kwargs): + # Get DATA + DATA = self.get_input() + DATA = DATA.as_array() + + pm = Ufo.PluginManager() + graph = Ufo.TaskGraph() + scheduler = Ufo.Scheduler() + + read = pm.get_task('memory-in') + read.props.pointer = DATA.__array_interface__['data'][0] + read.props.width = DATA.shape[1] + read.props.height = DATA.shape[0] + read.props.number = DATA.shape[2] + read.props.bitdepth = 32 + + fft = pm.get_task('fft') + fft.props.dimensions = 1 + + filter = pm.get_task('filter') + + ifft = pm.get_task('ifft') + ifft.props.dimensions = 1 + + stack = pm.get_task('stack') + stack.props.number = self.stack_num + + stacked_bp = pm.get_task('stacked-backproject') + stacked_bp.props.axis_pos = DATA.shape[1] / 2 + stacked_bp.props.precision_mode = self.precision_mode + + vol_arr = np.zeros(self.volume_geometry.shape, dtype=np.float32) + + write = pm.get_task('memory-out') + write.props.pointer = vol_arr.__array_interface__['data'][0] + write.props.max_size = vol_arr.nbytes + + graph.connect_nodes(read, fft) + graph.connect_nodes(fft, filter) + graph.connect_nodes(filter, ifft) + graph.connect_nodes(ifft, stack) + graph.connect_nodes(stack, stacked_bp) + graph.connect_nodes(stacked_bp, write) + + scheduler.run(graph) + vol_arr = ImageData(vol_arr, deep_copy=False, geometry=self.volume_geometry.copy(), suppress_warning=True) + return vol_arr |