def input_data(*iargs):
tasks = []
buffers = []
for i in range(len(iargs)):
task = Ufo.InputTask()
tasks.append(task)
buffers.append(None)
self.env.graph.connect_nodes_full(task, self.task, i)
while not self.env.started:
time.sleep(0.01)
for args in zip(*iargs):
for i, (task, data) in enumerate(zip(tasks, args)):
if buffers[i] is None:
buffers[i] = empty_like(data)
else:
buffers[i] = task.get_input_buffer()
buffers[i].set_host_array(
data.__array_interface__['data'][0], False)
task.release_input_buffer(buffers[i])
for task in tasks:
task.stop()
def __init__(self, graph, get_output=False):
self.output_task = None
self._started = False
if isinstance(graph, Ufo.TaskGraph):
self.graph = graph
roots = self.graph.get_roots()
elif isinstance(graph, Ufo.TaskNode):
self.graph = Ufo.TaskGraph()
roots = [graph]
else:
msg = 'graph is neither Ufo.TaskGraph nor Ufo.TaskNode'
raise ValueError(msg)
# Initialize inputs
self.input_tasks = {}
self.ufo_buffers = {}
for root in roots:
self.input_tasks[root] = []
self.ufo_buffers[root] = []
num_inputs = root.get_num_inputs()
for i in range(num_inputs):
self.input_tasks[root].append(Ufo.InputTask())
self.ufo_buffers[root].append(None)
self.graph.connect_nodes_full(self.input_tasks[root][i], root,
i)
if get_output:
self.output_task = Ufo.OutputTask()
leaves = self.graph.get_leaves()
self.graph.connect_nodes(leaves[0], self.output_task)
def __init__(self, graph, get_output=False, output_dims=2, scheduler=None, copy_inputs=False):
self.output_tasks = []
self.sched = scheduler if scheduler else Ufo.Scheduler()
self._started = False
self.copy_inputs = copy_inputs
if isinstance(graph, Ufo.TaskGraph):
self.graph = graph
roots = self.graph.get_roots()
elif isinstance(graph, Ufo.TaskNode):
self.graph = Ufo.TaskGraph()
roots = [graph]
else:
msg = 'graph is neither Ufo.TaskGraph nor Ufo.TaskNode'
raise ValueError(msg)
# Initialize inputs
self.input_tasks = {}
self.ufo_buffers = {}
for root in roots:
self.input_tasks[root] = []
self.ufo_buffers[root] = []
num_inputs = root.get_num_inputs()
for i in range(num_inputs):
self.input_tasks[root].append(Ufo.InputTask())
self.ufo_buffers[root].append(None)
self.graph.connect_nodes_full(self.input_tasks[root][i], root, i)
if get_output:
for i, leave in enumerate(self.graph.get_leaves()):
self.output_tasks.append(Ufo.OutputTask())
self.output_tasks[-1].props.num_dims = output_dims
self.graph.connect_nodes(leave, self.output_tasks[-1])
def __init__(self, graph):
self.input_task = Ufo.InputTask()
if isinstance(graph, Ufo.TaskGraph):
self.graph = graph
roots = self.graph.get_roots()
self.graph.connect_nodes(self.input_task, roots[0])
elif isinstance(graph, Ufo.TaskNode):
self.graph = Ufo.TaskGraph()
self.graph.connect_nodes(self.input_task, graph)
else:
msg = 'graph is neither Ufo.TaskGraph nor Ufo.TaskNode'
raise ValueError(msg)
self.ufo_buffer = None
def ufo_dfi(tomo, center, recon, theta, **kwargs):
"""
Reconstruct object using UFO's Direct Fourier pipeline
"""
import gi
gi.require_version('Ufo', '0.0')
from gi.repository import Ufo
theta = theta[1] - theta[0]
center = center[0]
g = Ufo.TaskGraph()
pm = Ufo.PluginManager()
sched = Ufo.Scheduler()
input_task = Ufo.InputTask()
output_task = Ufo.OutputTask()
pad = pm.get_task('zeropad')
fft = pm.get_task('fft')
ifft = pm.get_task('ifft')
dfi = pm.get_task('dfi-sinc')
swap_forward = pm.get_task('swap-quadrants')
swap_backward = pm.get_task('swap-quadrants')
pad.set_properties(oversampling=1, center_of_rotation=center)
fft.set_properties(dimensions=1, auto_zeropadding=False)
ifft.set_properties(dimensions=2)
dfi.set_properties(angle_step=theta)
g.connect_nodes(input_task, pad)
g.connect_nodes(pad, fft)
g.connect_nodes(fft, dfi)
g.connect_nodes(dfi, swap_forward)
g.connect_nodes(swap_forward, ifft)
g.connect_nodes(ifft, swap_backward)
g.connect_nodes(swap_backward, output_task)
args = (input_task, output_task, tomo, recon)
thread = threading.Thread(target=_process_data, args=args)
thread.start()
sched.run(g)
thread.join()
logger.info("UFO+DFI run time: {}s".format(sched.props.time))
def ufo_fbp(tomo, center, recon, theta, **kwargs):
"""
Reconstruct object using UFO's FBP pipeline
"""
import gi
gi.require_version('Ufo', '0.0')
from gi.repository import Ufo
width = tomo.shape[2]
theta = theta[1] - theta[0]
center = center[0]
g = Ufo.TaskGraph()
pm = Ufo.PluginManager()
sched = Ufo.Scheduler()
input_task = Ufo.InputTask()
output_task = Ufo.OutputTask()
fft = pm.get_task('fft')
ifft = pm.get_task('ifft')
fltr = pm.get_task('filter')
backproject = pm.get_task('backproject')
ifft.set_properties(crop_width=width)
backproject.set_properties(axis_pos=center,
angle_step=theta,
angle_offset=np.pi)
g.connect_nodes(input_task, fft)
g.connect_nodes(fft, fltr)
g.connect_nodes(fltr, ifft)
g.connect_nodes(ifft, backproject)
g.connect_nodes(backproject, output_task)
args = (input_task, output_task, tomo, recon)
thread = threading.Thread(target=_process_data, args=args)
thread.start()
sched.run(g)
thread.join()
logger.info("UFO+FBP run time: {}s".format(sched.props.time))
def __init__(self, args, resources=None, gpu_index=0, do_normalization=False,
region=None, copy_inputs=False):
if args.width is None or args.height is None:
raise GeneralBackprojectError('width and height must be set in GeneralBackprojectArgs')
scheduler = Ufo.FixedScheduler()
if resources:
scheduler.set_resources(resources)
gpu = scheduler.get_resources().get_gpu_nodes()[gpu_index]
self.args = copy.deepcopy(args)
x_region, y_region, z_region = get_reconstruction_regions(self.args, store=True,
dtype=float)
set_projection_filter_scale(self.args)
if region is not None:
self.args.region = region
LOG.debug('Creating reconstructor for gpu %d, region: %s', gpu_index, self.args.region)
geometry = CTGeometry(self.args)
if not self.args.disable_projection_crop:
geometry.optimize_args()
self.args = geometry.args
regions = make_runs([gpu], [gpu_index], x_region, y_region, self.args.region,
DTYPE_CL_SIZE[self.args.store_type],
slices_per_device=self.args.slices_per_device,
slice_memory_coeff=self.args.slice_memory_coeff,
data_splitting_policy=self.args.data_splitting_policy)
if len(regions) > 1:
raise GeneralBackprojectError('Region does not fit to the GPU memory')
graph = Ufo.TaskGraph()
# Normalization
self.ffc = None
self.do_normalization = do_normalization
if do_normalization:
self.ffc = get_task('flat-field-correct', processing_node=gpu)
self.ffc.props.fix_nan_and_inf = self.args.fix_nan_and_inf
self.ffc.props.absorption_correct = self.args.absorptivity
self._darks_averaged = False
self._flats_averaged = False
self.dark_avg = get_task('average', processing_node=gpu)
self.flat_avg = get_task('average', processing_node=gpu)
graph.connect_nodes_full(self.dark_avg, self.ffc, 1)
graph.connect_nodes_full(self.flat_avg, self.ffc, 2)
(first, last) = setup_graph(self.args, graph, x_region, y_region, self.args.region,
source=self.ffc, gpu=gpu, index=gpu_index, do_output=False,
make_reader=False)
output_dims = 2
if args.slice_metric:
output_dims = 1
metric = self.args.slice_metric
if args.slice_metric == 'sag':
metric = 'sum'
gradient_task = get_task('gradient', processing_node=gpu, direction='both_abs')
graph.connect_nodes(last, gradient_task)
last = gradient_task
measure_task = get_task('measure', processing_node=gpu, axis=-1, metric=metric)
graph.connect_nodes(last, measure_task)
elif first == last:
# There are no other processing steps other than back projection
LOG.debug('Only back projection, no other processing')
graph = first
super().__init__(graph, get_output=True, output_dims=output_dims, scheduler=scheduler,
copy_inputs=copy_inputs)
if self.do_normalization:
# Setup input tasks for normalization images averaging. Our parent picks up only the two
# averagers and not the ffc's zero port for projections.
self.input_tasks[self.ffc] = [Ufo.InputTask()]
self.ufo_buffers[self.ffc] = [None]
self.graph.connect_nodes_full(self.input_tasks[self.ffc][0], self.ffc, 0)