mridc.collections.reconstruction.models.unet_base package

Submodules

mridc.collections.reconstruction.models.unet_base.unet_block module

class mridc.collections.reconstruction.models.unet_base.unet_block.ConvBlock(in_chans: int, out_chans: int, drop_prob: float)

Bases: Module

A Convolutional Block that consists of two convolution layers each followed by instance normalization, LeakyReLU activation and dropout.

forward(image: Tensor) → Tensor

Parameters

image (Input 4D tensor of shape (N, in_chans, H, W).) –

Return type

Output tensor of shape (N, out_chans, H, W).

training: bool

class mridc.collections.reconstruction.models.unet_base.unet_block.NormUnet(chans: int, num_pools: int, in_chans: int = 2, out_chans: int = 2, drop_prob: float = 0.0, padding_size: int = 15, normalize: bool = True, norm_groups: int = 2)

Bases: Module

Normalized U-Net model.

This is the same as a regular U-Net, but with normalization applied to the input before the U-Net. This keeps the values more numerically stable during training.

static chan_complex_to_last_dim(x: Tensor) → Tensor

Convert the last dimension of the input to complex.

static complex_to_chan_dim(x: Tensor) → Tensor

Convert the last dimension of the input to complex.

forward(x: Tensor) → Tensor

Forward pass of the network.

norm(x: Tensor) → Tuple[Tensor, Tensor, Tensor]

Normalize the input.

pad(x: Tensor) → Tuple[Tensor, Tuple[List[int], List[int], int, int]]

Pad the input with zeros to make it square.

training: bool

unnorm(x: Tensor, mean: Tensor, std: Tensor) → Tensor

Unnormalize the input.

static unpad(x: Tensor, h_pad: List[int], w_pad: List[int], h_mult: int, w_mult: int) → Tensor

Unpad the input.

class mridc.collections.reconstruction.models.unet_base.unet_block.TransposeConvBlock(in_chans: int, out_chans: int)

Bases: Module

A Transpose Convolutional Block that consists of one convolution transpose layers followed by instance normalization and LeakyReLU activation.

forward(image: Tensor) → Tensor

Parameters

image (Input 4D tensor of shape (N, in_chans, H, W).) –

Return type

Output tensor of shape (N, out_chans, H*2, W*2).

training: bool

class mridc.collections.reconstruction.models.unet_base.unet_block.Unet(in_chans: int, out_chans: int, chans: int = 32, num_pool_layers: int = 4, drop_prob: float = 0.0)

Bases: Module

PyTorch implementation of a U-Net model, as presented in 1.

References

1

15. Ronneberger, P. Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer-assisted intervention, pages 234–241. Springer, 2015.

forward(image: Tensor) → Tensor

Parameters

image (Input 4D tensor of shape (N, in_chans, H, W).) –

Return type

Output tensor of shape (N, out_chans, H, W).

training: bool

Module contents