Skip to content

Commit

Permalink
Add vsharp (2D & 3D), 3D transformers/engine, unet 3D, CMR challenge …
Browse files Browse the repository at this point in the history
…code (dataset, cartesuan subsampling), new metrics, transformer models
  • Loading branch information
georgeyiasemis committed Sep 20, 2023
1 parent 2a4c293 commit 5860c77
Show file tree
Hide file tree
Showing 57 changed files with 8,523 additions and 233 deletions.
273 changes: 273 additions & 0 deletions direct/common/subsample.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,9 @@

__all__ = (
"CalgaryCampinasMaskFunc",
"CartesianRandomMaskFunc",
"CartesianEquispacedMaskFunc",
"CartesianMagicMaskFunc",
"FastMRIRandomMaskFunc",
"FastMRIEquispacedMaskFunc",
"FastMRIMagicMaskFunc",
Expand Down Expand Up @@ -243,6 +246,84 @@ def mask_func(
return torch.from_numpy(self._reshape_and_broadcast_mask(shape, mask))


class CartesianRandomMaskFunc(FastMRIRandomMaskFunc):
r"""Cartesian random vertical line mask function.
Similar to :class:`FastMRIRandomMaskFunc, but instead of center fraction (`center_fractions`) representing
the fraction of center lines to the original size, here, it represents the actual number of center lines.
"""

def __init__(
self,
accelerations: Union[List[Number], Tuple[Number, ...]],
center_fractions: Optional[Union[List[int], Tuple[int, ...]]] = None,
uniform_range: bool = False,
):
"""Inits :class:`CartesianRandomMaskFunc`.
Parameters
----------
accelerations: Union[List[Number], Tuple[Number, ...]]
Amount of under-sampling_mask. An acceleration of 4 retains 25% of the k-space, the method is given by
mask_type. Has to be the same length as center_fractions if uniform_range is not True.
center_fractions: list or tuple of ints, optional
Number of low-frequency (center) columns to be retained.
If multiple values are provided, then one of these numbers is chosen uniformly each time.
uniform_range: bool
If True then an acceleration will be uniformly sampled between the two values. Default: True.
"""
super().__init__(
accelerations=accelerations,
center_fractions=center_fractions,
uniform_range=uniform_range,
)

def mask_func(
self,
shape: Union[List[int], Tuple[int, ...]],
return_acs: bool = False,
seed: Optional[Union[int, Iterable[int]]] = None,
) -> torch.Tensor:
"""Creates a random vertical Cartesian mask.
Parameters
----------
shape: list or tuple of ints
The shape of the mask to be created. The shape should at least 3 dimensions.
Samples are drawn along the second last dimension.
return_acs: bool
Return the autocalibration signal region as a mask.
seed: int or iterable of ints or None (optional)
Seed for the random number generator. Setting the seed ensures the same mask is generated
each time for the same shape. Default: None.
Returns
-------
mask: torch.Tensor
The sampling mask.
"""
if len(shape) < 3:
raise ValueError("Shape should have 3 or more dimensions")

with temp_seed(self.rng, seed):
num_cols = shape[-2]

num_center_lines, acceleration = self.choose_acceleration()
print(num_center_lines, acceleration)

mask = self.center_mask_func(num_cols, num_center_lines)

if return_acs:
return torch.from_numpy(self._reshape_and_broadcast_mask(shape, mask))

# Create the mask
prob = (num_cols / acceleration - num_center_lines) / (num_cols - num_center_lines)
mask = mask | (self.rng.uniform(size=num_cols) < prob)

return torch.from_numpy(self._reshape_and_broadcast_mask(shape, mask))


class FastMRIEquispacedMaskFunc(FastMRIMaskFunc):
r"""Equispaced vertical line mask function.
Expand Down Expand Up @@ -324,6 +405,90 @@ def mask_func(
return torch.from_numpy(self._reshape_and_broadcast_mask(shape, mask))


class CartesianEquispacedMaskFunc(FastMRIEquispacedMaskFunc):
r"""Cartesian equispaced vertical line mask function.
Similar to :class:`FastMRIEquispacedMaskFunc, but instead of center fraction (`center_fractions`) representing
the fraction of center lines to the original size, here, it represents the actual number of center lines.
"""

def __init__(
self,
accelerations: Union[List[Number], Tuple[Number, ...]],
center_fractions: Optional[Union[List[int], Tuple[int, ...]]] = None,
uniform_range: bool = False,
):
"""Inits :class:`CartesianEquispacedMaskFunc`.
Parameters
----------
accelerations: Union[List[Number], Tuple[Number, ...]]
Amount of under-sampling_mask. An acceleration of 4 retains 25% of the k-space, the method is given by
mask_type. Has to be the same length as center_fractions if uniform_range is not True.
center_fractions: list or tuple of ints, optional
Number of low-frequency (center) columns to be retained.
If multiple values are provided, then one of these numbers is chosen uniformly each time.
uniform_range: bool
If True then an acceleration will be uniformly sampled between the two values. Default: True.
"""
super().__init__(
accelerations=accelerations,
center_fractions=center_fractions,
uniform_range=uniform_range,
)

def mask_func(
self,
shape: Union[List[int], Tuple[int, ...]],
return_acs: bool = False,
seed: Optional[Union[int, Iterable[int]]] = None,
) -> torch.Tensor:
"""Creates an equispaced vertical Cartesian mask.
Parameters
----------
shape: list or tuple of ints
The shape of the mask to be created. The shape should at least 3 dimensions.
Samples are drawn along the second last dimension.
return_acs: bool
Return the autocalibration signal region as a mask.
seed: int or iterable of ints or None (optional)
Seed for the random number generator. Setting the seed ensures the same mask is generated
each time for the same shape. Default: None.
Returns
-------
mask: torch.Tensor
The sampling mask.
"""
if len(shape) < 3:
raise ValueError("Shape should have 3 or more dimensions")

with temp_seed(self.rng, seed):
num_cols = shape[-2]

num_center_lines, acceleration = self.choose_acceleration()

num_center_lines = int(num_center_lines)
mask = self.center_mask_func(num_cols, num_center_lines)

if return_acs:
return torch.from_numpy(self._reshape_and_broadcast_mask(shape, mask))

# determine acceleration rate by adjusting for the number of low frequencies
adjusted_accel = (acceleration * (num_center_lines - num_cols)) / (
num_center_lines * acceleration - num_cols
)
offset = self.rng.randint(0, round(adjusted_accel))

accel_samples = np.arange(offset, num_cols - 1, adjusted_accel)
accel_samples = np.around(accel_samples).astype(np.uint)
mask[accel_samples] = True

return torch.from_numpy(self._reshape_and_broadcast_mask(shape, mask))


class FastMRIMagicMaskFunc(FastMRIMaskFunc):
"""Vertical line mask function as implemented in [1]_.
Expand Down Expand Up @@ -422,6 +587,114 @@ def mask_func(
return torch.from_numpy(self._reshape_and_broadcast_mask(shape, mask))


class CartesianMagicMaskFunc(FastMRIMagicMaskFunc):
r"""Cartesian equispaced mask function as implemented in [1]_.
Similar to :class:`FastMRIMagicMaskFunc, but instead of center fraction (`center_fractions`) representing
the fraction of center lines to the original size, here, it represents the actual number of center lines.
References
----------
.. [1] Defazio, Aaron. “Offset Sampling Improves Deep Learning Based Accelerated MRI Reconstructions by
Exploiting Symmetry.” ArXiv:1912.01101 [Cs, Eess], Feb. 2020. arXiv.org, http://arxiv.org/abs/1912.01101.
"""

def __init__(
self,
accelerations: Union[List[Number], Tuple[Number, ...]],
center_fractions: Optional[Union[List[int], Tuple[int, ...]]] = None,
uniform_range: bool = False,
):
"""Inits :class:`CartesianMagicMaskFunc`.
Parameters
----------
accelerations: Union[List[Number], Tuple[Number, ...]]
Amount of under-sampling_mask. An acceleration of 4 retains 25% of the k-space, the method is given by
mask_type. Has to be the same length as center_fractions if uniform_range is not True.
center_fractions: list or tuple of ints, optional
Number of low-frequency (center) columns to be retained.
If multiple values are provided, then one of these numbers is chosen uniformly each time.
uniform_range: bool
If True then an acceleration will be uniformly sampled between the two values. Default: True.
"""
super().__init__(
accelerations=accelerations,
center_fractions=center_fractions,
uniform_range=uniform_range,
)

def mask_func(
self,
shape: Union[List[int], Tuple[int, ...]],
return_acs: bool = False,
seed: Optional[Union[int, Iterable[int]]] = None,
) -> torch.Tensor:
r"""Creates an equispaced Cartesian mask that exploits conjugate symmetry.
Parameters
----------
shape: list or tuple of ints
The shape of the mask to be created. The shape should at least 3 dimensions.
Samples are drawn along the second last dimension.
return_acs: bool
Return the autocalibration signal region as a mask.
seed: int or iterable of ints or None (optional)
Seed for the random number generator. Setting the seed ensures the same mask is generated
each time for the same shape. Default: None.
Returns
-------
mask: torch.Tensor
The sampling mask.
"""
if len(shape) < 3:
raise ValueError("Shape should have 3 or more dimensions")

with temp_seed(self.rng, seed):
num_cols = shape[-2]

num_center_lines, acceleration = self.choose_acceleration()

# bound the number of low frequencies between 1 and target columns
target_cols_to_sample = int(round(num_cols / acceleration))
num_center_lines = max(min(num_center_lines, target_cols_to_sample), 1)

acs_mask = self.center_mask_func(num_cols, num_center_lines)

if return_acs:
return torch.from_numpy(self._reshape_and_broadcast_mask(shape, acs_mask))

# adjust acceleration rate based on target acceleration.
adjusted_target_cols_to_sample = target_cols_to_sample - num_center_lines
adjusted_acceleration = 0
if adjusted_target_cols_to_sample > 0:
adjusted_acceleration = int(round(num_cols / adjusted_target_cols_to_sample))

offset = self.rng.randint(0, high=adjusted_acceleration)

if offset % 2 == 0:
offset_pos = offset + 1
offset_neg = offset + 2
else:
offset_pos = offset - 1 + 3
offset_neg = offset - 1 + 0

poslen = (num_cols + 1) // 2
neglen = num_cols - (num_cols + 1) // 2
mask_positive = np.zeros(poslen, dtype=bool)
mask_negative = np.zeros(neglen, dtype=bool)

mask_positive[offset_pos::adjusted_acceleration] = True
mask_negative[offset_neg::adjusted_acceleration] = True
mask_negative = np.flip(mask_negative)

mask = np.fft.fftshift(np.concatenate((mask_positive, mask_negative)))
mask = np.logical_or(mask, acs_mask)

return torch.from_numpy(self._reshape_and_broadcast_mask(shape, mask))


class CalgaryCampinasMaskFunc(BaseMaskFunc):
BASE_URL = "https://s3.aiforoncology.nl/direct-project/calgary_campinas_masks/"
MASK_MD5S = {
Expand Down
Loading

0 comments on commit 5860c77

Please sign in to comment.