summaryrefslogtreecommitdiffstats
path: root/src/processors/FBP_Stacked.py
blob: b6310a127b2d0459e5a725e197cbb7d2a3776437 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
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