def tf_ortho_ifft2d(kspace, enable_multiprocessing=True):
axes = [len(kspace.shape) - 2, len(kspace.shape) - 1]
scaling_norm = tf.cast(
tf.math.sqrt(
tf.cast(tf.math.reduce_prod(tf.shape(kspace)[-2:]), 'float32')),
kspace.dtype)
if len(kspace.shape) == 4:
# multicoil case
ncoils = tf.shape(kspace)[1]
n_slices = tf.shape(kspace)[0]
k_shape_x = tf.shape(kspace)[-2]
k_shape_y = tf.shape(kspace)[-1]
shifted_kspace = ifftshift(kspace, axes=axes)
if enable_multiprocessing:
batched_shifted_kspace = tf.reshape(shifted_kspace,
(-1, k_shape_x, k_shape_y))
batched_shifted_image = tf.map_fn(
ifft2d,
batched_shifted_kspace,
parallel_iterations=multiprocessing.cpu_count(),
)
if len(kspace.shape) == 4:
# multicoil case
image_shape = [n_slices, ncoils, k_shape_x, k_shape_y]
elif len(kspace.shape) == 3:
image_shape = [n_slices, k_shape_x, k_shape_y]
else:
image_shape = [k_shape_x, k_shape_y]
shifted_image = tf.reshape(batched_shifted_image, image_shape)
else:
shifted_image = ifft2d(shifted_kspace)
image = fftshift(shifted_image, axes=axes)
return scaling_norm * image
def operator_and_matrix(self,
shape_info,
dtype,
use_placeholder,
ensure_self_adjoint_and_pd=False):
shape = shape_info.shape
# For this test class, we are creating Hermitian spectrums.
# We also want the spectrum to have eigenvalues bounded away from zero.
#
# pre_spectrum is bounded away from zero.
pre_spectrum = linear_operator_test_util.random_uniform(
shape=self._shape_to_spectrum_shape(shape),
dtype=dtype,
minval=1.,
maxval=2.)
pre_spectrum_c = _to_complex(pre_spectrum)
# Real{IFFT[pre_spectrum]}
# = IFFT[EvenPartOf[pre_spectrum]]
# is the IFFT of something that is also bounded away from zero.
# Therefore, FFT[pre_h] would be a well-conditioned spectrum.
pre_h = fft_ops.ifft2d(pre_spectrum_c)
# A spectrum is Hermitian iff it is the DFT of a real convolution kernel.
# So we will make spectrum = FFT[h], for real valued h.
h = math_ops.real(pre_h)
h_c = _to_complex(h)
spectrum = fft_ops.fft2d(h_c)
lin_op_spectrum = spectrum
if use_placeholder:
lin_op_spectrum = array_ops.placeholder_with_default(spectrum, shape=None)
operator = linalg.LinearOperatorCirculant2D(
lin_op_spectrum,
is_positive_definite=True if ensure_self_adjoint_and_pd else None,
is_self_adjoint=True if ensure_self_adjoint_and_pd else None,
input_output_dtype=dtype)
self.assertEqual(
operator.parameters,
{
"input_output_dtype": dtype,
"is_non_singular": None,
"is_positive_definite": (
True if ensure_self_adjoint_and_pd else None),
"is_self_adjoint": (
True if ensure_self_adjoint_and_pd else None),
"is_square": True,
"name": "LinearOperatorCirculant2D",
"spectrum": lin_op_spectrum,
})
mat = self._spectrum_to_circulant_2d(spectrum, shape, dtype=dtype)
return operator, mat
def tf_unmasked_adj_op(x, idx=0):
axes = [len(x.shape) - 3, len(x.shape) - 2]
scaling_norm = tf.dtypes.cast(
tf.math.sqrt(
tf.dtypes.cast(tf.math.reduce_prod(tf.shape(x)[-3:-1]),
'float32')), x.dtype)
return scaling_norm * tf.expand_dims(fftshift(
ifft2d(ifftshift(x[..., idx], axes=axes)), axes=axes),
axis=-1)
def ortho_ifft2d(kspace):
kspace = _order_for_ft(kspace)
shift_axes = [2, 3]
scaling_norm = _compute_scaling_norm(kspace)
shifted_kspace = ifftshift(kspace, axes=shift_axes)
image_shifted = ifft2d(shifted_kspace)
image_unnormed = fftshift(image_shifted, axes=shift_axes)
image = image_unnormed * scaling_norm
image = _order_after_ft(image)
return image
def extract_smaps(kspace, low_freq_percentage=8, background_thresh=4e-6):
n_low_freq = tf.cast(tf.shape(kspace)[-2:] * low_freq_percentage / 100, tf.int32)
center_dimension = tf.cast(tf.shape(kspace)[-2:] / 2, tf.int32)
low_freq_lower_locations = center_dimension - tf.cast(n_low_freq / 2, tf.int32)
low_freq_upper_locations = center_dimension + tf.cast(n_low_freq / 2, tf.int32)
###
# NOTE: the following stands for in numpy:
# low_freq_mask = np.zeros_like(kspace)
# low_freq_mask[
# ...,
# low_freq_lower_locations[0]:low_freq_upper_locations[0],
# low_freq_lower_locations[1]:low_freq_upper_locations[1]
# ] = 1
x_range = tf.range(low_freq_lower_locations[0], low_freq_upper_locations[0])
y_range = tf.range(low_freq_lower_locations[1], low_freq_upper_locations[1])
X_range, Y_range = tf.meshgrid(x_range, y_range)
X_range = tf.reshape(X_range, (-1,))
Y_range = tf.reshape(Y_range, (-1,))
low_freq_mask_indices = tf.stack([X_range, Y_range], axis=-1)
# we have to transpose because only the first dimension can be indexed in
# scatter_nd
scatter_nd_perm = [2, 3, 0, 1]
low_freq_mask = tf.scatter_nd(
indices=low_freq_mask_indices,
updates=tf.ones([
tf.size(X_range),
tf.shape(kspace)[0],
tf.shape(kspace)[1]],
),
shape=[tf.shape(kspace)[i] for i in scatter_nd_perm],
)
low_freq_mask = tf.transpose(low_freq_mask, perm=scatter_nd_perm)
###
low_freq_kspace = kspace * tf.cast(low_freq_mask, kspace.dtype)
shifted_kspace = ifftshift(low_freq_kspace, axes=[2, 3])
coil_image_low_freq_shifted = ifft2d(shifted_kspace)
coil_image_low_freq = fftshift(coil_image_low_freq_shifted, axes=[2, 3])
# no need to norm this since they all have the same norm
low_freq_rss = tf.norm(coil_image_low_freq, axis=1)
coil_smap = coil_image_low_freq / low_freq_rss[:, None]
# for now we do not perform background removal based on low_freq_rss
# could be done with 1D k-means or fixed background_thresh, with tf.where
return coil_smap
def operator_and_matrix(
self, shape_info, dtype, use_placeholder,
ensure_self_adjoint_and_pd=False):
shape = shape_info.shape
# For this test class, we are creating Hermitian spectrums.
# We also want the spectrum to have eigenvalues bounded away from zero.
#
# pre_spectrum is bounded away from zero.
pre_spectrum = linear_operator_test_util.random_uniform(
shape=self._shape_to_spectrum_shape(shape),
dtype=dtype,
minval=1.,
maxval=2.)
pre_spectrum_c = _to_complex(pre_spectrum)
# Real{IFFT[pre_spectrum]}
# = IFFT[EvenPartOf[pre_spectrum]]
# is the IFFT of something that is also bounded away from zero.
# Therefore, FFT[pre_h] would be a well-conditioned spectrum.
pre_h = fft_ops.ifft2d(pre_spectrum_c)
# A spectrum is Hermitian iff it is the DFT of a real convolution kernel.
# So we will make spectrum = FFT[h], for real valued h.
h = math_ops.real(pre_h)
h_c = _to_complex(h)
spectrum = fft_ops.fft2d(h_c)
lin_op_spectrum = spectrum
if use_placeholder:
lin_op_spectrum = array_ops.placeholder_with_default(spectrum, shape=None)
operator = linalg.LinearOperatorCirculant2D(
lin_op_spectrum,
is_positive_definite=True if ensure_self_adjoint_and_pd else None,
is_self_adjoint=True if ensure_self_adjoint_and_pd else None,
input_output_dtype=dtype)
mat = self._spectrum_to_circulant_2d(spectrum, shape, dtype=dtype)
return operator, mat
def adj_op(self, inputs):
if self.masked:
if self.multicoil:
kspace, mask, smaps = inputs
else:
kspace, mask = inputs
kspace = _mask_tf([kspace, mask])
else:
if self.multicoil:
kspace, smaps = inputs
else:
kspace = inputs
kspace = kspace[..., 0]
scaling_norm = _compute_scaling_norm(kspace)
shifted_kspace = ifftshift(kspace, axes=self.shift_axes)
image_shifted = ifft2d(shifted_kspace)
image_unnormed = fftshift(image_shifted, axes=self.shift_axes)
image = image_unnormed * scaling_norm
if self.multicoil:
image = tf.reduce_sum(image * tf.math.conj(smaps), axis=1)
image = image[..., None]
return image
def _spectrum_to_circulant_2d(self, spectrum, shape, dtype):
"""Creates a block circulant matrix from a spectrum.
Intentionally done in an explicit yet inefficient way. This provides a
cross check to the main code that uses fancy reshapes.
Args:
spectrum: Float or complex `Tensor`.
shape: Python list. Desired shape of returned matrix.
dtype: Type to cast the returned matrix to.
Returns:
Block circulant (batch) matrix of desired `dtype`.
"""
spectrum = _to_complex(spectrum)
spectrum_shape = self._shape_to_spectrum_shape(shape)
domain_dimension = spectrum_shape[-1]
if not domain_dimension:
return array_ops.zeros(shape, dtype)
block_shape = spectrum_shape[-2:]
# Explicitly compute the action of spectrum on basis vectors.
matrix_rows = []
for n0 in range(block_shape[0]):
for n1 in range(block_shape[1]):
x = np.zeros(block_shape)
# x is a basis vector.
x[n0, n1] = 1.0
fft_x = fft_ops.fft2d(math_ops.cast(x, spectrum.dtype))
h_convolve_x = fft_ops.ifft2d(spectrum * fft_x)
# We want the flat version of the action of the operator on a basis
# vector, not the block version.
h_convolve_x = array_ops.reshape(h_convolve_x, shape[:-1])
matrix_rows.append(h_convolve_x)
matrix = array_ops.stack(matrix_rows, axis=-1)
return math_ops.cast(matrix, dtype)
def _spectrum_to_circulant_2d(self, spectrum, shape, dtype):
"""Creates a block circulant matrix from a spectrum.
Intentionally done in an explicit yet inefficient way. This provides a
cross check to the main code that uses fancy reshapes.
Args:
spectrum: Float or complex `Tensor`.
shape: Python list. Desired shape of returned matrix.
dtype: Type to cast the returned matrix to.
Returns:
Block circulant (batch) matrix of desired `dtype`.
"""
spectrum = _to_complex(spectrum)
spectrum_shape = self._shape_to_spectrum_shape(shape)
domain_dimension = spectrum_shape[-1]
if not domain_dimension:
return array_ops.zeros(shape, dtype)
block_shape = spectrum_shape[-2:]
# Explicitly compute the action of spectrum on basis vectors.
matrix_rows = []
for n0 in range(block_shape[0]):
for n1 in range(block_shape[1]):
x = np.zeros(block_shape)
# x is a basis vector.
x[n0, n1] = 1.0
fft_x = fft_ops.fft2d(x.astype(np.complex64))
h_convolve_x = fft_ops.ifft2d(spectrum * fft_x)
# We want the flat version of the action of the operator on a basis
# vector, not the block version.
h_convolve_x = array_ops.reshape(h_convolve_x, shape[:-1])
matrix_rows.append(h_convolve_x)
matrix = array_ops.stack(matrix_rows, axis=-1)
return math_ops.cast(matrix, dtype)
def tf_unmasked_adj_op(x, idx=0):
scaling_norm = tf.dtypes.cast(tf.math.sqrt(tf.dtypes.cast(tf.math.reduce_prod(tf.shape(x)[1:3]), 'float32')), x.dtype)
return scaling_norm * tf.expand_dims(fftshift(ifft2d(ifftshift(x[..., idx], axes=[1, 2])), axes=[1, 2]), axis=-1)
def extract_smaps(kspace, low_freq_percentage=8, background_thresh=4e-6):
"""Extract raw sensitivity maps for kspaces
This function will first select a low frequency region in all the kspaces,
then Fourier invert it, and finally perform a normalisation by the root
sum-of-square.
kspace has to be of shape: nslices x ncoils x height x width
Arguments:
kspace (tf.Tensor): the kspace whose sensitivity maps you want extracted.
low_freq_percentage (int): the low frequency region to consider for
sensitivity maps extraction, given as a percentage of the width of
the kspace. In fastMRI, it's 8 for an acceleration factor of 4, and
4 for an acceleration factor of 8. Defaults to 8.
background_thresh (float): unused for now, will later allow to have
thresholded sensitivity maps.
Returns:
tf.Tensor: extracted raw sensitivity maps.
"""
n_low_freq = tf.cast(tf.shape(kspace)[-2:] * low_freq_percentage / 100, tf.int32)
center_dimension = tf.cast(tf.shape(kspace)[-2:] / 2, tf.int32)
low_freq_lower_locations = center_dimension - tf.cast(n_low_freq / 2, tf.int32)
low_freq_upper_locations = center_dimension + tf.cast(n_low_freq / 2, tf.int32)
###
# NOTE: the following stands for in numpy:
# low_freq_mask = np.zeros_like(kspace)
# low_freq_mask[
# ...,
# low_freq_lower_locations[0]:low_freq_upper_locations[0],
# low_freq_lower_locations[1]:low_freq_upper_locations[1]
# ] = 1
x_range = tf.range(low_freq_lower_locations[0], low_freq_upper_locations[0])
y_range = tf.range(low_freq_lower_locations[1], low_freq_upper_locations[1])
X_range, Y_range = tf.meshgrid(x_range, y_range)
X_range = tf.reshape(X_range, (-1,))
Y_range = tf.reshape(Y_range, (-1,))
low_freq_mask_indices = tf.stack([X_range, Y_range], axis=-1)
# we have to transpose because only the first dimension can be indexed in
# scatter_nd
scatter_nd_perm = [2, 3, 0, 1]
low_freq_mask = tf.scatter_nd(
indices=low_freq_mask_indices,
updates=tf.ones([
tf.size(X_range),
tf.shape(kspace)[0],
tf.shape(kspace)[1]],
),
shape=[tf.shape(kspace)[i] for i in scatter_nd_perm],
)
low_freq_mask = tf.transpose(low_freq_mask, perm=scatter_nd_perm)
###
low_freq_kspace = kspace * tf.cast(low_freq_mask, kspace.dtype)
shifted_kspace = ifftshift(low_freq_kspace, axes=[2, 3])
coil_image_low_freq_shifted = ifft2d(shifted_kspace)
coil_image_low_freq = fftshift(coil_image_low_freq_shifted, axes=[2, 3])
# no need to norm this since they all have the same norm
low_freq_rss = tf.norm(coil_image_low_freq, axis=1)
coil_smap = coil_image_low_freq / low_freq_rss[:, None]
# for now we do not perform background removal based on low_freq_rss
# could be done with 1D k-means or fixed background_thresh, with tf.where
return coil_smap