savu.plugins.reconstructions.tomobar.tomobar_recon_3D.TomobarRecon3d._set_max_frames() - Code Metrics - Inspection of "Merge pull request #911 from DiamondLightSource/da..." - DiamondLightSource/Savu - Measure and Improve Code Quality continuously with Scrutinizer

Test Failed

Push — master ( 877aea...3bbc8e )

by Daniil

created 2022-03-17 18:27 UTC

TomobarRecon3d._set_max_frames() A

↳ Parent: savu.plugins.reconstructions.tomobar.tomobar_recon_3D

Complexity

Conditions

Size

Total Lines	2
Code Lines	2

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	1
eloc	2
nop	2
dl	0
loc	2
rs	10
c	0
b	0
f	0

# Copyright 2019 Diamond Light Source Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
.. module:: tomobar_recon_3D
   :platform: Unix
   :synopsis: A wrapper around TOmographic MOdel-BAsed Reconstruction (ToMoBAR) software \
   for direct and advanced iterative image reconstruction using _3D_ capabilities of regularisation. \
   This plugin will divide 3D projection data into overlapping subsets using padding.

.. moduleauthor:: Daniil Kazantsev <[email protected]>
"""

from savu.plugins.reconstructions.base_recon import BaseRecon
from savu.plugins.driver.gpu_plugin import GpuPlugin

import numpy as np
from tomobar.methodsIR import RecToolsIR
from tomobar.methodsDIR import RecToolsDIR
from savu.plugins.utils import register_plugin
from savu.core.iterate_plugin_group_utils import enable_iterative_loop, \
    check_if_end_plugin_in_iterate_group, setup_extra_plugin_data_padding


@register_plugin
class TomobarRecon3d(BaseRecon, GpuPlugin):

    def __init__(self):
        super(TomobarRecon3d, self).__init__("TomobarRecon3d")
        self.Vert_det = None
        self.pad = None

    @setup_extra_plugin_data_padding

    def set_filter_padding(self, in_pData, out_pData):
        self.pad = self.parameters['padding']
        in_data = self.get_in_datasets()[0]
        det_y = in_data.get_data_dimension_by_axis_label('detector_y')
        pad_det_y = '%s.%s' % (det_y, self.pad)
        pad_dict = {'pad_directions': [pad_det_y], 'pad_mode': 'edge'}
        in_pData[0].padding = pad_dict
        out_pData[0].padding = pad_dict
        if len(self.get_in_datasets()) > 1:
            in_pData[1].padding = pad_dict

    @enable_iterative_loop
    def setup(self):
        in_dataset = self.get_in_datasets()[0]
        procs = self.exp.meta_data.get("processes")
        procs = len([i for i in procs if 'GPU' in i]) # calculates the total number of GPU processes
        dim = in_dataset.get_data_dimension_by_axis_label('detector_y')
        nSlices = int(np.ceil(in_dataset.get_shape()[dim] / float(procs)))
        # calculate the amount of slices than would fit the GPU memory
        gpu_available_mb = self.get_gpu_memory()[0]/procs  # get the free GPU memory of a first device if many
        det_x_dim = in_dataset.get_shape()[in_dataset.get_data_dimension_by_axis_label('detector_x')]
        rot_angles_dim = in_dataset.get_shape()[in_dataset.get_data_dimension_by_axis_label('rotation_angle')]
        slice_size_mbbytes = int(np.ceil(((det_x_dim * det_x_dim) * 1024 * 4) / (1024 ** 3)))

        if self.parameters['reconstruction_method'] == 'FISTA3D':
            # calculate the GPU memory required based on 3D regularisation restrictions (avoiding CUDA-error)
            if 'ROF_TV' in self.parameters['regularisation_method']:
                slice_size_mbbytes *= 8
            if 'FGP_TV' in self.parameters['regularisation_method']:
                slice_size_mbbytes *= 12
            if 'SB_TV' in self.parameters['regularisation_method']:
                slice_size_mbbytes *= 10
            if 'PD_TV' in self.parameters['regularisation_method']:
                slice_size_mbbytes *= 8
            if 'LLT_ROF' in self.parameters['regularisation_method']:
                slice_size_mbbytes *= 12
            if 'TGV' in self.parameters['regularisation_method']:
                slice_size_mbbytes *= 15
            if 'NDF' in self.parameters['regularisation_method']:
                slice_size_mbbytes *= 5
            if 'Diff4th' in self.parameters['regularisation_method']:
                slice_size_mbbytes *= 5
            if 'NLTV' in self.parameters['regularisation_method']:
                slice_size_mbbytes *= 8

        slices_fit_total = int(gpu_available_mb / slice_size_mbbytes) - 2*self.parameters['padding']
        if nSlices > slices_fit_total:
            nSlices = slices_fit_total
        self._set_max_frames(nSlices)
        # get experimental metadata of projection_shifts
        if 'projection_shifts' in list(self.exp.meta_data.dict.keys()):
            self.projection_shifts = self.exp.meta_data.dict['projection_shifts']
        super(TomobarRecon3d, self).setup()

    def pre_process(self):
        in_pData = self.get_plugin_in_datasets()[0]
        self.det_dimX_ind = in_pData.get_data_dimension_by_axis_label('detector_x')
        self.det_dimY_ind = in_pData.get_data_dimension_by_axis_label('detector_y')
        #  getting the value for padded vertical detector
        self.Vert_det = in_pData.get_shape()[self.det_dimY_ind] + 2 * self.pad

        # extract given parameters into dictionaries suitable for ToMoBAR input
        self._data_ = {'OS_number': self.parameters['algorithm_ordersubsets'],
                       'huber_threshold': self.parameters['data_Huber_thresh'],
                       'ringGH_lambda': self.parameters['data_full_ring_GH'],
                       'ringGH_accelerate': self.parameters['data_full_ring_accelerator_GH']}

        self._algorithm_ = {'iterations': self.parameters['algorithm_iterations'],
                            'nonnegativity': self.parameters['algorithm_nonnegativity'],
                            'mask_diameter': self.parameters['algorithm_mask'],
                            'verbose': self.parameters['algorithm_verbose']}

        self._regularisation_ = {'method': self.parameters['regularisation_method'],
                                 'regul_param': self.parameters['regularisation_parameter'],
                                 'iterations': self.parameters['regularisation_iterations'],
                                 'device_regulariser': self.parameters['regularisation_device'],
                                 'edge_threhsold': self.parameters['regularisation_edge_thresh'],
                                 'time_marching_step': self.parameters['regularisation_timestep'],
                                 'regul_param2': self.parameters['regularisation_parameter2'],
                                 'PD_LipschitzConstant': self.parameters['regularisation_PD_lip'],
                                 'NDF_penalty': self.parameters['regularisation_NDF_penalty'],
                                 'methodTV': self.parameters['regularisation_methodTV']}

    def process_frames(self, data):
        cor, angles, self.vol_shape, init = self.get_frame_params()
        self.anglesRAD = np.deg2rad(angles.astype(np.float32))
        projdata3D = data[0].astype(np.float32)
        dim_tuple = np.shape(projdata3D)
        self.Horiz_det = dim_tuple[self.det_dimX_ind]
        half_det_width = 0.5 * self.Horiz_det
        projdata3D[projdata3D > 10 ** 15] = 0.0
        projdata3D = np.swapaxes(projdata3D, 0, 1)
        self._data_.update({'projection_norm_data': projdata3D})

        # dealing with projection shifts and the CoR
        cor_astra = half_det_width - np.mean(cor)
        CenterOffset = cor_astra.item() - 0.5
        if np.sum(self.projection_shifts) != 0.0:
            CenterOffset = np.zeros(np.shape(self.projection_shifts))
            CenterOffset[:, 0] = (cor_astra.item() - 0.5) - self.projection_shifts[:, 0]
            CenterOffset[:, 1] = -self.projection_shifts[:, 1] - 0.5

        if self.parameters['reconstruction_method'] == 'FISTA3D':
            if self.parameters['regularisation_method'] == 'PD_TV':
	            self._regularisation_.update({'device_regulariser': self.parameters['GPU_index']})        
            # if one selects PWLS or SWLS models then raw data is also required (2 inputs)
            if (self.parameters['data_fidelity'] == 'PWLS') or (self.parameters['data_fidelity'] == 'SWLS'):
                rawdata3D = data[1].astype(np.float32)
                rawdata3D[rawdata3D > 10 ** 15] = 0.0
                rawdata3D = np.swapaxes(rawdata3D, 0, 1) / np.max(np.float32(rawdata3D))
                self._data_.update({'projection_raw_data': rawdata3D})
                self._data_.update({'beta_SWLS': self.parameters['data_beta_SWLS'] * np.ones(self.Horiz_det)})

            # set parameters and initiate a TomoBar class object for FISTA reconstruction
            RectoolsIter = RecToolsIR(DetectorsDimH=self.Horiz_det,  # DetectorsDimH # detector dimension (horizontal)
                                       DetectorsDimV=self.Vert_det,   # DetectorsDimV # detector dimension (vertical) for 3D case only
                                       CenterRotOffset=CenterOffset,  # The center of rotation combined with  the shift offsets
                                       AnglesVec=-self.anglesRAD,  # the vector of angles in radians
                                       ObjSize=self.vol_shape[0],  # a scalar to define the reconstructed object dimensions
                                       datafidelity=self.parameters['data_fidelity'], # data fidelity, choose LS, PWLS, SWLS
                                       device_projector=self.parameters['GPU_index'])

            # Run FISTA reconstruction algorithm here
            recon = RectoolsIter.FISTA(self._data_, self._algorithm_, self._regularisation_)

        if self.parameters['reconstruction_method'] == 'FBP3D':
            RectoolsDIR = RecToolsDIR(DetectorsDimH=self.Horiz_det,  # DetectorsDimH # detector dimension (horizontal)
                                       DetectorsDimV=self.Vert_det,  # DetectorsDimV # detector dimension (vertical) for 3D case only
                                       CenterRotOffset=CenterOffset,  # The center of rotation combined with the shift offsets
                                       AnglesVec=-self.anglesRAD,  # the vector of angles in radians
                                       ObjSize=self.vol_shape[0],  # a scalar to define the reconstructed object dimensions
                                       device_projector=self.parameters['GPU_index'])

            recon = RectoolsDIR.FBP(projdata3D) #perform FBP3D

        if self.parameters['reconstruction_method'] == 'CGLS3D':
            # set parameters and initiate a TomoBar class object for FISTA reconstruction
            RectoolsIter = RecToolsIR(DetectorsDimH=self.Horiz_det,  # DetectorsDimH # detector dimension (horizontal)
                                       DetectorsDimV=self.Vert_det,   # DetectorsDimV # detector dimension (vertical) for 3D case only
                                       CenterRotOffset=CenterOffset,  # The center of rotation combined with  the shift offsets
                                       AnglesVec=-self.anglesRAD,  # the vector of angles in radians
                                       ObjSize=self.vol_shape[0],  # a scalar to define the reconstructed object dimensions
                                       datafidelity=self.parameters['data_fidelity'], # data fidelity, choose LS, PWLS, SWLS
                                       device_projector=self.parameters['GPU_index'])

            # Run CGLS reconstruction algorithm here
            recon = RectoolsIter.CGLS(self._data_, self._algorithm_)

        recon = np.swapaxes(recon, 0, 1)

        return recon

    def nInput_datasets(self):
        return max(len(self.parameters['in_datasets']), 1)

    # total number of output datasets
    def nOutput_datasets(self):
        if check_if_end_plugin_in_iterate_group(self.exp):
            return 2
        else:
            return 1

    # total number of output datasets that are clones
    def nClone_datasets(self):
        if check_if_end_plugin_in_iterate_group(self.exp):
            return 1
        else:
            return 0

    def _set_max_frames(self, frames):
        self._max_frames = frames


1		# Copyright 2019 Diamond Light Source Ltd.
2		#
3		# Licensed under the Apache License, Version 2.0 (the "License");
4		# you may not use this file except in compliance with the License.
5		# You may obtain a copy of the License at
6		#
7		# http://www.apache.org/licenses/LICENSE-2.0
8		#
9		# Unless required by applicable law or agreed to in writing, software
10		# distributed under the License is distributed on an "AS IS" BASIS,
11		# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12		# See the License for the specific language governing permissions and
13		# limitations under the License.
14
15		"""
16		.. module:: tomobar_recon_3D
17		:platform: Unix
18		:synopsis: A wrapper around TOmographic MOdel-BAsed Reconstruction (ToMoBAR) software \
19		for direct and advanced iterative image reconstruction using _3D_ capabilities of regularisation. \
20		This plugin will divide 3D projection data into overlapping subsets using padding.
21
22		.. moduleauthor:: Daniil Kazantsev <[email protected]>
23		"""
24
25		from savu.plugins.reconstructions.base_recon import BaseRecon
26		from savu.plugins.driver.gpu_plugin import GpuPlugin
27
28		import numpy as np
29		from tomobar.methodsIR import RecToolsIR
30		from tomobar.methodsDIR import RecToolsDIR
31		from savu.plugins.utils import register_plugin
32		from savu.core.iterate_plugin_group_utils import enable_iterative_loop, \
33		check_if_end_plugin_in_iterate_group, setup_extra_plugin_data_padding
34
35
36		@register_plugin
37		class TomobarRecon3d(BaseRecon, GpuPlugin):
38
39		def __init__(self):
40		super(TomobarRecon3d, self).__init__("TomobarRecon3d")
41		self.Vert_det = None
42		self.pad = None
43
44	View Code Duplication	@setup_extra_plugin_data_padding
		0 ignored issues – show Duplication introduced 2022-03-02 13:24 UTC by Report Bug Copy Issue Report This code seems to be duplicated in your project. Loading history...
45		def set_filter_padding(self, in_pData, out_pData):
46		self.pad = self.parameters['padding']
47		in_data = self.get_in_datasets()[0]
48		det_y = in_data.get_data_dimension_by_axis_label('detector_y')
49		pad_det_y = '%s.%s' % (det_y, self.pad)
50		pad_dict = {'pad_directions': [pad_det_y], 'pad_mode': 'edge'}
51		in_pData[0].padding = pad_dict
52		out_pData[0].padding = pad_dict
53		if len(self.get_in_datasets()) > 1:
54		in_pData[1].padding = pad_dict
55
56		@enable_iterative_loop
57		def setup(self):
58		in_dataset = self.get_in_datasets()[0]
59		procs = self.exp.meta_data.get("processes")
60		procs = len([i for i in procs if 'GPU' in i]) # calculates the total number of GPU processes
61		dim = in_dataset.get_data_dimension_by_axis_label('detector_y')
62		nSlices = int(np.ceil(in_dataset.get_shape()[dim] / float(procs)))
63		# calculate the amount of slices than would fit the GPU memory
64		gpu_available_mb = self.get_gpu_memory()[0]/procs # get the free GPU memory of a first device if many
65		det_x_dim = in_dataset.get_shape()[in_dataset.get_data_dimension_by_axis_label('detector_x')]
66		rot_angles_dim = in_dataset.get_shape()[in_dataset.get_data_dimension_by_axis_label('rotation_angle')]
67		slice_size_mbbytes = int(np.ceil(((det_x_dim * det_x_dim) * 1024 * 4) / (1024 ** 3)))
68
69		if self.parameters['reconstruction_method'] == 'FISTA3D':
70		# calculate the GPU memory required based on 3D regularisation restrictions (avoiding CUDA-error)
71		if 'ROF_TV' in self.parameters['regularisation_method']:
72		slice_size_mbbytes *= 8
73		if 'FGP_TV' in self.parameters['regularisation_method']:
74		slice_size_mbbytes *= 12
75		if 'SB_TV' in self.parameters['regularisation_method']:
76		slice_size_mbbytes *= 10
77		if 'PD_TV' in self.parameters['regularisation_method']:
78		slice_size_mbbytes *= 8
79		if 'LLT_ROF' in self.parameters['regularisation_method']:
80		slice_size_mbbytes *= 12
81		if 'TGV' in self.parameters['regularisation_method']:
82		slice_size_mbbytes *= 15
83		if 'NDF' in self.parameters['regularisation_method']:
84		slice_size_mbbytes *= 5
85		if 'Diff4th' in self.parameters['regularisation_method']:
86		slice_size_mbbytes *= 5
87		if 'NLTV' in self.parameters['regularisation_method']:
88		slice_size_mbbytes *= 8
89
90		slices_fit_total = int(gpu_available_mb / slice_size_mbbytes) - 2*self.parameters['padding']
91		if nSlices > slices_fit_total:
92		nSlices = slices_fit_total
93		self._set_max_frames(nSlices)
94		# get experimental metadata of projection_shifts
95		if 'projection_shifts' in list(self.exp.meta_data.dict.keys()):
96		self.projection_shifts = self.exp.meta_data.dict['projection_shifts']
97		super(TomobarRecon3d, self).setup()
98
99		def pre_process(self):
100		in_pData = self.get_plugin_in_datasets()[0]
101		self.det_dimX_ind = in_pData.get_data_dimension_by_axis_label('detector_x')
102		self.det_dimY_ind = in_pData.get_data_dimension_by_axis_label('detector_y')
103		# getting the value for padded vertical detector
104		self.Vert_det = in_pData.get_shape()[self.det_dimY_ind] + 2 * self.pad
105
106		# extract given parameters into dictionaries suitable for ToMoBAR input
107		self._data_ = {'OS_number': self.parameters['algorithm_ordersubsets'],
108		'huber_threshold': self.parameters['data_Huber_thresh'],
109		'ringGH_lambda': self.parameters['data_full_ring_GH'],
110		'ringGH_accelerate': self.parameters['data_full_ring_accelerator_GH']}
111
112		self._algorithm_ = {'iterations': self.parameters['algorithm_iterations'],
113		'nonnegativity': self.parameters['algorithm_nonnegativity'],
114		'mask_diameter': self.parameters['algorithm_mask'],
115		'verbose': self.parameters['algorithm_verbose']}
116
117		self._regularisation_ = {'method': self.parameters['regularisation_method'],
118		'regul_param': self.parameters['regularisation_parameter'],
119		'iterations': self.parameters['regularisation_iterations'],
120		'device_regulariser': self.parameters['regularisation_device'],
121		'edge_threhsold': self.parameters['regularisation_edge_thresh'],
122		'time_marching_step': self.parameters['regularisation_timestep'],
123		'regul_param2': self.parameters['regularisation_parameter2'],
124		'PD_LipschitzConstant': self.parameters['regularisation_PD_lip'],
125		'NDF_penalty': self.parameters['regularisation_NDF_penalty'],
126		'methodTV': self.parameters['regularisation_methodTV']}
127
128		def process_frames(self, data):
129		cor, angles, self.vol_shape, init = self.get_frame_params()
130		self.anglesRAD = np.deg2rad(angles.astype(np.float32))
131		projdata3D = data[0].astype(np.float32)
132		dim_tuple = np.shape(projdata3D)
133		self.Horiz_det = dim_tuple[self.det_dimX_ind]
134		half_det_width = 0.5 * self.Horiz_det
135		projdata3D[projdata3D > 10 ** 15] = 0.0
136		projdata3D = np.swapaxes(projdata3D, 0, 1)
137		self._data_.update({'projection_norm_data': projdata3D})
138
139		# dealing with projection shifts and the CoR
140		cor_astra = half_det_width - np.mean(cor)
141		CenterOffset = cor_astra.item() - 0.5
142		if np.sum(self.projection_shifts) != 0.0:
143		CenterOffset = np.zeros(np.shape(self.projection_shifts))
144		CenterOffset[:, 0] = (cor_astra.item() - 0.5) - self.projection_shifts[:, 0]
145		CenterOffset[:, 1] = -self.projection_shifts[:, 1] - 0.5
146
147		if self.parameters['reconstruction_method'] == 'FISTA3D':
148		if self.parameters['regularisation_method'] == 'PD_TV':
149		self._regularisation_.update({'device_regulariser': self.parameters['GPU_index']})
150		# if one selects PWLS or SWLS models then raw data is also required (2 inputs)
151		if (self.parameters['data_fidelity'] == 'PWLS') or (self.parameters['data_fidelity'] == 'SWLS'):
152		rawdata3D = data[1].astype(np.float32)
153		rawdata3D[rawdata3D > 10 ** 15] = 0.0
154		rawdata3D = np.swapaxes(rawdata3D, 0, 1) / np.max(np.float32(rawdata3D))
155		self._data_.update({'projection_raw_data': rawdata3D})
156		self._data_.update({'beta_SWLS': self.parameters['data_beta_SWLS'] * np.ones(self.Horiz_det)})
157
158		# set parameters and initiate a TomoBar class object for FISTA reconstruction
159		RectoolsIter = RecToolsIR(DetectorsDimH=self.Horiz_det, # DetectorsDimH # detector dimension (horizontal)
160		DetectorsDimV=self.Vert_det, # DetectorsDimV # detector dimension (vertical) for 3D case only
161		CenterRotOffset=CenterOffset, # The center of rotation combined with the shift offsets
162		AnglesVec=-self.anglesRAD, # the vector of angles in radians
163		ObjSize=self.vol_shape[0], # a scalar to define the reconstructed object dimensions
164		datafidelity=self.parameters['data_fidelity'], # data fidelity, choose LS, PWLS, SWLS
165		device_projector=self.parameters['GPU_index'])
166
167		# Run FISTA reconstruction algorithm here
168		recon = RectoolsIter.FISTA(self._data_, self._algorithm_, self._regularisation_)
169
170		if self.parameters['reconstruction_method'] == 'FBP3D':
171		RectoolsDIR = RecToolsDIR(DetectorsDimH=self.Horiz_det, # DetectorsDimH # detector dimension (horizontal)
172		DetectorsDimV=self.Vert_det, # DetectorsDimV # detector dimension (vertical) for 3D case only
173		CenterRotOffset=CenterOffset, # The center of rotation combined with the shift offsets
174		AnglesVec=-self.anglesRAD, # the vector of angles in radians
175		ObjSize=self.vol_shape[0], # a scalar to define the reconstructed object dimensions
176		device_projector=self.parameters['GPU_index'])
177
178		recon = RectoolsDIR.FBP(projdata3D) #perform FBP3D
179
180		if self.parameters['reconstruction_method'] == 'CGLS3D':
181		# set parameters and initiate a TomoBar class object for FISTA reconstruction
182		RectoolsIter = RecToolsIR(DetectorsDimH=self.Horiz_det, # DetectorsDimH # detector dimension (horizontal)
183		DetectorsDimV=self.Vert_det, # DetectorsDimV # detector dimension (vertical) for 3D case only
184		CenterRotOffset=CenterOffset, # The center of rotation combined with the shift offsets
185		AnglesVec=-self.anglesRAD, # the vector of angles in radians
186		ObjSize=self.vol_shape[0], # a scalar to define the reconstructed object dimensions
187		datafidelity=self.parameters['data_fidelity'], # data fidelity, choose LS, PWLS, SWLS
188		device_projector=self.parameters['GPU_index'])
189
190		# Run CGLS reconstruction algorithm here
191		recon = RectoolsIter.CGLS(self._data_, self._algorithm_)
192
193		recon = np.swapaxes(recon, 0, 1)
		0 ignored issues – show introduced 2022-03-16 17:09 UTC by Report Bug Copy Issue Report The variable `recon` does not seem to be defined in case `SubscriptNode == 'FISTA3D'` on line `147` is `False`. Are you sure this can never be the case? Loading history...
194		return recon
195
196		def nInput_datasets(self):
197		return max(len(self.parameters['in_datasets']), 1)
198
199		# total number of output datasets
200		def nOutput_datasets(self):
201		if check_if_end_plugin_in_iterate_group(self.exp):
202		return 2
203		else:
204		return 1
205
206		# total number of output datasets that are clones
207		def nClone_datasets(self):
208		if check_if_end_plugin_in_iterate_group(self.exp):
209		return 1
210		else:
211		return 0
212
213		def _set_max_frames(self, frames):
214		self._max_frames = frames
215

DiamondLightSource / Savu

Push — master ( 877aea...3bbc8e )

TomobarRecon3d._set_max_frames() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like