def align_3D_tomo_file_mpi_specific(file_save_path, files_recon=[], files_ref='', binning=1, circle_mask_ratio=0.8, file_type='.h5', align_coarse=1, align_method=1, hdf_attr='img', num_cpu=4):
'''
align_method:
1: old method
2: 3D cross-correlation
'''
from multiprocessing import Pool, cpu_count
from functools import partial
num_cpu = min(round(cpu_count() * 0.8), num_cpu)
print(f'align_3D_tomo using {num_cpu:2d} CPUs')
# save ref image
file_path = os.path.abspath(file_save_path)
img_ref, scan_id, X_eng = get_tomo_image(files_ref, file_type, hdf_attr)
if binning > 1:
img_ref = pyxas.bin_image(img_ref, binning)
if circle_mask_ratio < 1:
img_ref = pyxas.circ_mask(img_ref, axis=0, ratio=circle_mask_ratio)
if align_method == 1:
img_ref = pyxas.move_3D_to_center(img_ref, circle_mask_ratio=circle_mask_ratio)
fn_save = f'{file_save_path}/ali_recon_{scan_id}_bin_{binning}.h5'
pyxas.save_hdf_file(fn_save, 'img', img_ref.astype(np.float32), 'scan_id', scan_id, 'X_eng', X_eng)
# start align
pool = Pool(num_cpu)
pool.map(partial(align_3D_tomo_file_mpi_sub,
img_ref=img_ref,
file_path=file_path,
binning=binning,
circle_mask_ratio=circle_mask_ratio,
file_type=file_type,
align_coarse=align_coarse,
align_method=align_method,
hdf_attr=hdf_attr),
files_recon)
pool.close()
def align_3D_tomo_file_specific(file_save_path='.', files_recon=[], files_ref='', binning=1, circle_mask_ratio=0.9, file_type='.h5', align_coarse=1, align_method=1, hdf_attr='img'):
'''
align_method:
1: old method
2: 3D cross-correlation
'''
import time
file_path = os.path.abspath(file_save_path)
img_ref, scan_id, X_eng = get_tomo_image(files_ref, file_type, hdf_attr)
if binning > 1:
img_ref = pyxas.bin_image(img_ref, binning)
if circle_mask_ratio < 1:
img_ref = pyxas.circ_mask(img_ref, axis=0, ratio=circle_mask_ratio)
if align_method == 1:
img_ref = pyxas.move_3D_to_center(img_ref, circle_mask_ratio=circle_mask_ratio)
fn_save = f'{file_path}/ali_recon_{scan_id}_bin_{binning}.h5'
pyxas.save_hdf_file(fn_save, 'img', img_ref.astype(np.float32), 'scan_id', scan_id, 'X_eng', X_eng)
time_start = time.time()
num_file = len(files_recon)
for i in range(num_file):
fn = files_recon[i]
align_3D_tomo_file_mpi_sub(fn, img_ref, file_path, binning, circle_mask_ratio, file_type, align_coarse, align_method, hdf_attr)
print(f'time elasped: {time.time() - time_start:05.1f}\n')
def align_3D_tomo_file_mpi_sub(files_recon, img_ref, file_path='.', binning=1, circle_mask_ratio=0.9, file_type='.h5', align_coarse=1, align_method=1, hdf_attr='img'):
'''
align_method:
1: old method
2: 3D cross-correlation
'''
bin_info = ''
fn = files_recon
fn_short = fn.split('/')[-1]
print(f'aligning {fn_short} ...')
img1, scan_id, X_eng = get_tomo_image(files_recon, file_type, hdf_attr)
if circle_mask_ratio < 1:
img1 = pyxas.circ_mask(img1, axis=0, ratio=circle_mask_ratio)
if binning > 1:
img1 = pyxas.bin_image(img1, binning)
bin_info == f'_bin_{binning}'
img_ali = align_3D_tomo_image(img1, img_ref, circle_mask_ratio, align_method, align_coarse)
if X_eng == 0: # read tiff file
fn_save = f'{file_path}/ali_{fn_short}'
print(f'saving aligned file: {fn_save.split("/")[-1]}\n')
io.imsave(fn_save, img_ali.astype(np.float32))
else:
fn_save = f'{file_path}/ali_recon_{scan_id}{bin_info}.h5'
print(f'saving aligned file: {fn_save.split("/")[-1]}\n')
pyxas.save_hdf_file(fn_save, 'img', img_ali.astype(np.float32), 'scan_id', scan_id, 'X_eng', X_eng)
def align_two_tomo_prj_mpi(fn_ref, fn_target, num_cpu=20):
from multiprocessing import Pool, cpu_count
from functools import partial
res_ref = pyxas.get_img_from_hdf_file(fn_ref, 'angle', 'img_tomo',
'img_bkg_avg', 'img_dark_avg',
'scan_id', 'X_eng')
fn_current = fn_ref.split('/')[:-1]
if not len(fn_current):
fn_current = '.'
fn_current = '/'.join(tmp for tmp in fn_current)
fn_short = fn_ref.split('/')[-1]
fn_save = f'{fn_current}/ali_{fn_short}'
ref_img = res_ref['img_tomo']
ref_angle = res_ref['angle']
res = pyxas.get_img_from_hdf_file(fn_target, 'angle', 'img_tomo',
'img_bkg_avg', 'img_dark_avg', 'scan_id',
'X_eng')
img_bkg = res['img_bkg_avg']
img_dark = res['img_dark_avg']
target_angle = res['angle']
fn_target_short = fn_target.split('/')[-1]
fn_save = f'{fn_current}/ali_norm_{fn_target_short}'
current_angle_index = np.arange(len(target_angle))
time_s = time.time()
pool = Pool(num_cpu)
res_mpi = pool.map(
partial(align_tomo_prj_mpi_sub,
fn_ref=fn_ref,
fn_target=fn_target,
img_dark=img_dark,
img_bkg=img_bkg), current_angle_index)
pool.close()
print(f'aligning {fn_target_short} using {time.time() - time_s:4.1f} sec')
img_ali_norm = np.zeros(res['img_tomo'].shape)
r_shift = np.zeros(len(ref_angle))
c_shift = r_shift.copy()
for i in range(len(target_angle)):
r_shift[i] = res_mpi[i]['r']
c_shift[i] = res_mpi[i]['c']
img_ali_norm[i] = res_mpi[i]['img_ali']
pyxas.save_hdf_file(fn_save, 'angle', res['angle'], 'img_tomo',
np.array(img_ali_norm, dtype=np.float32), 'scan_id',
res['scan_id'], 'X_eng', res['X_eng'], 'r_shift',
r_shift, 'c_shift', c_shift)
def align_tomo_proj_serial(file_path,
file_prefix='fly',
file_type='.h5',
ref_index=-1):
files_scan = pyxas.retrieve_file_type(file_path, file_prefix, file_type)
n = len(files_scan)
if ref_index == -1:
ref_index = len(files_scan) - 1
fn_ref = files_scan[ref_index]
res_ref = pyxas.get_img_from_hdf_file(fn_ref, 'angle', 'img_tomo',
'img_bkg_avg', 'img_dark_avg',
'scan_id', 'X_eng')
fn_current = fn_ref.split('/')[:-1]
fn_current = '/'.join(tmp for tmp in fn_current)
fn_short = fn_ref.split('/')[-1]
fn_save = f'{fn_current}/ali_{fn_short}'
ref_img = res_ref['img_tomo']
ref_angle = res_ref['angle']
r_shift = np.zeros(len(ref_angle))
c_shift = r_shift.copy()
pyxas.save_hdf_file(fn_save, 'angle', res_ref['angle'], 'img_tomo',
res_ref['img_tomo'], 'img_bkg_avg',
res_ref['img_bkg_avg'], 'img_dark_avg',
res_ref['img_dark_avg'], 'scan_id', res_ref['scan_id'],
'X_eng', res_ref['X_eng'], 'r_shift', r_shift,
'c_shift', c_shift)
for i in range(len(files_scan)):
fn = files_scan[i]
if i == ref_index:
continue
time_s = time.time()
res = pyxas.get_img_from_hdf_file(fn, 'angle', 'img_tomo',
'img_bkg_avg', 'img_dark_avg',
'scan_id', 'X_eng')
res_angle = res['angle']
prj_ali = res['img_tomo'].copy()
r_shift = np.zeros(len(res_angle))
c_shift = r_shift.copy()
for j in range(len(res_angle)):
angle_id = pyxas.find_nearest(ref_angle, res_angle[j])
img_ref = res_ref['img_tomo'][angle_id]
prj_ali[j], r_shift[j], c_shift[j] = pyxas.align_img_stackreg(
img_ref, prj_ali[j], align_flag=1)
print(
f'{fn.split("/")[-1]} proj #{j}: r_shift = {r_shift[j]:3.1f}, c_shift = {c_shift[j]:3.1f}'
)
fn_short = fn.split('/')[-1]
fn_save = f'{fn_current}/ali_{fn_short}'
pyxas.save_hdf_file(fn_save, 'angle', res['angle'], 'img_tomo',
np.array(prj_ali, dtype=np.float32), 'img_bkg_avg',
res['img_bkg_avg'], 'img_dark_avg',
res['img_dark_avg'], 'scan_id', res['scan_id'],
'X_eng', res['X_eng'], 'r_shift', r_shift,
'c_shift', c_shift)
print(
f'{fn_save.split("/")[-1]} saved, time elaped: {time.time() - time_s:4.2f} sec'
)