def pyramid_data(self):
# get the pyramid data
# method 1, pyramid then stack the image, which means the template window size is increasing for different pyramid level
# get the pyramid wrapping images
ref_pyramid = []
img_pyramid = []
prColor(
'obtain pyramid image and stack the window with pyramid level: {}'.
format(self.pyramid_level), 'green')
ref_pyramid.append(self.ref_data)
img_pyramid.append(self.img_data)
for kk in range(self.pyramid_level):
ref_pyramid.append(
pywt.dwtn(ref_pyramid[kk], 'db3', mode='zero',
axes=(-2, -1))['aa'])
img_pyramid.append(
pywt.dwtn(img_pyramid[kk], 'db3', mode='zero',
axes=(-2, -1))['aa'])
normlize_std = lambda img: (
(img - np.ndarray.mean(img, axis=0)) / np.ndarray.std(img, axis=0))
ref_pyramid = [
normlize_std(self.template_stack(img_data))
for img_data in ref_pyramid
]
img_pyramid = [
normlize_std(self.template_stack(img_data))
for img_data in img_pyramid
]
return ref_pyramid, img_pyramid
def load_image(file_path):
if os.path.exists(file_path):
img = np.array(Image.open(file_path))
else:
prColor('Error: wrong data path. No data is loaded.', 'red')
sys.exit()
return np.array(img)
def load_images(Folder_path, filename_format='*.tif'):
f_list = glob.glob(os.path.join(Folder_path, filename_format))
f_list = sorted(f_list)
img = []
for f_single in f_list:
img.append(np.array(Image.open(f_single)))
# prColor('load image: {}'.format(f_single), 'green')
if len(img) == 0:
prColor('Error: wrong data path. No data is loaded.', 'red')
sys.exit()
return np.array(img)
def pyramid_data(self):
# data normalization
ref_data = ((self.ref_data - np.ndarray.mean(self.ref_data, axis=0)) /
np.ndarray.std(self.ref_data, axis=0))
img_data = ((self.img_data - np.ndarray.mean(self.img_data, axis=0)) /
np.ndarray.std(self.img_data, axis=0))
ref_pyramid = []
img_pyramid = []
prColor(
'obtain pyramid image with pyramid level: {}'.format(
self.pyramid_level), 'green')
ref_pyramid.append(ref_data)
img_pyramid.append(img_data)
for kk in range(self.pyramid_level):
ref_pyramid.append(
pywt.dwtn(ref_pyramid[kk], 'db3', mode='zero',
axes=(-2, -1))['aa'])
img_pyramid.append(
pywt.dwtn(img_pyramid[kk], 'db3', mode='zero',
axes=(-2, -1))['aa'])
return ref_pyramid, img_pyramid
def solver(self):
ref_pyramid, img_pyramid = self.wavelet_data()
transmission = self.img_data / self.ref_data
for attr in ('img_data', 'ref_data'):
self.__dict__.pop(attr, None)
cores = ms.cpu_count()
prColor('Computer available cores: {}'.format(cores), 'green')
if cores > self.n_cores:
cores = self.n_cores
else:
cores = ms.cpu_count()
prColor('Use {} cores'.format(cores), 'light_purple')
prColor('Process group number: {}'.format(self.n_group),
'light_purple')
if cores * self.n_group > self.M_image:
n_tasks = 4
else:
n_tasks = cores * self.n_group
start_time = time.time()
# use pyramid wrapping
max_pyramid_searching_window = int(
np.ceil(self.cal_half_window / 2**self.pyramid_level))
searching_window_pyramid_list = [self.N_s_extend
] * self.pyramid_level + [
int(max_pyramid_searching_window)
]
m, n, c = img_pyramid[0].shape
displace = [np.zeros((m, n)), np.zeros((m, n))]
for k_iter in range(self.n_iter):
# iteration to approximating the results
displace = [img / 2**self.pyramid_level for img in displace]
m, n, c = img_pyramid[-1].shape
displace[0] = self.resampling_spline(displace[0], (m, n))
displace[1] = self.resampling_spline(displace[1], (m, n))
prColor(
'down sampling the dispalce to size: {}'.format(
displace[0].shape), 'green')
displace = [
np.fmax(
np.fmin(displace[0],
self.cal_half_window / 2**self.pyramid_level),
-self.cal_half_window / 2**self.pyramid_level),
np.fmax(
np.fmin(displace[1],
self.cal_half_window / 2**self.pyramid_level),
-self.cal_half_window / 2**self.pyramid_level)
]
for p_level in range(self.pyramid_level, -1, -1):
# first pyramid, searching the window. Then search nearby
if p_level == self.pyramid_level:
pyramid_seaching_window = searching_window_pyramid_list[
p_level]
m, n, c = img_pyramid[p_level].shape
displace_pyramid = [np.round(img) for img in displace]
n_pad = int(np.ceil(self.cal_half_window / 2**p_level))
else:
pyramid_seaching_window = searching_window_pyramid_list[
p_level]
m, n, c = img_pyramid[p_level].shape
displace_pyramid = [
np.round(self.resampling_spline(img * 2, (m, n)))
for img in displace
]
displace_pyramid = [
np.fmax(
np.fmin(displace_pyramid[0],
self.cal_half_window / 2**p_level),
-self.cal_half_window / 2**p_level),
np.fmax(
np.fmin(displace_pyramid[1],
self.cal_half_window / 2**p_level),
-self.cal_half_window / 2**p_level)
]
n_pad = int(np.ceil(self.cal_half_window / 2**p_level))
prColor(
'pyramid level: {}\nImage size: {}\nsearching window:{}'.
format(p_level, ref_pyramid[p_level].shape,
pyramid_seaching_window), 'cyan')
y_axis = np.arange(ref_pyramid[p_level].shape[0])
chunks_idx_y = np.array_split(y_axis, n_tasks)
dim = img_pyramid[p_level].shape
ref_wa_pad = np.pad(ref_pyramid[p_level],
((n_pad + pyramid_seaching_window,
n_pad + pyramid_seaching_window),
(n_pad + pyramid_seaching_window,
n_pad + pyramid_seaching_window),
(0, 0)),
'constant',
constant_values=(0, 0))
# use CPU parallel to calculate
result_list = []
'''
calculate the pixel displacement for the pyramid images
'''
with concurrent.futures.ProcessPoolExecutor(
max_workers=cores) as executor:
futures = []
for y_list in chunks_idx_y:
# get the stack data
img_wa_stack = img_pyramid[p_level][y_list, :, :]
ref_wa_stack = ref_wa_pad[
y_list[0]:y_list[-1] + 2 *
(n_pad + pyramid_seaching_window) + 1, :, :]
# start the jobs
futures.append(
executor.submit(self.displace_wavelet, y_list,
img_wa_stack, ref_wa_stack,
(displace_pyramid[0][y_list, :],
displace_pyramid[1][y_list, :]),
pyramid_seaching_window, n_pad))
for future in concurrent.futures.as_completed(futures):
try:
result_list.append(future.result())
# display the status of the program
Total_iter = cores * self.n_group
Current_iter = len(result_list)
percent_iter = Current_iter / Total_iter * 100
str_bar = '>' * (int(np.ceil(
percent_iter / 2))) + ' ' * (int(
(100 - percent_iter) // 2))
prColor(
'\r' + str_bar + 'processing: [%3.1f%%] ' %
(percent_iter), 'purple')
except:
prColor('Error in the parallel calculation', 'red')
disp_y_list = [item[0] for item in result_list]
disp_x_list = [item[1] for item in result_list]
y_list = [item[2] for item in result_list]
displace_y = np.zeros((dim[0], dim[1]))
displace_x = np.zeros((dim[0], dim[1]))
for y, disp_x, disp_y in zip(y_list, disp_x_list, disp_y_list):
displace_x[y, :] = disp_x
displace_y[y, :] = disp_y
displace = [
np.fmax(
np.fmin(displace_y, self.cal_half_window / 2**p_level),
-self.cal_half_window / 2**p_level),
np.fmax(
np.fmin(displace_x, self.cal_half_window / 2**p_level),
-self.cal_half_window / 2**p_level)
]
prColor('displace map wrapping: {}'.format(displace[0].shape),
'green')
print('max of displace: {}, min of displace: {}'.format(
np.amax(displace[0]), np.amin(displace[1])))
end_time = time.time()
prColor(
'\r' + 'Processing time: {:0.3f} s'.format(end_time - start_time),
'light_purple')
displace[0] = -displace[0][self.cal_half_window:-self.cal_half_window,
self.cal_half_window:-self.cal_half_window]
displace[1] = -displace[1][self.cal_half_window:-self.cal_half_window,
self.cal_half_window:-self.cal_half_window]
DPC_y = (displace[0] - np.mean(displace[0])) * p_x / z
DPC_x = (displace[1] - np.mean(displace[1])) * p_x / z
phase = -frankotchellappa(
DPC_x, DPC_y) * self.p_x * 2 * np.pi / self.wavelength
self.time_cost = end_time - start_time
return displace, [DPC_y, DPC_x], phase, transmission
def wavelet_data(self):
# process the data to get the wavelet transform
ref_pyramid, img_pyramid = self.pyramid_data()
if self.use_wavelet:
prColor('obtain wavelet data...', 'green')
wavelet_method = 'db2'
# wavelet_method = 'bior1.3'
# wavelet wrapping level. 2 is half, 3 is 1/3 of the size
max_wavelet_level = pywt.dwt_max_level(ref_pyramid[0].shape[0],
wavelet_method)
prColor('max wavelet level: {}'.format(max_wavelet_level), 'green')
self.wavelet_level = max_wavelet_level
coefs_level = self.wavelet_level + 1 - self.wavelet_level_cut
if ref_pyramid[0].shape[0] > 150:
self.wavelet_add_list = [0, 0, 0, 0, 0, 0]
elif ref_pyramid[0].shape[0] > 50:
self.wavelet_add_list = [0, 0, 1, 2, 2, 2]
else:
self.wavelet_add_list = [2, 2, 2, 2, 2, 2]
# wavelet transform and cut for the pyramid images
start_time = time.time()
for p_level in range(len(img_pyramid)):
if p_level > len(self.wavelet_add_list):
wavelevel_add = 2
else:
wavelevel_add = self.wavelet_add_list[p_level]
img_wa, level_name = Wavelet_transform(
img_pyramid[p_level],
wavelet_method,
w_level=self.wavelet_level,
return_level=coefs_level + wavelevel_add)
img_pyramid[p_level] = img_wa
ref_wa, level_name = Wavelet_transform(
ref_pyramid[p_level],
wavelet_method,
w_level=self.wavelet_level,
return_level=coefs_level + wavelevel_add)
ref_pyramid[p_level] = ref_wa
prColor(
'pyramid level: {}\nvector length: {}\nUse wavelet coef: {}'
.format(p_level, ref_wa.shape[2], level_name), 'green')
end_time = time.time()
print('wavelet time: {}'.format(end_time - start_time))
else:
img_pyramid = [
np.moveaxis(img_data, 0, -1) for img_data in img_pyramid
]
ref_pyramid = [
np.moveaxis(img_data, 0, -1) for img_data in ref_pyramid
]
self.wavelet_level = None
self.wavelet_add_list = None
self.wavelet_level_cut = None
return ref_pyramid, img_pyramid
elif len(sys.argv) == 4:
File_img = sys.argv[1]
File_ref = sys.argv[2]
Folder_result = sys.argv[3]
# [image_size, template_window, cal_half_window, n_group, n_cores, energy, pixel_size, distance, wavelet_ct, pyramid level, n_iteration]
parameter_wavelet = [
1500, 5, 20, 4, 4, 14e3, 0.65e-6, 500e-3, 1, 2, 2, 1
]
elif len(sys.argv) == 16:
File_img = sys.argv[1]
File_ref = sys.argv[2]
Folder_result = sys.argv[3]
parameter_wavelet = sys.argv[4:]
else:
prColor('Wrong parameters! should be: sample, ref, result', 'red')
prColor('folder: {}'.format(Folder_result), 'green')
# roi of the images
M_image = int(parameter_wavelet[0])
# template window, the N_s nearby pixels used to represent the local pixel, 2*N_s+1
N_s = int(parameter_wavelet[1])
# the number of the area to calculate for each pixel, 2*cal_half_window X 2*cal_half_window
cal_half_window = int(parameter_wavelet[2])
# the calculation window for high order pyramid
N_s_extend = 4
# process number for parallel
n_cores = int(parameter_wavelet[3])
# number to reduce the each memory use
n_group = int(parameter_wavelet[4])
# [image_size, cal_half_window, ues_parallel, n_group, n_cores, energy, pixel_size, distance, wavelet_cut_level]
parameter_wavelet = [1500, 10, 1, 4, 4, 14e3, 0.65e-6, 500e-3, 1]
elif len(sys.argv) == 4:
Folder_img = sys.argv[1]
Folder_ref = sys.argv[2]
Folder_result = sys.argv[3]
# [image_size, cal_half_window, ues_parallel, n_group, n_cores, energy, pixel_size, distance, wavelet_cut_level]
parameter_wavelet = [1500, 10, 1, 4, 4, 14e3, 0.65e-6, 500e-3, 1]
elif len(sys.argv) == 14:
Folder_img = sys.argv[1]
Folder_ref = sys.argv[2]
Folder_result = sys.argv[3]
parameter_wavelet = sys.argv[4:]
else:
prColor('Wrong parameters! should be: sample, ref, result', 'red')
prColor('folder: {}'.format(Folder_result), 'green')
# roi of the images
M_image = int(parameter_wavelet[0])
# do sub pixel calculation, if it's 1, no sub pixel
pixel_sample = 1
# the number of the area to calculate for each pixel, 2*cal_half_window X 2*cal_half_window
cal_half_window = int(parameter_wavelet[1])
# parallel calculation or not
use_parallel = int(parameter_wavelet[2])
# process number for parallel
n_cores = int(parameter_wavelet[3])
# number to reduce the each memory use
n_group = int(parameter_wavelet[4])
def run(self, result_path=None):
self.displace, self.DPC, self.phase, self.transmission = self.solver()
if result_path is not None:
'''
save the calculation results
'''
if not os.path.exists(result_path):
os.makedirs(result_path)
self.result_filename = 'WSVT' + str(self.M_image) + '_px_' + str(
self.wavelet_level_cut) + 'wavelet_Cutlevel_' + str(
self.pyramid_level) + 'pyramid_level'
kk = 1
while os.path.exists(
os.path.join(
result_path,
self.result_filename + '.hdf5')) or os.path.exists(
os.path.join(result_path,
self.result_filename + '.json')):
self.result_filename = 'WSVT' + str(
self.M_image) + '_px_' + str(
self.wavelet_level_cut) + 'wavelet_Cutlevel_' + str(
self.pyramid_level
) + 'pyramid_level' + '_{}'.format(kk)
kk += 1
write_h5(
result_path, self.result_filename, {
'displace_x': self.displace[1],
'displace_y': self.displace[0],
'DPC_x': self.DPC[1],
'DPC_y': self.DPC[0],
'phase': self.phase,
'transmission_image': self.transmission
})
phase_rms = np.std(self.phase)
phase_pv = np.amax(self.phase) - np.amin(self.phase)
prColor(
'phase rms: {:0.2f}, phase PV: {:0.2f}'.format(
phase_rms, phase_pv), 'purple')
parameter_dict = {
'M_image': self.M_image,
'N_s extend': self.N_s_extend,
'half_window': self.cal_half_window,
'energy': self.energy,
'wavelength': self.wavelength,
'pixel_size': self.p_x,
'z_distance': self.z,
'cpu_cores': self.n_cores,
'n_group': self.n_group,
'wavelet_level': self.wavelet_level,
'pyramid_level': self.pyramid_level,
'n_iter': self.n_iter,
'time_cost': self.time_cost,
'use_wavelet': self.use_wavelet,
'wavelet_level_cut': self.wavelet_level_cut,
'wavelet_add': self.wavelet_add_list,
'phase_rms': phase_rms,
'phase_pv': phase_pv
}
write_json(result_path, self.result_filename, parameter_dict)