def deconvolve(vol_a: Volume, vol_b: Volume, n: int, psf_a: np.ndarray,
psf_b: np.ndarray) -> Volume:
"""
Perform joint Richardson-Lucy deconvolution on two volumes using two specified PSFs on the CPU
:param vol_a: The first volume
:param vol_b: The second volume
:param n: The number of Richardson-Lucy iterations
:param psf_a: The PSF for the first volume
:param psf_b: The PSF for the second volume
"""
# from astropy.convolution import convolve_fft
from functools import partial
from scipy.signal import fftconvolve
view_a, view_b = vol_a.astype(np.float), vol_b.astype(np.float)
psf_a = psf_a.astype(np.float) / np.sum(psf_a).astype(np.float)
psf_b = psf_b.astype(np.float) / np.sum(psf_b).astype(np.float)
psf_Ai = psf_a[::-1, ::-1, ::-1]
psf_Bi = psf_b[::-1, ::-1, ::-1]
estimate = (view_a + view_b) / 2
convolve = partial(fftconvolve, mode='same')
with progressbar.ProgressBar(max_value=n, redirect_stderr=True) as bar:
for _ in bar(range(n)):
estimate = estimate * convolve(
view_a / (convolve(estimate, psf_a) + 1e-6), psf_Ai)
estimate = estimate * convolve(
view_b / (convolve(estimate, psf_b) + 1e-6), psf_Bi)
CURSOR_UP_ONE = '\x1b[1A'
ERASE_LINE = '\x1b[2K'
print(CURSOR_UP_ONE + ERASE_LINE + CURSOR_UP_ONE)
# TODO: Rescaling might be unwanted
e_min, e_max = np.percentile(estimate, [0.05, 99.95])
estimate = ((np.clip(estimate, e_min, e_max) - e_min) / (e_max - e_min) *
(2**16 - 1)).astype(np.uint16)
return Volume(estimate,
inverted=False,
spacing=vol_a.spacing,
is_skewed=False)
def deconvolve_single_gpu(vol_a: Volume, n: int, psf_a: np.ndarray) -> Volume:
"""
Perform joint Richardson-Lucy deconvolution on two volumes using two specified PSFs on the GPU
:param vol_a: The first volume
:param n: The number of Richardson-Lucy iterations
:param psf_a: The PSF for the first volume
:return: The fused RL deconvolution
"""
from functools import partial
from dispim.metrics import DECONV_MSE_DELTA
import arrayfire as af
print(vol_a.shape)
view_a = vol_a.astype(np.float)
psf_a = psf_a.astype(np.float) / np.sum(psf_a).astype(np.float)
psf_Ai = psf_a[::-1, ::-1, ::-1]
view_a = af.cast(af.from_ndarray(np.array(view_a)), af.Dtype.f32)
psf_a = af.cast(af.from_ndarray(psf_a), af.Dtype.f32)
psf_Ai = af.cast(af.from_ndarray(psf_Ai), af.Dtype.f32)
estimate = view_a
convolve = partial(af.fft_convolve3)
with progressbar.ProgressBar(max_value=n, redirect_stderr=True) as bar:
for _ in bar(range(n)):
if metrack.is_tracked(DECONV_MSE_DELTA):
prev = estimate
estimate = estimate * convolve(
view_a / (convolve(estimate, psf_a) + 1), psf_Ai)
if metrack.is_tracked(DECONV_MSE_DELTA):
metrack.append_metric(DECONV_MSE_DELTA,
(_, float(np.mean(
(prev - estimate)**2))))
CURSOR_UP_ONE = '\x1b[1A'
ERASE_LINE = '\x1b[2K'
print(CURSOR_UP_ONE + ERASE_LINE + CURSOR_UP_ONE)
logger.debug(
f'Deconved min: {np.min(estimate)}, max: {np.max(estimate)}, has nan: {np.any(np.isnan(estimate))}'
)
result = estimate.to_ndarray()
del estimate
return Volume(result.astype(np.uint16),
inverted=False,
spacing=vol_a.spacing,
is_skewed=False)
def deconvolve_gpu_blind(vol_a: Volume, vol_b: Volume, n: int, m: int,
psf_a: np.ndarray, psf_b: np.ndarray) -> Volume:
"""
Perform blind joint Richardson-Lucy deconvolution on two volumes using two specified estimates of the PSF on the
GPU
:param vol_a: The first volume
:param vol_b: The second volume
:param n: The number of Richardson-Lucy iterations
:param m: The number of sub-iterations per RL iteration
:param psf_a: The initial PSF estimate for the first volume
:param psf_b: The initial PSF estimate for the second volume
:return: The fused RL deconvolution
"""
from functools import partial
import arrayfire as af
view_a, view_b = vol_a.astype(np.float), vol_b.astype(np.float)
psf_a = psf_a.astype(np.float) / np.sum(psf_a).astype(np.float)
psf_b = psf_b.astype(np.float) / np.sum(psf_b).astype(np.float)
padding = tuple(
(int(s // 2 - psf_a.shape[i]), int((s - s // 2) - psf_a.shape[i]))
for i, s in enumerate(view_a.shape))
psf_a = np.pad(
psf_a,
tuple(((s - psf_a.shape[i]) // 2,
(s - psf_a.shape[i]) - (s - psf_a.shape[i]) // 2)
for i, s in enumerate(view_a.shape)), 'constant')
psf_b = np.pad(
psf_b,
tuple(((s - psf_b.shape[i]) // 2,
(s - psf_b.shape[i]) - (s - psf_b.shape[i]) // 2)
for i, s in enumerate(view_b.shape)), 'constant')
view_a = af.cast(af.from_ndarray(view_a), af.Dtype.u16)
view_b = af.cast(af.from_ndarray(view_b), af.Dtype.u16)
psf_a = af.cast(af.from_ndarray(psf_a), af.Dtype.f32)
psf_b = af.cast(af.from_ndarray(psf_b), af.Dtype.f32)
estimate = (view_a + view_b) / 2
convolve = partial(af.fft_convolve3)
lamb = 0.002
with progressbar.ProgressBar(max_value=n, redirect_stderr=True) as bar:
for _ in bar(range(n)):
for j in range(m):
psf_a = psf_a * convolve(
view_a / (convolve(psf_a, estimate) + 1e-1),
estimate[::-1, ::-1, ::-1])
for j in range(m):
estimate = estimate * convolve(
view_a /
(convolve(estimate, psf_a) + 10), psf_a[::-1, ::-1, ::-1])
for j in range(m):
psf_b = psf_b * convolve(
view_b / (convolve(psf_b, estimate) + 1e-1),
estimate[::-1, ::-1, ::-1])
for j in range(m):
estimate = estimate * convolve(
view_b /
(convolve(estimate, psf_b) + 10), psf_b[::-1, ::-1, ::-1])
del psf_a, psf_b, view_a, view_b
CURSOR_UP_ONE = '\x1b[1A'
ERASE_LINE = '\x1b[2K'
print(CURSOR_UP_ONE + ERASE_LINE + CURSOR_UP_ONE)
return Volume(estimate.to_ndarray(),
inverted=False,
spacing=vol_a.spacing,
is_skewed=False)
