def process_intermediate_pred(self, pred, sensitivity_maps, target):
"""
Process the intermediate prediction.
Parameters
----------
pred: Intermediate prediction.
torch.Tensor, shape [batch_size, n_coils, n_x, n_y, 2]
sensitivity_maps: Coil sensitivity maps.
torch.Tensor, shape [batch_size, n_coils, n_x, n_y, 2]
target: Target data to crop to size.
torch.Tensor, shape [batch_size, n_x, n_y, 2]
Returns
-------
pred: torch.Tensor, shape [batch_size, n_x, n_y, 2]
Processed prediction.
"""
pred = ifft2(
pred, centered=self.fft_centered, normalization=self.fft_normalization, spatial_dims=self.spatial_dims
)
pred = coil_combination(pred, sensitivity_maps, method=self.coil_combination_method, dim=self.coil_dim)
pred = torch.view_as_complex(pred)
_, pred = center_crop_to_smallest(target, pred)
return pred
def forward(
self,
y: torch.Tensor,
sensitivity_maps: torch.Tensor,
mask: torch.Tensor,
init_pred: torch.Tensor,
target: torch.Tensor,
) -> torch.Tensor:
"""
Forward pass of the network.
Parameters
----------
y: Subsampled k-space data.
torch.Tensor, shape [batch_size, n_coils, n_x, n_y, 2]
sensitivity_maps: Coil sensitivity maps.
torch.Tensor, shape [batch_size, n_coils, n_x, n_y, 2]
mask: Sampling mask.
torch.Tensor, shape [1, 1, n_x, n_y, 1]
init_pred: Initial prediction.
torch.Tensor, shape [batch_size, n_x, n_y, 2]
target: Target data to compute the loss.
torch.Tensor, shape [batch_size, n_x, n_y, 2]
Returns
-------
pred: list of torch.Tensor, shape [batch_size, n_x, n_y, 2], or torch.Tensor, shape [batch_size, n_x, n_y, 2]
If self.accumulate_loss is True, returns a list of all intermediate estimates.
If False, returns the final estimate.
"""
estimation = y.clone()
for cascade in self.cascades:
# Forward pass through the cascades
estimation = cascade(estimation, y, sensitivity_maps, mask)
estimation = ifft2(
estimation,
centered=self.fft_centered,
normalization=self.fft_normalization,
spatial_dims=self.spatial_dims,
)
estimation = coil_combination(estimation,
sensitivity_maps,
method=self.coil_combination_method,
dim=self.coil_dim)
estimation = torch.view_as_complex(estimation)
_, estimation = center_crop_to_smallest(target, estimation)
return estimation
def forward(
self,
y: torch.Tensor,
sensitivity_maps: torch.Tensor,
mask: torch.Tensor,
init_pred: torch.Tensor,
target: torch.Tensor,
) -> torch.Tensor:
"""
Forward pass of the network.
Parameters
----------
y: Subsampled k-space data.
torch.Tensor, shape [batch_size, n_coils, n_x, n_y, 2]
sensitivity_maps: Coil sensitivity maps.
torch.Tensor, shape [batch_size, n_coils, n_x, n_y, 2]
mask: Sampling mask.
torch.Tensor, shape [1, 1, n_x, n_y, 1]
init_pred: Initial prediction.
torch.Tensor, shape [batch_size, n_x, n_y, 2]
target: Target data to compute the loss.
torch.Tensor, shape [batch_size, n_x, n_y, 2]
Returns
-------
pred: list of torch.Tensor, shape [batch_size, n_x, n_y, 2], or torch.Tensor, shape [batch_size, n_x, n_y, 2]
If self.accumulate_loss is True, returns a list of all intermediate estimates.
If False, returns the final estimate.
"""
sensitivity_maps = self.sens_net(
y, mask) if self.use_sens_net else sensitivity_maps
image = self.model(y, sensitivity_maps, mask)
image = torch.view_as_complex(
coil_combination(
ifft2(
image,
centered=self.fft_centered,
normalization=self.fft_normalization,
spatial_dims=self.spatial_dims,
),
sensitivity_maps,
method=self.coil_combination_method,
dim=self.coil_dim,
))
_, image = center_crop_to_smallest(target, image)
return image
def forward(
self,
y: torch.Tensor,
sensitivity_maps: torch.Tensor,
mask: torch.Tensor,
init_pred: torch.Tensor,
target: torch.Tensor,
) -> torch.Tensor:
"""
Forward pass of the network.
Parameters
----------
y: Subsampled k-space data.
torch.Tensor, shape [batch_size, n_coils, n_x, n_y, 2]
sensitivity_maps: Coil sensitivity maps.
torch.Tensor, shape [batch_size, n_coils, n_x, n_y, 2]
mask: Sampling mask.
torch.Tensor, shape [1, 1, n_x, n_y, 1]
init_pred: Initial prediction.
torch.Tensor, shape [batch_size, n_x, n_y, 2]
target: Target data to compute the loss.
torch.Tensor, shape [batch_size, n_x, n_y, 2]
Returns
-------
pred: list of torch.Tensor, shape [batch_size, n_x, n_y, 2], or torch.Tensor, shape [batch_size, n_x, n_y, 2]
If self.accumulate_loss is True, returns a list of all intermediate estimates.
If False, returns the final estimate.
"""
image = ifft2(y,
centered=self.fft_centered,
normalization=self.fft_normalization,
spatial_dims=self.spatial_dims)
if hasattr(self, "standardization"):
image = self.standardization(image, sensitivity_maps)
output_image = self._compute_model_per_coil(
self.unet, image.permute(0, 1, 4, 2, 3)).permute(0, 1, 3, 4, 2)
output_image = coil_combination(output_image,
sensitivity_maps,
method=self.coil_combination_method,
dim=self.coil_dim)
output_image = torch.view_as_complex(output_image)
_, output_image = center_crop_to_smallest(target, output_image)
return output_image
def forward(
self,
y: torch.Tensor,
sensitivity_maps: torch.Tensor,
mask: torch.Tensor,
target: torch.Tensor = None,
) -> Union[list, Any]:
"""
Forward pass of the zero-filled method.
Parameters
----------
y: Subsampled k-space data.
torch.Tensor, shape [batch_size, n_coils, n_x, n_y, 2]
sensitivity_maps: Coil sensitivity maps.
torch.Tensor, shape [batch_size, n_coils, n_x, n_y, 2]
mask: Sampling mask.
torch.Tensor, shape [1, 1, n_x, n_y, 1]
init_pred: Initial prediction.
torch.Tensor, shape [batch_size, n_x, n_y, 2]
target: Target data to compute the loss.
torch.Tensor, shape [batch_size, n_x, n_y, 2]
Returns
-------
pred: torch.Tensor, shape [batch_size, n_x, n_y, 2]
Predicted data.
"""
pred = coil_combination(
ifft2(y,
centered=self.fft_centered,
normalization=self.fft_normalization,
spatial_dims=self.spatial_dims),
sensitivity_maps,
method=self.coil_combination_method.upper(),
dim=self.coil_dim,
)
pred = check_stacked_complex(pred)
_, pred = center_crop_to_smallest(target, pred)
return pred
def forward(
self,
y: torch.Tensor,
sensitivity_maps: torch.Tensor,
mask: torch.Tensor,
init_pred: torch.Tensor,
target: torch.Tensor,
**kwargs,
) -> torch.Tensor:
"""
Forward pass of the network.
Parameters
----------
y: Subsampled k-space data.
torch.Tensor, shape [batch_size, n_coils, n_x, n_y, 2]
sensitivity_maps: Coil sensitivity maps.
torch.Tensor, shape [batch_size, n_coils, n_x, n_y, 2]
mask: Sampling mask.
torch.Tensor, shape [1, 1, n_x, n_y, 1]
init_pred: Initial prediction.
torch.Tensor, shape [batch_size, n_x, n_y, 2]
target: Target data to compute the loss.
torch.Tensor, shape [batch_size, n_x, n_y, 2]
Returns
-------
pred: list of torch.Tensor, shape [batch_size, n_x, n_y, 2], or torch.Tensor, shape [batch_size, n_x, n_y, 2]
If self.accumulate_loss is True, returns a list of all intermediate estimates.
If False, returns the final estimate.
"""
previous_state: Optional[torch.Tensor] = None
if self.initializer is not None:
if self.initializer_initialization == "sense":
initializer_input_image = (complex_mul(
ifft2(
y,
centered=self.fft_centered,
normalization=self.fft_normalization,
spatial_dims=self.spatial_dims,
),
complex_conj(sensitivity_maps),
).sum(self.coil_dim).unsqueeze(self.coil_dim))
elif self.initializer_initialization == "input_image":
if "initial_image" not in kwargs:
raise ValueError(
"`'initial_image` is required as input if initializer_initialization "
f"is {self.initializer_initialization}.")
initializer_input_image = kwargs["initial_image"].unsqueeze(
self.coil_dim)
elif self.initializer_initialization == "zero_filled":
initializer_input_image = ifft2(
y,
centered=self.fft_centered,
normalization=self.fft_normalization,
spatial_dims=self.spatial_dims,
)
previous_state = self.initializer(
fft2(
initializer_input_image,
centered=self.fft_centered,
normalization=self.fft_normalization,
spatial_dims=self.spatial_dims,
).sum(1).permute(0, 3, 1, 2))
kspace_prediction = y.clone()
for step in range(self.num_steps):
block = self.block_list[
step] if self.no_parameter_sharing else self.block_list[0]
kspace_prediction, previous_state = block(
kspace_prediction,
y,
mask,
sensitivity_maps,
previous_state,
)
eta = ifft2(
kspace_prediction,
centered=self.fft_centered,
normalization=self.fft_normalization,
spatial_dims=self.spatial_dims,
)
eta = coil_combination(eta,
sensitivity_maps,
method=self.coil_combination_method,
dim=self.coil_dim)
eta = torch.view_as_complex(eta)
_, eta = center_crop_to_smallest(target, eta)
return eta