def deconvolve_gpu_chunked(vol_a: Volume, vol_b: Volume, n: int, psf_a: np.ndarray, psf_b: np.ndarray, nchunks: int, blind: bool = False) -> Volume: """ Perform joint Richardson-Lucy deconvolution on two volumes using two specified PSFs on the GPU in chunks along the z-axis :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 :param blind: Whether to perform blind RL deconvolution using the given PSFs as initial estimates :param nchunks: The number of chunks to subdivide the volume into :return: The fused RL deconvolution """ import arrayfire as af result = np.zeros(vol_a.shape, np.float32) chunk_size = vol_a.shape[2] // nchunks for i in range(nchunks): start = i * chunk_size end = (i + 1) * chunk_size if i < nchunks - 1 else vol_a.shape[2] lpad = int(psf_a.shape[2] * 4) rpad = int(psf_a.shape[2] * 4) start_exp = max(0, start - lpad) end_exp = min(vol_a.shape[2], end + rpad) with metrack.Context(f'Chunk {i}'): if not blind: chunk_est = deconvolve_gpu( Volume(vol_a[:, :, start_exp:end_exp], False, (1, 1, 1)), Volume(vol_b[:, :, start_exp:end_exp], False, (1, 1, 1)), n, psf_a, psf_b) else: chunk_est = deconvolve_gpu_blind( Volume(vol_a[:, :, start_exp:end_exp], False, (1, 1, 1)), Volume(vol_b[:, :, start_exp:end_exp], False, (1, 1, 1)), n, 5, psf_a, psf_b) af.device_gc() if end != end_exp: result[:, :, start:end] = chunk_est[:, :, start - start_exp:end - end_exp] else: result[:, :, start:end] = chunk_est[:, :, start - start_exp:] # FIXME: Proper outlier clipping! e_min, e_max = np.percentile(result, [0.002, 99.998]) result = ((np.clip(result, e_min, e_max) - e_min) / (e_max - e_min) * (2**16 - 1)).astype(np.uint16) return Volume(result, inverted=False, spacing=(1, 1, 1))
def apply_registration(vol_a: Volume, vol_b: Volume, order: int = 2) -> Tuple[Volume, Volume]: from scipy.ndimage import affine_transform min_res = min(np.min(vol_a.spacing), np.min(vol_b.spacing)) logger.debug(f"Min res: {min_res}") logger.debug(f"Grid-to-world A:\n{vol_a.grid_to_world}") logger.debug(f"Grid-to-world B:\n{vol_b.grid_to_world}") logger.debug(f"World transform:\n{vol_b.world_transform}") grid_to_world_final = np.eye(4) * np.array([min_res, min_res, min_res, 1]) transform_a = np.linalg.inv(vol_a.grid_to_world) transform_a = transform_a @ grid_to_world_final logger.debug(f'Final A transform:\n{transform_a}') # TODO: Remove this, unnecessary final_shape = (np.ceil( np.linalg.inv(transform_a) @ np.array([vol_a.shape[0], vol_a.shape[1], vol_a.shape[2], 1]))).astype( np.int)[:3] # psf = vol_a.psf # # psf_shape_old = psf.shape # psf_shape_transformed = (np.ceil( # np.linalg.inv(transform_a) @ np.array([psf.shape[0], psf.shape[1], psf.shape[2], 1]))).astype( # np.int)[:3] # # psf = affine_transform(psf, transform_a, output_shape=psf_shape_transformed, order=3) # psf = center_psf(psf, psf_shape_old[0]) vol_a = Volume(affine_transform(vol_a, transform_a, output_shape=final_shape, order=order), spacing=(min_res, min_res, min_res), is_skewed=False, inverted=False, ) transform_b = np.linalg.inv(vol_b.grid_to_world) transform_b = transform_b @ vol_b.world_transform transform_b = transform_b @ grid_to_world_final logger.debug(f'Final B transform:\n{transform_b}') # psf = vol_b.psf # # psf_shape_old = psf.shape # psf_shape_transformed = (np.ceil( # np.linalg.inv(transform_b) @ np.array([psf.shape[0], psf.shape[1], psf.shape[2], 1]))).astype( # np.int)[:3] # # psf = affine_transform(psf, transform_b, output_shape=psf_shape_transformed, order=3) # psf = center_psf(psf, psf_shape_old[0]) transformed = affine_transform(vol_b, transform_b, output_shape=final_shape, order=order) logger.debug(f'Transformed A average value: {np.mean(vol_a)}') logger.debug(f'Transformed B average value: {np.mean(transformed)}') return vol_a, Volume(transformed, spacing=vol_a.spacing, is_skewed=False, inverted=False)
def test_reigster_dipy(test_data, offset): a = Volume(test_data[:42, :42, :42], is_skewed=False) b = Volume(test_data[offset[0]:42 + offset[0], offset[1]:42 + offset[1], offset[2]:42 + offset[2]], is_skewed=False) assert np.allclose(a.grid_to_world, np.eye(4)) assert np.allclose(b.grid_to_world, np.eye(4)) dispim.register.register_com(a, b) dispim.register.register_dipy(a, b, crop=1.0) dispim.register.register_dipy(a, b, crop=1.0) estimated_offset = -(b.world_transform[:3, 3]) assert np.allclose(estimated_offset, offset, atol=0.2, rtol=0.05)
def unshift_fast_diag(vol: Volume, invert: bool = False, estimate_true_interval: bool = True, rotate: bool = True) -> Volume: if estimate_true_interval: interval = compute_true_interval(vol, invert) vol = Volume(vol, spacing=vol.spacing[:2] + (interval,)) logger.debug('Estimated volume interval: {}'.format(interval)) # FIXME: Metadata is lost here if invert: data = unshift_fast_numbai_diag(np.array(vol), vol.spacing) if rotate: data = np.rot90(data, k=2, axes=(1, 2)) return Volume(data, inverted=False, is_skewed=False) else: return Volume(unshift_fast_numba_diag(np.array(vol), vol.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 generate_spheres_test_volume(name: str, size: int = 32, num_beads: int = 32, sigmas: Tuple[float, float, float] = (0, 0, 2)): data = np.zeros((size, size, size)) for i in range(num_beads): p = np.random.randint(8, size - 8, 3) data[p[0] - 8:p[0] + 8, p[1] - 8:p[1] + 8, p[1] - 8:p[1] + 8] = pymrt.geometry.sphere(16, 0.5, 8) vol = Volume(data, (1, 1, 1)) vol.save_tiff('test_' + name) vol.save_tiff_single('test_single_' + name) blurred = gaussian_filter(data, sigmas, mode='constant') blurred_vol = Volume(blurred, (1, 1, 1)) blurred_vol.save_tiff('btest_' + name) blurred_vol.save_tiff_single('btest_single_' + name) return data
def generate_bead_test_volume(name: str, size: Tuple[int, int, int] = (32, 32, 32), num_beads: int = 32, sigmas: Tuple[float, float, float] = (0, 0, 2)): data = np.zeros(size) mask = np.array([True] * num_beads + [False] * (np.prod(data.size) - num_beads)) np.random.shuffle(mask) mask = mask.reshape(data.shape) data[mask] = 1 blurred = gaussian_filter(data, sigmas, mode='constant') blurred_vol = Volume(blurred, (1, 1, 1)) blurred_vol.save_tiff('test_' + name) blurred_vol.save_tiff_single('test_single_' + name)
def test_register2d(fractal_noise, offset): from scipy.ndimage import affine_transform a = Volume(fractal_noise[:32, :32, :32], is_skewed=False) b = Volume(fractal_noise[offset[0]:32 + offset[0], offset[1]:32 + offset[1], offset[2]:32 + offset[2]], is_skewed=False) assert np.allclose(a.grid_to_world, np.eye(4)) assert np.allclose(b.grid_to_world, np.eye(4)) dispim.register.register_2d(a, b, axis=2) b_transformed = affine_transform(b, b.world_transform, order=1) estimated_offset = -(b.world_transform[:3, 3]) expected_mismatches = 32 * offset[0] + 32 * offset[1] - offset[0] * offset[ 1] expected_mismatches *= 32 expected_mismatches *= 1.02 expected_mismatches += 1500 assert np.allclose(estimated_offset, offset, atol=0.2, rtol=0.05)
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 test_io(): try: from dispim.io import save_tiff, load_tiff a = Volume((np.random.rand(64, 64, 64) * 2000).astype(np.uint16), spacing=(0.3, 0.3, 0.7)) save_tiff(a, 'temp.tif') b: Volume = load_tiff('temp.tif', step_size=0.7) assert np.allclose(a, b) assert a.dtype == np.uint16 assert b.dtype == np.uint16 assert type(b) == Volume assert np.allclose(a.spacing, b.spacing) finally: import os if os.path.exists("temp.tif"): os.remove("temp.tif")
def rand_vol_a(): return Volume( np.random.randint(0, 2**16 - 1, size=(64, 64, 64), dtype=np.uint16))
def process(self, data: ProcessData) -> ProcessData: if len(data) == 2: return dispim.make_isotropic(data[0], data[1]) else: return dispim.make_isotropic(data[0], Volume(np.empty((1, 1, 1)), False, (1, 1, 1)))[0],
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)
def register_dipy( vol_a: Volume, vol_b: Volume, sampling_prop: float = 1.0, crop: float = 0.8, transform_cls: type = TranslationTransform3D) -> Tuple[Volume, Volume]: """ Perform registration on two volumes :param vol_a: The fixed volume :param vol_b: The moving volume :param sampling_prop: Proportion of data to be used (0-1) :param crop: The value to use for cropping both volumes before registration (0-1 :param transform_cls: The type of registration transform to compute :return: The updated volumes """ from dipy.align.imaffine import (MutualInformationMetric, AffineRegistration) from dispim.util import crop_view from dispim.metrics import MUTUAL_INFORMATION_METRIC, MUTUAL_INFORMATION_GRADIENT_METRIC logger.debug('Sampling prop: ' + str(sampling_prop)) logger.debug('Crop: ' + str(crop)) subvol_a = crop_view(vol_a, crop, center_crop=False) subvol_b = crop_view(vol_b, crop, center_crop=False) logger.debug('Sub-volume A size: ' + str(subvol_a.shape)) logger.debug('Sub-volume B size: ' + str(subvol_b.shape)) level_iters = [40000, 10000, 1000, 500] sigmas = [7.0, 3.0, 1.0, 0.0] factors = [4, 2, 1, 1] if metrack.is_tracked(MUTUAL_INFORMATION_METRIC) or metrack.is_tracked( MUTUAL_INFORMATION_GRADIENT_METRIC): def callback(value: float, gradient: float): metrack.append_metric(MUTUAL_INFORMATION_METRIC, (None, value)) metrack.append_metric(MUTUAL_INFORMATION_GRADIENT_METRIC, (None, gradient)) else: callback = None affreg = AffineRegistration(metric=MutualInformationMetric( 32, None if sampling_prop > 0.99 else sampling_prop, sampling_type='grid'), level_iters=level_iters, sigmas=sigmas, factors=factors) transform = transform_cls() params0 = None starting_affine = vol_b.world_transform affine = affreg.optimize(subvol_a, subvol_b, transform, params0, subvol_a.grid_to_world, subvol_b.grid_to_world, starting_affine=starting_affine) logger.debug('Registration transform: ' + str(transform)) vol_b = Volume(vol_b, world_transform=np.array(affine.affine)) return vol_a, vol_b
def test_mut(): a = Volume((np.random.rand(64, 64, 64) * 2000).astype(np.uint16), spacing=(0.3, 0.3, 0.7)) with pytest.raises(ValueError): a += (np.random.rand(64, 64, 64) * 2000).astype(np.uint16)
def test_volume_update(rand_vol_a: Volume): assert not rand_vol_a.inverted rand_vol_a = Volume(rand_vol_a, inverted=True) assert rand_vol_a.inverted rand_vol_a = Volume(rand_vol_a) assert rand_vol_a.inverted
psf = np.zeros((9, 9, 9)) for x in range(9): for y in range(9): for z in range(9): psf[x, y, z] = multivariate_normal.pdf((x, y, z), mean=(4, 4, 4), cov=np.array([[4, 0, 0], [0, 4, 0], [0, 0, 4]])) blurred = convolve(vol, psf) blurred += (np.random.poisson(lam=25, size=blurred.shape) - 10) / 250. blurred_vol = Volume(blurred, (1, 1, 1)) blurred_vol.save_tiff('blurred_thing') # estimate = data # for i in range(200): # estimate = estimate * blur(data/(blur(estimate)+1e-6)) # estimate = restoration.richardson_lucy(blurred, psf, iterations=120) # # estimate_vol = Volume(estimate, (1, 1, 1)) # estimate_vol.save_tiff('estimate') # import numpy as np # import matplotlib.pyplot as plt # # from scipy.signal import convolve2d as conv2
def load_tiff(path: str, series: int = 0, channel: int = 0, inverted: bool = False, flipped: Tuple[bool] = (False, False, False), pixel_size: float = None, step_size: float = None, is_ome=True) -> Union[Tuple[Volume, Volume], Tuple[Volume]]: import json with tifffile.TiffFile(path, is_ome=is_ome) as f: data = f.asarray(series=series) logger.debug(f'Data shape is {data.shape}') # TODO: Figure out how to properly handle the axis order if data.ndim == 3: data = np.moveaxis(data, [0, 1, 2], [2, 0, 1]) elif data.ndim == 4: data = np.moveaxis(data, [0, 1, 2, 3], [3, 0, 1, 2])[2 * channel:2 * channel + 2] else: raise ValueError(f'Invalid data shape: {data.shape}') # TODO: Automatically extract metadata if pixel_size is None: pixel_size = f.pages[0].tags['XResolution'].value[1] / f.pages[ 0].tags['XResolution'].value[0] if f.pages[0].tags['YResolution'].value[1] / f.pages[0].tags[ 'YResolution'].value[0] != pixel_size: raise ValueError( f'X and Y resolution differ in metadata ({pixel_size}x{f.pages[0].tags["YResolution"].value[1] / f.pages[0].tags["YResolution"].value[0]})' ) if f.pages[0].tags['ResolutionUnit'].value != 3: raise ValueError( f'Unsupported resolution unit: {f.pages[0].tags["ResolutionUnit"].value}' ) pixel_size *= 10000 if step_size is None: try: step_size = json.load(f.micromanager_metadata['Summary'] ['SPIMAcqSettings'])['stepSizeUm'] except (ValueError, TypeError): raise ValueError( 'No stage step size specified and metadata cannot be accessed' ) if data.ndim == 3: if flipped[0]: data = data[::-1, :, :] if flipped[1]: data = data[:, ::-1, :] if flipped[2]: data = data[:, :, ::-1] return Volume(data, spacing=(pixel_size, pixel_size, step_size), is_skewed=True, flipped=flipped, inverted=inverted) else: if flipped[0]: data[1] = data[1, ::-1, :, :] if flipped[1]: data[1] = data[1, :, ::-1, :] if flipped[2]: data[1] = data[1, :, :, ::-1] return (Volume(data[0], spacing=(pixel_size, pixel_size, step_size), is_skewed=True, inverted=inverted), Volume(data[1], spacing=(pixel_size, pixel_size, step_size), is_skewed=True, flipped=flipped, inverted=True))
def process(self, data: ProcessData): psf_a = dispim.deconvolve.extract_psf(data[0], self.min_size, self.max_size, self.psf_half_width) psf_b = dispim.deconvolve.extract_psf(data[1], self.min_size, self.max_size, self.psf_half_width) return Volume(data[0], psf=psf_a), Volume(data[1], psf=psf_b)