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 measure_ufo(out_path, metric, axis, width, height):
pm = Ufo.PluginManager()
sched = Ufo.Scheduler()
graph = Ufo.TaskGraph()
input_path = 'data/measure.tif'
image = make_input(width, height)
tifffile.imsave(input_path, image)
reader = pm.get_task('read')
measure = pm.get_task('measure')
output = Ufo.OutputTask()
reader.props.path = input_path
measure.props.axis = axis
measure.props.metric = metric
graph.connect_nodes(reader, measure)
graph.connect_nodes(measure, output)
sched.run(graph)
buf = output.get_output_buffer()
gpu_result = ufo.numpy.asarray(buf)
write_image(out_path, gpu_result)
def _run(params, x_region, y_region, regions, index):
"""Execute one pass on all possible GPUs with slice ranges given by *regions*."""
from gi.repository import Ufo
pm = Ufo.PluginManager()
graph = Ufo.TaskGraph()
scheduler = Ufo.FixedScheduler()
gpus = scheduler.get_resources().get_gpu_nodes()
num_gpus = len(gpus)
broadcast = Ufo.CopyTask()
source = _setup_source(params, pm, graph)
graph.connect_nodes(source, broadcast)
for i, region in enumerate(regions):
subindex = index * num_gpus + i
_setup_graph(pm,
graph,
subindex,
x_region,
y_region,
region,
params,
broadcast,
gpu=gpus[i])
scheduler.run(graph)
duration = scheduler.props.time
LOG.info('Execution time: {} s'.format(duration))
return duration
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 _update(self):
"""Update the regions and volume sizes based on changed args or region."""
st = time.perf_counter()
x_region, y_region, z_region = get_reconstruction_regions(self.args)
if not self._resources:
self._resources = [Ufo.Resources()]
gpus = np.array(self._resources[0].get_gpu_nodes())
gpu_indices = np.array(self.args.gpus or list(range(len(gpus))))
if min(gpu_indices) < 0 or max(gpu_indices) > len(gpus) - 1:
raise ValueError('--gpus contains invalid indices')
gpus = gpus[gpu_indices]
if self.regions is None:
self._regions = make_runs(gpus, gpu_indices, x_region, y_region, z_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,
num_gpu_threads=self.args.num_gpu_threads)
else:
self._regions = self.regions
offset = 0
for batch in self._regions:
for i, region in batch:
if len(self._resources) < len(batch):
self._resources.append(Ufo.Resources())
offset += len(np.arange(*region))
if self.args.slice_metric:
shape = (offset,)
else:
shape = (offset, len(np.arange(*y_region)), len(np.arange(*x_region)))
if self.volume is None or shape != self.volume.shape:
self.volume = np.empty(shape, dtype=np.float32)
LOG.log(PERFDEBUG, 'Backprojector manager update duration: %g s', time.perf_counter() - st)
def compute_phase_spectrum_density(args):
graph = Ufo.TaskGraph()
scheduler = Ufo.Scheduler()
build_power_spectrum_density_pipeline(args, graph, scheduler)
# scheduler.run(graph) # runned inside for now
duration = scheduler.props.time
return duration
def run_flat_correct(args):
graph = Ufo.TaskGraph()
sched = Ufo.Scheduler()
pm = Ufo.PluginManager()
out_task = get_writer(args)
flat_task = create_flat_correct_pipeline(args, graph)
graph.connect_nodes(flat_task, out_task)
sched.run(graph)
def target():
from gi.repository import Ufo
if self._config:
sched = Ufo.Scheduler(config=self._config)
else:
sched = Ufo.Scheduler()
sched.run(self._graph)
def run_preprocessing(args):
graph = Ufo.TaskGraph()
sched = Ufo.Scheduler()
pm = Ufo.PluginManager()
out_task = get_writer(args)
current = create_preprocessing_pipeline(args, graph)
graph.connect_nodes(current, out_task)
sched.run(graph)
def run_flat_correct(args):
graph = Ufo.TaskGraph()
sched = Ufo.Scheduler()
pm = Ufo.PluginManager()
out_task = pm.get_task('write')
out_task.props.filename = args.output
flat_task = create_flat_correct_pipeline(args, graph)
graph.connect_nodes(flat_task, out_task)
sched.run(graph)
def generate_partial(append=False):
graph = Ufo.TaskGraph()
sched = Ufo.Scheduler()
args.output_append = append
writer = get_writer(args)
sinos = create_sinogram_pipeline(args, graph)
graph.connect_nodes(sinos, writer)
sched.run(graph)
def generate_partial(append=False):
pm = Ufo.PluginManager()
graph = Ufo.TaskGraph()
sched = Ufo.Scheduler()
writer = pm.get_task('write')
writer.props.filename = args.output
writer.props.append = append
sinos = create_sinogram_pipeline(args, graph)
graph.connect_nodes(sinos, writer)
sched.run(graph)
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 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 initialize(self):
self.data = self.get("DataStore").data()
self.graph = Ufo.TaskGraph()
self.sched = Ufo.Scheduler()
manager = Ufo.PluginManager()
self.read = manager.get_task('memory-in')
self.sino = manager.get_task('transpose-projections')
self.pad = manager.get_task('pad')
self.fft = manager.get_task('fft')
self.fltr = manager.get_task('filter')
self.ifft = manager.get_task('ifft')
self.bp = manager.get_task('backproject')
self.crop = manager.get_task('crop')
self.write= manager.get_task('memory-out')
self.LogInfo("initialized, UFO Reconstruction")
return True
def test_wrong_connection():
from ufo import Read, FlatFieldCorrect, Write
darks = Read()
flats = Read()
radios = Read()
write = Write()
ffc = FlatFieldCorrect()
g = Ufo.TaskGraph()
g.connect_nodes_full(radios.task, ffc.task, 0)
g.connect_nodes_full(darks.task, ffc.task, 1)
g.connect_nodes_full(flats.task, ffc.task, 0)
g.connect_nodes(ffc.task, write.task)
sched = Ufo.Scheduler()
sched.run(g)
def start_one(index):
gpu_index, region = regions[index]
scheduler = Ufo.FixedScheduler()
scheduler.set_resources(resources[index])
graph = Ufo.TaskGraph()
gpu = scheduler.get_resources().get_gpu_nodes()[gpu_index]
region_index = run_number * len(resources) + index
geometry = CTGeometry(args)
if (len(args.center_position_z) == 1
and np.modf(args.center_position_z[0])[0] == 0
and geometry.is_simple_parallel_tomo):
LOG.info(
'Simple tomography with integer z center, changing to center_position_z + 0.5 '
'to avoid interpolation')
geometry.args.center_position_z = (
geometry.args.center_position_z[0] + 0.5, )
if not args.disable_projection_crop:
if not args.dry_run and (args.y or args.height):
LOG.debug(
'--y or --height specified, not optimizing projection region'
)
else:
geometry.optimize_args(region=region)
opt_args = geometry.args
if args.dry_run:
source = get_task('dummy-data',
number=args.number,
width=args.width,
height=args.height)
else:
source = None
setup_graph(opt_args,
graph,
x_region,
y_region,
region,
source=source,
gpu=gpu,
index=region_index,
make_reader=True)
LOG.debug('Pass: %d, device: %d, region: %s', run_number + 1,
gpu_index, region)
scheduler.run(graph)
return scheduler.props.time
def test_resource_info():
resources = Ufo.Resources()
nodes = resources.get_gpu_nodes()
assert (nodes)
node = nodes[0]
assert (node.get_info(Ufo.GpuNodeInfo.LOCAL_MEM_SIZE) > 0)
assert (node.get_info(Ufo.GpuNodeInfo.MAX_MEM_ALLOC_SIZE) > 0)
assert (node.get_info(Ufo.GpuNodeInfo.GLOBAL_MEM_SIZE) > node.get_info(
Ufo.GpuNodeInfo.LOCAL_MEM_SIZE))
def create_preprocessing_pipeline(args,
graph,
source=None,
processing_node=None):
pm = Ufo.PluginManager()
if not (args.width and args.height):
width, height = determine_shape(args, args.projections)
if not width:
raise RuntimeError("Could not determine width from the input")
if not args.width:
args.width = width
if not args.height:
args.height = height - args.y
LOG.debug('Image width x height: %d x %d', args.width, args.height)
if source:
current = source
elif args.darks and args.flats:
current = create_flat_correct_pipeline(args,
graph,
processing_node=processing_node)
else:
current = get_task('read')
set_node_props(current, args)
if not args.projections:
raise RuntimeError('--projections not set')
setup_read_task(current, args.projections, args)
if args.absorptivity:
absorptivity = get_task('calculate',
processing_node=processing_node)
absorptivity.props.expression = '-log(v)'
graph.connect_nodes(current, absorptivity)
current = absorptivity
if args.transpose_input:
transpose = get_task('transpose')
graph.connect_nodes(current, transpose)
current = transpose
tmp = args.width
args.width = args.height
args.height = tmp
if args.projection_filter != 'none':
pf_first, pf_last = create_projection_filtering_pipeline(
args, graph, processing_node=processing_node)
graph.connect_nodes(current, pf_first)
current = pf_last
if args.energy is not None and args.propagation_distance is not None:
pr_first, pr_last = create_phase_retrieval_pipeline(
args, graph, processing_node=processing_node)
graph.connect_nodes(current, pr_first)
current = pr_last
return current
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 transpose_ufo(width, height, path):
pm = Ufo.PluginManager()
graph = Ufo.TaskGraph()
sched = Ufo.Scheduler()
input_path = 'data/transpose.tif'
image = make_input(width, height)
tifffile.imsave(input_path, image)
reader = pm.get_task('read')
transpose = pm.get_task('transpose')
writer = pm.get_task('write')
reader.props.path = input_path
writer.props.filename = path
graph.connect_nodes(reader, transpose)
graph.connect_nodes(transpose, writer)
sched.run(graph)
def genreco(args):
st = time.time()
_fill_missing_args(args)
_convert_angles_to_rad(args)
set_projection_filter_scale(args)
x_region, y_region, z_region = get_reconstruction_regions(args, store=True)
resources = [Ufo.Resources()]
gpus = np.array(resources[0].get_gpu_nodes())
gpu_indices = np.array(args.gpus or range(len(gpus)))
if min(gpu_indices) < 0 or max(gpu_indices) > len(gpus) - 1:
raise ValueError('--gpus contains invalid indices')
gpus = gpus[gpu_indices]
duration = 0
for i, gpu in enumerate(gpus):
print 'Max mem for {}: {:.2f} GB'.format(i, gpu.get_info(0) / 2**30)
runs = make_runs(gpus,
gpu_indices,
x_region,
y_region,
z_region,
DTYPE_CL_SIZE[args.store_type],
slices_per_device=args.slices_per_device,
slice_memory_coeff=args.slice_memory_coeff,
data_splitting_policy=args.data_splitting_policy)
for i in range(len(runs[0]) - 1):
resources.append(Ufo.Resources())
LOG.info('Number of passes: %d', len(runs))
for i, regions in enumerate(runs):
duration += _run(resources, args, x_region, y_region, regions, i)
vol_shape = get_reconstructed_cube_shape(x_region, y_region, z_region)
num_gupdates = vol_shape[0] * vol_shape[1] * vol_shape[
2] * args.number * 1e-9
total_duration = time.time() - st
LOG.debug('UFO duration: %.2f s', duration)
LOG.debug('Total duration: %.2f s', total_duration)
LOG.debug('UFO performance: %.2f GUPS', num_gupdates / duration)
LOG.debug('Total performance: %.2f GUPS', num_gupdates / total_duration)
def test_core_issue_64_fixed_expansion():
g = Ufo.TaskGraph()
pm = Ufo.PluginManager()
sched = Ufo.FixedScheduler()
arch = Ufo.ArchGraph()
gpus = arch.get_gpu_nodes()
sched.set_gpu_nodes(arch, gpus)
generate = pm.get_task('generate')
null = pm.get_task('null')
generate.set_properties(number=5, width=512, height=512)
for gpu in gpus:
median = pm.get_task('median-filter')
median.set_proc_node(gpu)
g.connect_nodes(generate, median)
g.connect_nodes(median, null)
sched.run(g)
def test_broadcast():
from ufo import Generate, Writer
import glob
with tempdir() as d:
generate = Generate(number=5, width=512, height=512)
write1 = Writer(filename=d.path('foo-%i.tif'))
write2 = Writer(filename=d.path('bar-%i.tif'))
g = Ufo.TaskGraph()
g.connect_nodes(generate.task, write1.task)
g.connect_nodes(generate.task, write2.task)
sched = Ufo.Scheduler()
sched.run(g)
foos = glob.glob(d.path('foo-*'))
bars = glob.glob(d.path('bar-*'))
assert (len(foos) == 5)
assert (len(bars) == 5)
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 _create_runs(params, queue):
"""Workaround function to get the number of gpus and compute regions. gi.repository must always
be called in a separate process, otherwise the resources return None gpus.
"""
#TODO: remove the whole function after memory leak fix!
from gi.repository import Ufo
scheduler = Ufo.FixedScheduler()
gpus = scheduler.get_resources().get_gpu_nodes()
num_gpus = len(gpus)
x_region, y_region, regions = _split_regions(params, gpus)
LOG.info('Using {} GPUs in {} passes'.format(min(len(regions), num_gpus), len(regions)))
queue.put((x_region, y_region, regions, num_gpus))
def create_phase_retrieval_pipeline(args, graph, processing_node=None):
LOG.debug('Creating phase retrieval pipeline')
pm = Ufo.PluginManager()
# Retrieve phase
phase_retrieve = get_task('retrieve-phase',
processing_node=processing_node)
pad_phase_retrieve = get_task('pad', processing_node=processing_node)
crop_phase_retrieve = get_task('crop', processing_node=processing_node)
fft_phase_retrieve = get_task('fft', processing_node=processing_node)
ifft_phase_retrieve = get_task('ifft', processing_node=processing_node)
width = args.width
height = args.height
default_padded_width = next_power_of_two(width)
default_padded_height = next_power_of_two(height)
if not args.retrieval_padded_width:
args.retrieval_padded_width = default_padded_width
if not args.retrieval_padded_height:
args.retrieval_padded_height = default_padded_height
fmt = 'Phase retrieval padding: {}x{} -> {}x{}'
LOG.debug(
fmt.format(width, height, args.retrieval_padded_width,
args.retrieval_padded_height))
x = (args.retrieval_padded_width - width) / 2
y = (args.retrieval_padded_height - height) / 2
pad_phase_retrieve.props.x = x
pad_phase_retrieve.props.y = y
pad_phase_retrieve.props.width = args.retrieval_padded_width
pad_phase_retrieve.props.height = args.retrieval_padded_height
pad_phase_retrieve.props.addressing_mode = args.retrieval_padding_mode
crop_phase_retrieve.props.x = x
crop_phase_retrieve.props.y = y
crop_phase_retrieve.props.width = width
crop_phase_retrieve.props.height = height
phase_retrieve.props.method = args.retrieval_method
phase_retrieve.props.energy = args.energy
phase_retrieve.props.distance = args.propagation_distance
phase_retrieve.props.pixel_size = args.pixel_size
phase_retrieve.props.regularization_rate = args.regularization_rate
phase_retrieve.props.thresholding_rate = args.thresholding_rate
fft_phase_retrieve.props.dimensions = 2
ifft_phase_retrieve.props.dimensions = 2
graph.connect_nodes(pad_phase_retrieve, fft_phase_retrieve)
graph.connect_nodes(fft_phase_retrieve, phase_retrieve)
graph.connect_nodes(phase_retrieve, ifft_phase_retrieve)
graph.connect_nodes(ifft_phase_retrieve, crop_phase_retrieve)
return (pad_phase_retrieve, crop_phase_retrieve)
def __init__(self, axis_pos=None):
self.pm = PluginManager()
self.fft = self.pm.get_task('fft', dimensions=1)
self.ifft = self.pm.get_task('ifft', dimensions=1)
self.fltr = self.pm.get_task('filter')
self.backprojector = self.pm.get_task('backproject')
if axis_pos:
self.backprojector.props.axis_pos = axis_pos
graph = Ufo.TaskGraph()
graph.connect_nodes(self.fft, self.fltr)
graph.connect_nodes(self.fltr, self.ifft)
graph.connect_nodes(self.ifft, self.backprojector)
super(Backproject, self).__init__(graph, get_output=True)
self.output_task.props.num_dims = 2
def start(self, arch=None, gpu=None):
"""
Run the processing in a new thread.
Use :meth:`.push` to insert data into the processing chaing and
:meth:`~InjectProcess.wait` to wait until processing has finished."""
def run_scheduler(sched):
sched.run(self.graph)
if arch and gpu:
sched = Ufo.FixedScheduler()
sched.set_gpu_nodes(arch, [gpu])
else:
sched = self.sched
self.thread = threading.Thread(target=run_scheduler, args=(sched,), daemon=True)
self.thread.start()
if not self._started:
self._started = True
def items(self):
results = queue.Queue()
def processed(task):
buf = output_task.get_output_buffer()
results.put(asarray(buf).copy())
output_task.release_output_buffer(buf)
output_task = Ufo.OutputTask()
output_task.connect('processed', processed)
self.env.graph.connect_nodes(self.task, output_task)
self.run()
while not self.env.done or not results.empty():
try:
yield results.get(True, 0.01)
except queue.Empty:
pass
self.join()