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)