def tilt_mask(size, tilt_ang1, tilt_ang2=None, tilt_axis=1, light_axis=2, sphere_mask=True):
assert tilt_axis != light_axis
if tilt_ang2 is None:
tilt_ang2 = float(N.abs(tilt_ang1))
tilt_ang1 = -tilt_ang2
else:
assert tilt_ang1 < 0
assert tilt_ang2 > 0
tilt_ang1 = (tilt_ang1 / 180.0) * N.pi
tilt_ang2 = (tilt_ang2 / 180.0) * N.pi
g = AIVU.grid_displacement_to_center(size=size, mid_co=IVU.fft_mid_co(siz=size))
plane_axis = set([0,1,2])
plane_axis.difference_update([light_axis,tilt_axis])
assert len(plane_axis) == 1
plane_axis = list(plane_axis)[0]
x_light = g[light_axis]
x_plane = g[plane_axis]
m = N.zeros(size, dtype=float)
m[ N.logical_and(x_light <= (N.tan(tilt_ang1) * x_plane), x_light >= (N.tan(tilt_ang2) * x_plane))] = 1.0
m[ N.logical_and(x_light >= (N.tan(tilt_ang1) * x_plane), x_light <= (N.tan(tilt_ang2) * x_plane))] = 1.0
if sphere_mask: m *= MU.sphere_mask(m.shape)
return m
def bandpass_denoising(a, name, save_flag=False):
b_time = time.time()
grid = GV.grid_displacement_to_center(a.shape, GV.fft_mid_co(a.shape))
rad = GV.grid_distance_to_center(grid)
rad = np.round(rad).astype(np.int)
# create a mask that only center frequencies components will be left
curve = np.zeros(rad.shape)
# TODO: change the curve value as desired
curve[int(rad.shape[0] / 8) * 3:int(rad.shape[0] / 8) * 5,
int(rad.shape[1] / 8) * 3:int(rad.shape[1] / 8) * 5,
int(rad.shape[2] / 8) * 3:int(rad.shape[2] / 8) * 5] = 1
#perform FFT and filter the data with the mask and then transform the filtered data back
vf = ifftn(ifftshift((fftshift(fftn(a)) * curve)))
vf = np.real(vf)
end_time = time.time()
print('Bandpass de-noise takes', end_time - b_time, 's')
if save_flag:
img = (vf[:, :, int(vf.shape[2] / 2)]).copy()
# TODO: Change the image and tomogram saving path
img_path = '/Users/apple/Desktop/Lab/Zach_Project/Denoising_Result/Bandpass/' + str(
name) + '_BP.png'
plt.imsave(img_path, img, cmap='gray')
mrc_path = '/Users/apple/Desktop/Lab/Zach_Project/Denoising_Result/Bandpass/' + str(
name) + '_BP.mrc'
io_file.put_mrc_data(vf, mrc_path)
return img
def sphere_mask(shape, center=None, radius=None, smooth_sigma=None):
shape = N.array(shape)
v = N.zeros(shape)
if center is None:
center = (
shape - 1
) / 2.0 # IMPORTANT: following python convension, in index starts from 0 to size-1 !!! So (siz-1)/2 is real symmetry center of the volume
center = N.array(center)
if radius is None: radius = N.min(shape / 2.0)
grid = gv.grid_displacement_to_center(shape, mid_co=center)
dist = gv.grid_distance_to_center(grid)
v[dist <= radius] = 1.0
if smooth_sigma is not None:
assert smooth_sigma > 0
v_s = N.exp(-((dist - radius) / smooth_sigma)**2)
v_s[v_s < N.exp(
-3
)] = 0.0 # use a cutoff of -3 looks nicer, although the tom toolbox uses -2
v[dist >= radius] = v_s[dist >= radius]
return v
def wedge_mask(size,
ang1,
ang2=None,
tilt_axis=1,
sphere_mask=True,
verbose=False):
# should define both tilt axis and electron beam (missing wedge) direction
warnings.warn("The definition of wedge mask is still ambiguous")
if ang2 is None:
ang2 = float(N.abs(ang1))
ang1 = -ang2
else:
assert ang1 < 0
assert ang2 > 0
if verbose:
print('image.vol.wedge.util.wedge_mask()', 'ang1', ang1, 'ang2', ang2,
'tilt_axis', tilt_axis, 'sphere_mask', sphere_mask)
ang1 = (ang1 / 180.0) * N.pi
ang2 = (ang2 / 180.0) * N.pi
g = AIVU.grid_displacement_to_center(size=size,
mid_co=AIVU.fft_mid_co(siz=size))
if tilt_axis == 0:
# y-z plane
# y axis
x0 = g[1]
# z axis
x1 = g[2]
elif tilt_axis == 1:
# x-z plane
# x axis
x0 = g[0]
# z axis
x1 = g[2]
elif tilt_axis == 2:
# x-y plane
# x axis
x0 = g[0]
# y axis
x1 = g[1]
m = N.zeros(size, dtype=float)
m[N.logical_and(x0 >= (N.tan(ang2) * x1), x0 >= (N.tan(ang1) * x1))] = 1.0
m[N.logical_and(x0 <= (N.tan(ang1) * x1), x0 <= (N.tan(ang2) * x1))] = 1.0
if sphere_mask:
m *= MU.sphere_mask(m.shape)
return m
def filter_given_curve(v, curve):
grid = GV.grid_displacement_to_center(v.shape, GV.fft_mid_co(v.shape))
rad = GV.grid_distance_to_center(grid)
rad = N.round(rad).astype(N.int)
b = N.zeros(rad.shape)
for (i, a) in enumerate(curve):
b[(rad == i)] = a
vf = ifftn(ifftshift((fftshift(fftn(v)) * b)))
vf = N.real(vf)
return vf
def gauss_function(size, sigma):
grid = gv.grid_displacement_to_center(size)
dist_sq = gv.grid_distance_sq_to_center(grid)
del grid
# gauss function
g = (1 / ((2 * N.pi)**(3.0 / 2.0) *
(sigma**3))) * N.exp(-(dist_sq) / (2.0 * (sigma**2)))
return g
def difference_of_gauss_function(size, sigma1, sigma2):
grid = AIVU.grid_displacement_to_center(size)
dist_sq = AIVU.grid_distance_sq_to_center(grid)
del grid
# gauss function
dog = (1 / ((2 * N.pi)**(3.0 / 2.0) *
(sigma1**3))) * N.exp(-(dist_sq) / (2.0 * (sigma1**2)))
dog -= (1 / ((2 * N.pi)**(3.0 / 2.0) *
(sigma2**3))) * N.exp(-(dist_sq) / (2.0 * (sigma2**2)))
return dog
def ssnr__get_rad(siz):
grid = GV.grid_displacement_to_center(siz, GV.fft_mid_co(siz))
rad = GV.grid_distance_to_center(grid)
return rad