def test_pad_image_to_fft(self):
input_image = image.Image(np.zeros((7, 25)))
padded_image = image.pad_image_to_fft(input_image)
self.assertEqual(padded_image.shape, (8, 27))
input_image = image.Image(np.zeros((30, 27)))
padded_image = image.pad_image_to_fft(input_image)
self.assertEqual(padded_image.shape, (32, 27))
input_image = image.Image(np.zeros((300, 400)))
padded_image = image.pad_image_to_fft(input_image)
self.assertEqual(padded_image.shape, (324, 432))
def get(self, illuminated_volume, limits, fields, **kwargs):
''' Convolves the image with a pupil function
'''
# Pad volume
padded_volume, limits = self._pad_volume(illuminated_volume,
limits=limits,
**kwargs)
# Extract indexes of the output region
pad = kwargs.get("padding", (0, 0, 0, 0))
output_region = np.array(
kwargs.get("upscaled_output_region", (None, None, None, None)))
output_region[0] = None if output_region[0] is None else int(
output_region[0] - limits[0, 0] - pad[0])
output_region[1] = None if output_region[1] is None else int(
output_region[1] - limits[1, 0] - pad[1])
output_region[2] = None if output_region[2] is None else int(
output_region[2] - limits[0, 0] + pad[2])
output_region[3] = None if output_region[3] is None else int(
output_region[3] - limits[1, 0] + pad[3])
padded_volume = padded_volume[output_region[0]:output_region[2],
output_region[1]:output_region[3], :]
z_limits = limits[2, :]
output_image = Image(np.zeros((*padded_volume.shape[0:2], 1)))
index_iterator = range(padded_volume.shape[2])
z_iterator = np.linspace(z_limits[0],
z_limits[1],
num=padded_volume.shape[2],
endpoint=False)
zero_plane = np.all(padded_volume == 0, axis=(0, 1), keepdims=False)
# z_values = z_iterator[~zero_plane]
volume = pad_image_to_fft(padded_volume, axes=(0, 1))
voxel_size = kwargs['voxel_size']
pupils = (self._pupil(
volume.shape[:2], defocus=[1], include_aberration=False, **kwargs)
+ self._pupil(volume.shape[:2],
defocus=[-z_limits[1]],
include_aberration=True,
**kwargs))
pupil_step = np.fft.fftshift(pupils[0])
if "illumination" in kwargs:
light_in = np.ones(volume.shape[:2], dtype=np.complex)
light_in = kwargs["illumination"].resolve(light_in, **kwargs)
light_in = np.fft.fft2(light_in)
else:
light_in = np.zeros(volume.shape[:2], dtype=np.complex)
light_in[0, 0] = light_in.size
K = 2 * np.pi / kwargs["wavelength"]
field_z = [_get_position(field, return_z=True)[-1] for field in fields]
field_offsets = [
field.get_property("offset_z", default=0) for field in fields
]
z = z_limits[1]
for i, z in zip(index_iterator, z_iterator):
light_in = light_in * pupil_step
to_remove = []
for idx, fz in enumerate(field_z):
if fz < z:
propagation_matrix = self._pupil(
fields[idx].shape,
defocus=[z - fz - field_offsets[idx] / voxel_size[-1]],
include_aberration=False,
**kwargs)[0]
propagation_matrix = propagation_matrix * np.exp(
1j * voxel_size[-1] * 2 * np.pi / kwargs["wavelength"]
* kwargs["refractive_index_medium"] * (z - fz))
light_in += np.fft.fft2(
fields[idx][:, :,
0]) * np.fft.fftshift(propagation_matrix)
to_remove.append(idx)
for idx in reversed(to_remove):
fields.pop(idx)
field_z.pop(idx)
field_offsets.pop(idx)
if zero_plane[i]:
continue
ri_slice = volume[:, :, i]
light = np.fft.ifft2(light_in)
light_out = light * np.exp(1j * ri_slice * voxel_size[-1] * K)
light_in = np.fft.fft2(light_out)
# Add remaining fields
for idx, fz in enumerate(field_z):
prop_dist = z - fz - field_offsets[idx] / voxel_size[-1]
propagation_matrix = self._pupil(fields[idx].shape,
defocus=[prop_dist],
include_aberration=False,
**kwargs)[0]
propagation_matrix = propagation_matrix * np.exp(
-1j * voxel_size[-1] * 2 * np.pi / kwargs["wavelength"] *
kwargs["refractive_index_medium"] * prop_dist)
light_in += np.fft.fft2(
fields[idx][:, :, 0]) * np.fft.fftshift(propagation_matrix)
light_in_focus = light_in * np.fft.fftshift(pupils[-1])
output_image = np.fft.ifft2(
light_in_focus)[:padded_volume.shape[0], :padded_volume.shape[1]]
output_image = np.expand_dims(output_image, axis=-1)
output_image = Image(output_image[pad[0]:-pad[2], pad[1]:-pad[3]])
if not kwargs.get("return_field", False):
output_image = np.square(np.abs(output_image))
output_image.properties = illuminated_volume.properties
return output_image
def get(self, illuminated_volume, limits, **kwargs):
''' Convolves the image with a pupil function
'''
# Pad volume
padded_volume, limits = self._pad_volume(illuminated_volume,
limits=limits,
**kwargs)
# Extract indexes of the output region
pad = kwargs.get("padding", (0, 0, 0, 0))
output_region = np.array(
kwargs.get("upscaled_output_region", (None, None, None, None)))
output_region[0] = None if output_region[0] is None else int(
output_region[0] - limits[0, 0] - pad[0])
output_region[1] = None if output_region[1] is None else int(
output_region[1] - limits[1, 0] - pad[1])
output_region[2] = None if output_region[2] is None else int(
output_region[2] - limits[0, 0] + pad[2])
output_region[3] = None if output_region[3] is None else int(
output_region[3] - limits[1, 0] + pad[3])
padded_volume = padded_volume[output_region[0]:output_region[2],
output_region[1]:output_region[3], :]
z_limits = limits[2, :]
output_image = Image(np.zeros((*padded_volume.shape[0:2], 1)))
index_iterator = range(padded_volume.shape[2])
# Get planes in volume where not all values are 0.
z_iterator = np.linspace(z_limits[0],
z_limits[1],
num=padded_volume.shape[2],
endpoint=False)
zero_plane = np.all(padded_volume == 0, axis=(0, 1), keepdims=False)
z_values = z_iterator[~zero_plane]
# Further pad image to speed up fft
volume = pad_image_to_fft(padded_volume, axes=(0, 1))
pupils = self._pupil(volume.shape[:2], defocus=z_values, **kwargs)
pupil_iterator = iter(pupils)
# Loop through voluma and convole sample with pupil function
for i, z in zip(index_iterator, z_iterator):
if zero_plane[i]:
continue
image = volume[:, :, i]
pupil = Image(next(pupil_iterator))
psf = np.square(np.abs(np.fft.ifft2(np.fft.fftshift(pupil))))
optical_transfer_function = np.fft.fft2(psf)
fourier_field = np.fft.fft2(image)
convolved_fourier_field = fourier_field * optical_transfer_function
field = Image(np.fft.ifft2(convolved_fourier_field))
# Discard remaining imaginary part (should be 0 up to rounding error)
field = np.real(field)
output_image[:, :, 0] += field[:padded_volume.
shape[0], :padded_volume.shape[1]]
output_image = output_image[pad[0]:-pad[2], pad[1]:-pad[3]]
try:
output_image.properties = illuminated_volume.properties + pupil.properties
except UnboundLocalError:
output_image.properties = illuminated_volume.properties
return output_image
def get(self, illuminated_volume, limits, **kwargs):
''' Convolves the image with a pupil function
'''
# Pad volume
padded_volume, limits = self._pad_volume(illuminated_volume,
limits=limits,
**kwargs)
# Extract indexes of the output region
pad = kwargs.get("padding", (0, 0, 0, 0))
output_region = np.array(
kwargs.get("output_region", (None, None, None, None)))
output_region[0] = None if output_region[0] is None else int(
output_region[0] - limits[0, 0] - pad[0])
output_region[1] = None if output_region[1] is None else int(
output_region[1] - limits[1, 0] - pad[1])
output_region[2] = None if output_region[2] is None else int(
output_region[2] - limits[0, 0] + pad[2])
output_region[3] = None if output_region[3] is None else int(
output_region[3] - limits[1, 0] + pad[3])
padded_volume = padded_volume[output_region[0]:output_region[2],
output_region[1]:output_region[3], :]
z_limits = limits[2, :]
output_image = Image(np.zeros((*padded_volume.shape[0:2], 1)))
index_iterator = range(padded_volume.shape[2])
z_iterator = np.linspace(z_limits[0],
z_limits[1],
num=padded_volume.shape[2],
endpoint=False)
zero_plane = np.all(padded_volume == 0, axis=(0, 1), keepdims=False)
# z_values = z_iterator[~zero_plane]
volume = pad_image_to_fft(padded_volume, axes=(0, 1))
voxel_size = kwargs['voxel_size']
pupils = (self._pupil(
volume.shape[:2], defocus=[1], include_aberration=False, **kwargs)
+ self._pupil(volume.shape[:2],
defocus=[-z_limits[1]],
include_aberration=True,
**kwargs))
pupil_step = np.fft.fftshift(pupils[0])
if "illumination" in kwargs:
light_in = np.ones(volume.shape[:2])
light_in = kwargs["illumination"].resolve(light_in, **kwargs)
light_in = np.fft.fft2(light_in)
else:
light_in = np.zeros(volume.shape[:2])
light_in[0, 0] = light_in.size
K = 2 * np.pi / kwargs["wavelength"]
for i, z in zip(index_iterator, z_iterator):
light_in = light_in * pupil_step
if zero_plane[i]:
continue
ri_slice = volume[:, :, i]
light = np.fft.ifft2(light_in)
light_out = light * np.exp(1j * ri_slice * voxel_size[-1] * K)
light_in = np.fft.fft2(light_out)
light_in_focus = light_in * np.fft.fftshift(pupils[-1])
output_image = np.fft.ifft2(
light_in_focus)[:padded_volume.shape[0], :padded_volume.shape[1]]
output_image = np.expand_dims(output_image, axis=-1)
output_image = Image(output_image[pad[0]:-pad[2], pad[1]:-pad[3]])
if not kwargs.get("return_field", False):
output_image = np.square(np.abs(output_image))
output_image.properties = illuminated_volume.properties
return output_image