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