"""Core reconstruction functions of LiberTEM-holo.
This includes both FFT-based and phase-shifting approaches.
"""
# Functions freq_array, aperture_function, estimate_sideband_position
# estimate_sideband_size are adopted from Hyperspy
# and are subject of following copyright:
#
# Copyright 2007-2016 The HyperSpy developers
#
# HyperSpy is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# HyperSpy is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with HyperSpy. If not, see <http://www.gnu.org/licenses/>.
#
# Copyright 2019 The LiberTEM developers
#
# LiberTEM is distributed under the terms of the GNU General
# Public License as published by the Free Software Foundation,
# version 3 of the License.
# see: https://github.com/LiberTEM/LiberTEM
from __future__ import annotations
import typing
import time
from sparseconverter import NUMPY, for_backend
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LogNorm
import logging
from libertem_holo.base.filters import phase_unwrap
from libertem_holo.base.utils import get_slice_fft, HoloParams
log = logging.getLogger(__name__)
XPType = typing.Any # Union[Module("numpy"), Module("cupy")]
[docs]def reconstruct_frame(
frame: np.ndarray,
sb_pos: tuple[float, float],
aperture: np.ndarray,
slice_fft: tuple[slice, slice],
*,
precision: bool = True,
xp: XPType = np,
) -> np.ndarray:
"""Reconstruct a single hologram.
Parameters
----------
frame
A numpy or cupy array containing the input hologram
sb_pos
The sideband position, for example as returned by
`estimate_sideband_position`
aperture
A numpy or cupy array containing the aperture to apply in fourier space
slice_fft
A slice to crop to the selected output shape in fourier space,
as returned by `get_slice_fft`.
precision
Defines precision of the reconstruction, True for complex128 for the
resulting complex wave, otherwise results will be complex64
xp
Pass in either the numpy or cupy module to select CPU or GPU processing
"""
frame = xp.array(frame)
if not precision:
frame = frame.astype(np.float32)
frame_size = frame.shape
fft_frame = xp.fft.fft2(frame) / np.prod(frame_size)
fft_frame = xp.roll(fft_frame, xp.array(sb_pos), axis=(0, 1))
fft_frame = xp.fft.fftshift(xp.fft.fftshift(fft_frame)[slice_fft])
fft_frame = fft_frame * aperture
return xp.fft.ifft2(fft_frame) * np.prod(frame_size)
[docs]def reconstruct_double_resolution(
frames: np.ndarray,
sb_pos: tuple[float, float],
aperture: np.ndarray,
slice_fft: tuple[slice, slice],
*,
precision: bool = True,
xp: XPType = np,
) -> np.ndarray:
"""Reconstruct a stack of phase shifted holography with double resolution method.
Parameters
----------
frames : array_like
Two holograms taken at a phase offset of π; shape (2, height, width)
sb_pos
The sideband position, for example as returned by
`estimate_sideband_position`
aperture
A numpy or cupy array containing the aperture to apply in fourier space
slice_fft
A slice to crop to the selected output shape in fourier space,
as returned by `get_slice_fft`.
precision
Defines precision of the reconstruction, True for complex128 for the
resulting complex wave, otherwise results will be complex64
xp
Pass in either the numpy or cupy module to select CPU or GPU processing
Returns
-------
wav : complex array
the reconstructed complex image
"""
image_double_resolution = frames[1] - frames[0]
return reconstruct_frame(
image_double_resolution,
sb_pos,
aperture,
slice_fft,
precision=precision,
xp=xp,
)
[docs]def estimate_omega(
image: np.ndarray,
sideband_position: tuple[float, float],
*,
flip: bool = False,
) -> tuple[float, float]:
"""Estimates the frequency carrier of the hologram.
Parameters
----------
image : array
Holographic data array
sideband_position : tuple
The sideband position (y, x), referred to the non-shifted FFT.
flip
??? TODO
Returns
-------
omega : tuple
frequency carrier in y and x axis
"""
width = image.shape[-1]
if sideband_position[1] >= width / 2:
omega = (-sideband_position[0], (width - sideband_position[1]))
elif sideband_position[1] < width / 2:
omega = (-sideband_position[0], sideband_position[1])
if flip:
omega = (-omega[0], -omega[1])
return omega
[docs]def reconstruct_direct_euler(
image: np.ndarray,
omega: tuple[float, float],
) -> np.ndarray:
"""Reconstruct a stack of phase shifted holograms.
This is using the direct reconstruction method by Ru et.al. 1994 (euler
form)
Parameters
----------
image : array
Holographic data array
omega: tuple
frequency carrier in y and x axis
Returns
-------
phase final : 2D array
the reconstructed phase image
"""
number_of_images = image.shape[1]
phase_initial_euler = np.zeros(number_of_images, dtype="complex128")
n = np.arange(number_of_images)
phase_initial = 2 * np.pi * n / number_of_images
phase_initial_euler = np.exp(1j * phase_initial)
c22 = 0
for i in range(number_of_images):
c22 = c22 + (image[i] * phase_initial_euler[i])
c22 = c22 / number_of_images
x = np.linspace(0, omega[1], image.shape[2], endpoint=False)
y = np.linspace(0, omega[0], image.shape[1], endpoint=False)
irow, icol = np.meshgrid(x, y, indexing="xy")
ramp_carrier = np.exp(1j * 2 * np.pi * (irow + icol))
c22 = c22 / ramp_carrier
return np.angle(c22)
[docs]def reconstruct_direct(
stack: np.ndarray,
omega: tuple[float, float],
) -> np.ndarray:
"""Reconstruct a stack of phase shifted holograms.
This is using the direct reconstruction method.
Parameters
----------
stack : array_like
Stack of holograms
number_of_images : int
The number of images inside the stack
omega: tuple
frequency carrier in y and x axis
Returns
-------
phase final : 2D array
the reconstructed phase image
"""
xspace = omega[1]
yspace = omega[0]
compfront = np.zeros((stack.shape[1], stack.shape[2]), dtype="complex64")
compcar = np.zeros((stack.shape[1], stack.shape[2]), dtype="complex64")
sin_value_sum = np.zeros((stack.shape[1], stack.shape[2]), dtype="complex64")
cos_value_sum = np.zeros((stack.shape[1], stack.shape[2]), dtype="complex64")
coscar = np.zeros((stack.shape[1], stack.shape[2]), dtype="complex64")
sincar = np.zeros((stack.shape[1], stack.shape[2]), dtype="complex64")
number_of_images = stack.shape[0]
cos_value_sum = np.array(0)
sin_value_sum = np.array(0)
n = np.arange(number_of_images)
initial_phase_change = 2 * np.pi * (n / number_of_images)
for i in range(number_of_images):
cos_value = stack[i] * np.cos(initial_phase_change[i])
sin_value = stack[i] * np.sin(initial_phase_change[i])
cos_value_sum = cos_value_sum + cos_value
sin_value_sum = sin_value_sum + sin_value
compfront = cos_value_sum + sin_value_sum * 1j
x = np.linspace(0, xspace, stack.shape[2], endpoint=False)
y = np.linspace(0, yspace, stack.shape[1], endpoint=False)
irow, icol = np.meshgrid(x, y, indexing="xy")
coscar = np.cos(2 * np.pi * (icol + irow)).astype("float32")
sincar = np.sin(2 * np.pi * (icol + irow)).astype("float32")
compcar = coscar + sincar * 1j
compfinal = compfront / compcar
return np.arctan2(np.imag(compfinal), np.real(compfinal))
[docs]def display_fft_image(
image: np.ndarray,
sb_position: tuple,
slice_fft: np.ndarray,
mask: np.ndarray = 1,
detail: bool = True,
) -> None:
"""Display an fft image.
This function helps to show the steps of the reconstruction and to define
the best length and width of line mask
Parameters
----------
image : array
the input image
sb_position : tuple
the sideband position (y, x), referred to the non-shifted FFT.
slice fft : array
contain minimum and maximum value in y and x axis.
mask : array
contain the aperture and line mask
detail : bolean
"""
fft_original_image = np.fft.fft2(image) / np.prod(image.shape)
fft_original_image1 = np.roll(fft_original_image, sb_position, axis=(0, 1))
fft_original_image2 = np.fft.fftshift(fft_original_image1)
fft_original_image3 = fft_original_image2[slice_fft]
fft_original_image4 = np.fft.fftshift(fft_original_image3)
fft_with_aperture = fft_original_image4 * mask
if detail:
plt.figure()
plt.imshow(np.abs(fft_original_image1), norm=LogNorm(), cmap="gray")
plt.figure()
plt.imshow(np.abs(fft_original_image2), norm=LogNorm(), cmap="gray")
plt.figure()
plt.imshow(np.abs(fft_original_image3), norm=LogNorm(), cmap="gray")
plt.figure()
plt.imshow(np.abs(fft_original_image4), norm=LogNorm(), cmap="gray")
figure, (ax1, ax2) = plt.subplots(1, 2)
ax1.imshow(np.abs(fft_original_image4), norm=LogNorm(), cmap="gray")
ax1.set_title("Without Aperture Mask")
ax2.imshow(np.log1p(np.abs(fft_with_aperture)), cmap="gray")
ax2.set_title("FFT with Aperture Mask")
else:
figure, (ax1, ax2) = plt.subplots(1, 2)
ax1.imshow(np.abs(fft_original_image4), norm=LogNorm(), cmap="gray")
ax1.set_title("Without Aperture Mask")
ax2.imshow(np.abs(fft_with_aperture), vmax=0.01, cmap="gray")
ax2.set_title("FFT with Aperture Mask")
[docs]def get_phase(
hologram: np.ndarray,
params: HoloParams,
xp: XPType = np,
) -> np.ndarray:
"""Reconstruct hologram using HoloParams and extract and unwrap phase."""
t0 = time.perf_counter()
slice_fft = get_slice_fft(params.out_shape, hologram.shape)
phase_amp = reconstruct_frame(
hologram,
sb_pos=params.sb_position,
aperture=params.aperture,
slice_fft=slice_fft,
xp=xp
)
# phase_unwrap is numpy-only:
phase = for_backend(np.angle(phase_amp), NUMPY)
t1 = time.perf_counter()
phase_unwrapped = phase_unwrap(phase)
t2 = time.perf_counter()
log.debug(
f"time: total={t2-t0:.3f}s, "
f"reconstruction={t1-t0:.3f}s, unwrapping={t2-t1:.3f}s"
)
return phase_unwrapped
[docs]def reconstruct_bf(
frame: np.ndarray,
aperture: np.ndarray,
slice_fft: tuple[slice, slice],
*,
xp=np,
) -> np.ndarray:
"""Reconstruct a brightfield image from a hologram.
Please use `libertem_holo.base.filter.central_line_filter` to
filter out fresnel fringes as appropriate.
"""
frame = xp.array(frame)
fft_frame = xp.fft.fft2(frame)
fft_frame = xp.fft.fftshift(xp.fft.fftshift(fft_frame)[slice_fft])
fft_frame = fft_frame * xp.array(aperture)
return xp.fft.ifft2(fft_frame)