def testAnisotropicNoiseVariance(self):
noise_estimator = AnisotropicNoiseEstimator(self.sim, batchSize=512)
noise_variance = noise_estimator.estimate()
self.assertTrue(
np.allclose(
0.04534846544265747,
noise_variance,
atol=utest_tolerance(self.sim.dtype),
))
def denoise(
data_folder,
starfile_in,
starfile_out,
pixel_size,
max_rows,
max_resolution,
noise_type,
denoise_method,
):
"""
Denoise the images and output the clean images using the default CWF method.
"""
# Create a source object for 2D images
logger.info(
f"Read in images from {starfile_in} and preprocess the images.")
source = RelionSource(starfile_in,
data_folder,
pixel_size=pixel_size,
max_rows=max_rows)
logger.info(f"Set the resolution to {max_resolution} X {max_resolution}")
if max_resolution < source.L:
# Downsample the images
source.downsample(max_resolution)
else:
logger.warn(
f"Unable to downsample to {max_resolution}, using {source.L}")
source.cache()
# Specify the fast FB basis method for expending the 2D images
basis = FFBBasis2D((source.L, source.L))
# Estimate the noise of images
noise_estimator = None
if noise_type == "White":
logger.info("Estimate the noise of images using white noise method")
noise_estimator = WhiteNoiseEstimator(source)
elif noise_type == "Anisotropic":
logger.info("Estimate the noise of images using anisotropic method")
noise_estimator = AnisotropicNoiseEstimator(source)
else:
raise RuntimeError(f"Unsupported noise_type={noise_type}")
# Whiten the noise of images
logger.info("Whiten the noise of images from the noise estimator")
source.whiten(noise_estimator.filter)
if denoise_method == "CWF":
logger.info("Denoise the images using CWF cov2D method.")
denoiser = DenoiserCov2D(source, basis)
denoised_src = denoiser.denoise(batch_size=512)
denoised_src.save(starfile_out,
batch_size=512,
save_mode="single",
overwrite=False)
else:
raise NotImplementedError(
"Currently only covariance Wiener filtering method is supported")
def testWhiten(self):
noise_estimator = AnisotropicNoiseEstimator(self.sim)
self.sim.whiten(noise_estimator.filter)
imgs_wt = self.sim.images(start=0, num=self.n).asnumpy()
# calculate correlation between two neighboring pixels from background
corr_coef = np.corrcoef(imgs_wt[:, self.L - 1, self.L - 1],
imgs_wt[:, self.L - 2, self.L - 1])
# correlation matrix should be close to identity
self.assertTrue(np.allclose(np.eye(2), corr_coef, atol=1e-1))
def testParseval(self):
"""
Here we construct a source of white noise.
Then code tests that the average noise power in the real domain,
is equivalent to the sum of the magnitudes squared
of all frequency coefficients in the Fourier domain.
These are calculated by WhiteNoiseEstimator and
AnisotropicNoiseEstimator respectively.
See Parseval/Plancherel's Theorem.
"""
wht_noise = np.random.randn(1024, 128, 128).astype(self.dtype)
src = ArrayImageSource(wht_noise)
wht_noise_estimator = WhiteNoiseEstimator(src, batchSize=512)
wht_noise_variance = wht_noise_estimator.estimate()
noise_estimator = AnisotropicNoiseEstimator(src, batchSize=512)
noise_variance = noise_estimator.estimate()
self.assertTrue(np.allclose(noise_variance, wht_noise_variance))
source = RelionSource(
"data/sample_relion_data.star",
pixel_size=1.338,
max_rows=10000,
)
# Reduce the resolution
L = 12 # You may try 16 but it takes a significant amount of time.
source.downsample(L)
# %%
# Noise Estimation and Whitening
# ------------------------------
# Estimate noise in the ImageSource instance
noise_estimator = AnisotropicNoiseEstimator(source)
# Apply whitening to ImageSource
source.whiten(noise_estimator.filter)
# Display subset of the images
images = source.images(0, 10)
images.show()
# %%
# Estimate Mean Volume
# --------------------
# We'll create a 3D Fourier Bessel Basis corresponding to volume resolution L.
basis = FBBasis3D((L, L, L))
# Estimate mean Volume
mean_estimator = MeanEstimator(source, basis, batch_size=8192)
# In this case, we'll just use our Image in an ArrayImageSource to run a small experiment.
#
# If you were happy with the experiment design on an array of test data,
# the source is easily swapped out to something like RelionSource,
# which might point at a stack of images too large to fit in memory at once.
# The implementation of batching for memory management
# would be managed behind the scenes for you.
imgs_src = ArrayImageSource(imgs_with_noise)
# We'll copy the orginals for comparison later, before we process them further.
noisy_imgs_copy = imgs_src.images(0, n_imgs)
# One of the tools we can use is a NoiseEstimator,
# which consumes from a Source.
noise_estimator = AnisotropicNoiseEstimator(imgs_src)
# Once we have the estimator instance,
# we can use it in a transform applied to our Source.
imgs_src.whiten(noise_estimator.filter)
# Peek at two whitened images and their corresponding spectrum.
fig, axs = plt.subplots(2, 2)
for i, img in enumerate(imgs_src.images(0, 2)):
axs[0, i].imshow(img, cmap=plt.cm.gray)
axs[0, i].set_title(f"Whitened Noisy Image {i}")
img_with_noise_f = np.abs(np.fft.fftshift(np.fft.fft2(img)))
axs[1, i].imshow(np.log(1 + img_with_noise_f), cmap=plt.cm.gray)
axs[1, i].set_title(f"Whitened Noisy Image Spectrum {i}")
plt.tight_layout()
def testAnisotropicNoisePSD(self):
noise_estimator = AnisotropicNoiseEstimator(self.sim, batchSize=512)
noise_psd = noise_estimator.estimate_noise_psd()
self.assertTrue(
np.allclose(
noise_psd,
[
[
+0.00112473,
+0.00106200,
+0.00118618,
+0.00495772,
+0.00495797,
+0.00495772,
+0.00118618,
+0.00106200,
],
[
+0.00099063,
+0.00113591,
+0.00160692,
+0.00462546,
+0.00621764,
+0.00475203,
+0.00153705,
+0.00116489,
],
[
+0.00113134,
+0.00148855,
+0.00267187,
+0.00505812,
+0.01086790,
+0.00520619,
+0.00271356,
+0.00157493,
],
[
+0.00485551,
+0.00453407,
+0.00499355,
+0.00672553,
+0.01090170,
+0.00696211,
+0.00501925,
+0.00460892,
],
[
+0.00506158,
+0.00629060,
+0.01099897,
+0.01099300,
+0.04534847,
+0.01099300,
+0.01099897,
+0.00629060,
],
[
+0.00485551,
+0.00460892,
+0.00501925,
+0.00696211,
+0.01090170,
+0.00672553,
+0.00499355,
+0.00453407,
],
[
+0.00113134,
+0.00157493,
+0.00271356,
+0.00520619,
+0.01086790,
+0.00505812,
+0.00267187,
+0.00148855,
],
[
+0.00099063,
+0.00116489,
+0.00153705,
+0.00475203,
+0.00621764,
+0.00462546,
+0.00160692,
+0.00113591,
],
],
))
# Create yet another Simulation source to tinker with.
sim4 = Simulation(L=v2.resolution,
n=num_imgs,
vols=v2,
noise_filter=custom_noise_filter)
sim4.images(0, 10).show()
# %%
# Noise Whitening
# ---------------
#
# Applying the ``Simulation.whiten()`` method just requires passing the filter corresponding to the estimated noise instance.
# Then we can inspect some of the whitened images. While noise is still present, we can see a dramatic change.
# Estimate noise.
aiso_noise_estimator = AnisotropicNoiseEstimator(sim4)
# Whiten based on the estimated noise
sim4.whiten(aiso_noise_estimator.filter)
# What do the whitened images look like...
sim4.images(0, 10).show()
# %%
# Homework Task 4
# ^^^^^^^^^^^^^^^
#
# Try some other image preprocessing methods exposed by the ``Simulation``/``ImageSource`` classes.
#
# Try some other custom function to add noise or other corruptions to the images.