def batch_recon(file_path,
file_prefix='fly',
file_type='.h5',
rot_cen=[],
binning=1,
sli=[],
block_list=[],
txm_normed_flag=0,
read_full_memory=1):
files_scan = pyxas.retrieve_file_type(file_path, file_prefix, file_type)
n = len(files_scan)
if not type(rot_cen) is list:
rot_cen = [rot_cen] * n
for i in range(n):
time_s = time.time()
fn = files_scan[i]
print(f'recon {i+1}/{n}: {fn.split("/")[-1]}')
recon(fn,
rot_cen[i],
sli=sli,
binning=binning,
txm_normed_flag=txm_normed_flag,
block_list=block_list,
read_full_memory=read_full_memory)
print(
f'************* take {time.time()-time_s:4.1f} sec ***************\n'
)
def align_3D_tomo_folder_mpi(file_path='.', ref_index=-1, binning=1, circle_mask_ratio=0.8, file_prefix='recon', file_type='.h5', align_coarse=1, align_method=1, hdf_attr='img', num_cpu=4):
'''
align_method:
1: old method
2: 3D cross-correlation
'''
file_path = os.path.abspath(file_path)
files_recon = pyxas.retrieve_file_type(file_path, file_prefix=file_prefix, file_type=file_type)
files_ref = files_recon[ref_index]
align_3D_tomo_file_mpi_specific(file_path, files_recon, files_ref, binning, circle_mask_ratio, file_type, align_coarse, align_method, hdf_attr, num_cpu)
def retrieve_h5_xanes_energy(file_path='.', file_prefix='recon', file_type='.h5', attr_eng='X_eng'):
file_path = os.path.abspath(file_path)
files_list = pyxas.retrieve_file_type(file_path, file_prefix=file_prefix, file_type=file_type)
num_file = len(files_list)
eng_list = np.zeros(num_file)
for i in range(num_file):
f = h5py.File(files_list[i], 'r')
eng_list[i] = float(np.array(f[attr_eng]))
np.savetxt('eng_list.txt', eng_list)
return eng_list
def get_eng_from_file(file_path, file_prefix='fly', file_type='.h5'):
import pandas as pd
files = pyxas.retrieve_file_type(file_path, file_prefix, file_type)
eng_list = []
scan_id_list =[]
for fn in files:
f = h5py.File(fn, 'r')
eng_list.append(float(np.array(f['X_eng'])))
scan_id_list.append(float(np.array(f['scan_id'])))
f.close()
res = {'scan_id_list': scan_id_list, 'eng_list':eng_list}
df = pd.DataFrame(res)
return df
def align_tomo_prj_mpi(file_path,
file_prefix='fly',
file_type='.h5',
ref_index=-1,
num_cpu=8):
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]
for i in range(n):
fn_target = files_scan[i]
if i == ref_index:
continue
print(f'aligning #{i}/{n} {fn_ref.split("/")[-1]} ')
align_two_tomo_prj_mpi(fn_ref, fn_target, num_cpu)
def batch_test():
from scipy.ndimage import generate_binary_structure
from scipy.ndimage.morphology import binary_dilation
from scipy.ndimage.filters import gaussian_filter as gf
from scipy.ndimage.filters import median_filter as mf
from skimage import io
n = 3
fn_root = '/media/mingyuan/Seagate Backup Plus Drive/TXM_2019/Ruoqian/xanes_3D'
file_scan = pyxas.retrieve_file_type(fn_root, 'recon', 'h5')
for i in range(len(file_scan)):
print(f'processing {i+1} / {len(file_scan)} ...')
ts = time.time()
scan_id = file_scan[i].split('.')[0][-5:]
fn = file_scan[i]
img = pyxas.get_img_from_hdf_file(fn, 'img')['img']
img_blur = gf(img, 5)
img_blur[img_blur < 0.0005] = 0
#res = rm_image_bkg(img_blur, multiply_factor=1, n_jobs=4)
#mask = extract_particles(res['img_labels'], n_particles=n)
t = np.array(img_blur > 0, dtype=np.int8)
mask = extract_particles(t, n_particles=n)
fn1 = fn_root + f'/1/{scan_id}'
fn2 = fn_root + f'/2/{scan_id}'
fn3 = fn_root + f'/3/{scan_id}'
coord = {}
coord['0'], coord['1'], coord['2'] = [], [], []
for j in range(n):
t = mask[f'{j}'].image
s1 = t.shape
t1 = np.zeros([500, 500])
t1[:s1[1], :s1[2]] = t[len(t) // 2]
io.imsave(f'{fn_root}/{j}/{scan_id}.tiff',
np.array(t1, dtype=np.float32))
coord[f'{j}'].append(list(mask[f'{j}'].bbox))
te = time.time()
print(f'#{i+1} / {len(file_scan)} takes {te-ts:4.1f} sec')
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'
)
def fit_xanes2D_align_tomo_recon_along_axis(file_path='.', ref_index=-1, circle_mask_ratio=0.6, file_prefix='ali', file_type='.h5', axes=[0]):
'''
Aligning the reconstructed tomo with assigned 3D reconstruction along given axis. It will project the 3D data along given axis to find the shifts
Inputs:
-----------
file_path: str
Directory contains all "fly_scans"
binning: int
binning of reconstruction
ref_index: int
index of "fly_scans" which is assigned as reference projections
this fly_scan should has has good reconstruction quality and fare full list of rotation angles
if -1: use the last scan (sorted by alphabetic file name)
file_prefix: str
prefix of the "fly_scan"
e.g., 'fly'
file_type: str
e.g. '.h5'
axis: int
along which axis to project the 3D reconstruction to find image shifts
0, or 1, or 2
Output:
----------------
None, will save aligned 3D reconstruction in folder of "{file_path}/align_axis_{axis}"
'''
file_path = os.path.abspath(file_path)
files_recon = pyxas.retrieve_file_type(file_path, file_prefix=file_prefix, file_type=file_type)
f = h5py.File(files_recon[ref_index], 'r')
rec0 = np.array(f['img'])
#rec0 = pyxas.circ_mask(rec0, axis=0, ratio=0.8, val=0)
f.close()
s = rec0.shape
stack_range = [int(s[0]*(0.5-ratio/2)), int(s[0]*(0.5+ratio/2))]
#prj0 = np.sum(rec0[stack_range[0]:stack_range[1]], axis=axis)
num_files = len(files_recon)
sr = StackReg(StackReg.TRANSLATION)
for j in range(len(files_recon)):
fn = files_recon[j]
f = h5py.File(fn, 'r')
rec1 = np.array(f['img'])
scan_id = np.array(f['scan_id'])
eng1 = np.array(f['XEng'])
angle1 = np.array(f['angle'])
f.close()
tmp = rec1.copy()
for ax in axes:
tmp, shift_matrix = align_3D_cause_axes(rec0, tmp, circle_mask_ratio, ax)
print(f'along axes = {ax}, shift by: {shift_matrix}')
new_dir = f'{file_path}/align_axis_{axis}'
fn_save = f'{new_dir}/fn'
with h5py.File(fn_save, 'w') as hf:
hf.create_dataset('img', data = tmp)
hf.create_dataset('scan_id', data = scan_id)
hf.create_dataset('XEng', data = eng1)
hf.create_dataset('angle', data = angle1)
hf.create_dataset('shift', data = np.array([r, c]))
hf.create_dataset('axis', data = axis)
print(f'{fn_save} saved \n')
np.savetxt(f'{new_dir}/rshft_axis_{axis}.txt', rshft)
np.savetxt(f'{new_dir}/cshft_axis_{axis}.txt', cshft)
del tmp
def fit_xanes2D_align_tomo_proj(file_path='.', files_scan=[], binning=2, ref_index=-1, ref_rot_cen=-1, block_list=[], sli=[], ratio=0.8, file_prefix='fly', file_type='.h5'):
'''
Aligning the tomo-scan projections with assigned scan file, and generate 3D reconstruction.
Inputs:
-----------
file_path: str
Directory contains all "fly_scans"
binning: int
binning of reconstruction
ref_index: int
index of "fly_scans" which is assigned as reference projections
this fly_scan should has has good reconstruction quality and fare full list of rotation angles
if -1: use the last scan (sorted by alphabetic file name)
ref_rot_cen: float
rotation center for the referenced "fly_scan"
if -1: find rotation center using cross-corelationship at angle-0 and angle-180
block_list: list
indexes of bad projection
e.g., list(np.arange(380,550)
ratio: float: (0 < ratio < 1)
faction of projection image to be use to align projections
e.g., 0.6
file_prefix: str
prefix of the "fly_scan"
e.g., 'fly'
file_type: str
e.g. '.h5'
Output:
----------------
None, will save aligned 3D reconstruction in folder of "{file_path}/ali_recon"
'''
file_path = os.path.abspath(file_path)
binning = int(binning)
if len(files_scan) == 0:
files_scan = pyxas.retrieve_file_type(file_path, file_prefix=file_prefix, file_type=file_type)
num_files = len(files_scan)
# block_list=list(np.arange(380,550))
fn_ref = files_scan[ref_index]
f_ref = h5py.File(fn_ref, 'r')
img_ref, _, angle_ref = pyxas.retrieve_norm_tomo_image(fn_ref, index=ref_index, binning=binning)
theta_ref = angle_ref / 180 * np.pi
f_ref.close()
if ref_rot_cen == -1:
rot_cen = find_rot(fn_ref) / binning
else:
rot_cen = ref_rot_cen / binning
sr = StackReg(StackReg.TRANSLATION)
new_dir = f'{file_path}/ali_recon'
if not os.path.exists(new_dir):
os.makedirs(new_dir)
for i in range(num_files):
fn = files_scan[i]
f1 = h5py.File(fn, 'r')
scan_id = np.array(f1['scan_id'])
angle1 = np.array(f1['angle'])
theta1 = angle1 / 180 * np.pi
f1.close()
num_angle = len(angle1)
s = time.time()
img1_ali, eng1, rshft, cshft = pyxas.align_proj_sub(fn_ref, angle_ref, fn, angle1, binning, ratio=ratio, sli=sli, ali_method='stackreg')
print(f'#{i}/{num_files}, time elapsed: {time.time()-s}\n')
img1_ali = pyxas.norm_txm(img1_ali)
rec = pyxas.recon_sub(img1_ali, theta1, rot_cen, block_list)
rec = pyxas.circ_mask(rec, axis=0, ratio=ratio, val=0)
print('saving files...\n')
fn_save = f'{new_dir}/ali_recon_{scan_id}.h5'
with h5py.File(fn_save, 'w') as hf:
hf.create_dataset('img', data = rec.astype(np.float32))
hf.create_dataset('scan_id', data = scan_id)
hf.create_dataset('XEng', data = eng1)
hf.create_dataset('angle', data = angle1)
print(f'{fn_save} saved\n')