direct.nn.varsplitnet package

Submodules

direct.nn.varsplitnet.config module

class direct.nn.varsplitnet.config.MRIVarSplitNetConfig(model_name='???', engine_name=None, num_steps_reg=8, num_steps_dc=8, image_init='sense', no_parameter_sharing=True, kspace_no_parameter_sharing=True, image_model_architecture=ModelName.UNET, kspace_model_architecture=None, image_resnet_hidden_channels=128, image_resnet_num_blocks=15, image_resnet_batchnorm=True, image_resnet_scale=0.1, image_unet_num_filters=32, image_unet_num_pool_layers=4, image_unet_dropout=0.0, image_didn_hidden_channels=16, image_didn_num_dubs=6, image_didn_num_convs_recon=9, kspace_resnet_hidden_channels=64, kspace_resnet_num_blocks=1, kspace_resnet_batchnorm=True, kspace_resnet_scale=0.1, kspace_unet_num_filters=16, kspace_unet_num_pool_layers=4, kspace_unet_dropout=0.0, kspace_didn_hidden_channels=8, kspace_didn_num_dubs=6, kspace_didn_num_convs_recon=9, image_conv_hidden_channels=64, image_conv_n_convs=15, image_conv_activation=ActivationType.RELU, image_conv_batchnorm=False, kspace_conv_hidden_channels=64, kspace_conv_n_convs=15, kspace_conv_activation=ActivationType.PRELU, kspace_conv_batchnorm=False)

Bases: ModelConfig

num_steps_reg = 8

num_steps_dc = 8

image_init = 'sense'

no_parameter_sharing = True

kspace_no_parameter_sharing = True

image_model_architecture = 'unet'

kspace_model_architecture = None

image_resnet_hidden_channels = 128

image_resnet_num_blocks = 15

image_resnet_batchnorm = True

image_resnet_scale = 0.1

image_unet_num_filters = 32

image_unet_num_pool_layers = 4

image_unet_dropout = 0.0

image_didn_hidden_channels = 16

image_didn_num_dubs = 6

image_didn_num_convs_recon = 9

kspace_resnet_hidden_channels = 64

kspace_resnet_num_blocks = 1

kspace_resnet_batchnorm = True

kspace_resnet_scale = 0.1

kspace_unet_num_filters = 16

kspace_unet_num_pool_layers = 4

kspace_unet_dropout = 0.0

kspace_didn_hidden_channels = 8

kspace_didn_num_dubs = 6

kspace_didn_num_convs_recon = 9

image_conv_hidden_channels = 64

image_conv_n_convs = 15

image_conv_activation = 'relu'

image_conv_batchnorm = False

kspace_conv_hidden_channels = 64

kspace_conv_n_convs = 15

kspace_conv_activation = 'prelu'

kspace_conv_batchnorm = False

__init__(model_name='???', engine_name=None, num_steps_reg=8, num_steps_dc=8, image_init='sense', no_parameter_sharing=True, kspace_no_parameter_sharing=True, image_model_architecture=ModelName.UNET, kspace_model_architecture=None, image_resnet_hidden_channels=128, image_resnet_num_blocks=15, image_resnet_batchnorm=True, image_resnet_scale=0.1, image_unet_num_filters=32, image_unet_num_pool_layers=4, image_unet_dropout=0.0, image_didn_hidden_channels=16, image_didn_num_dubs=6, image_didn_num_convs_recon=9, kspace_resnet_hidden_channels=64, kspace_resnet_num_blocks=1, kspace_resnet_batchnorm=True, kspace_resnet_scale=0.1, kspace_unet_num_filters=16, kspace_unet_num_pool_layers=4, kspace_unet_dropout=0.0, kspace_didn_hidden_channels=8, kspace_didn_num_dubs=6, kspace_didn_num_convs_recon=9, image_conv_hidden_channels=64, image_conv_n_convs=15, image_conv_activation=ActivationType.RELU, image_conv_batchnorm=False, kspace_conv_hidden_channels=64, kspace_conv_n_convs=15, kspace_conv_activation=ActivationType.PRELU, kspace_conv_batchnorm=False)

direct.nn.varsplitnet.varsplitnet module

class direct.nn.varsplitnet.varsplitnet.MRIVarSplitNet(forward_operator, backward_operator, num_steps_reg, num_steps_dc, image_init=InitType.SENSE, no_parameter_sharing=True, image_model_architecture=ModelName.UNET, kspace_no_parameter_sharing=True, kspace_model_architecture=None, **kwargs)

Bases: Module

MRI reconstruction network that solves the variable split optimisation problem.

It solves the following:

\[ \begin{align}\begin{aligned}z^{i-1} = \arg \min_{z} \mu * ||x^{i-1} - z||_2^2 + \mathcal{R}(z)\\x^{i} = \arg \min_{x} ||y - A(x)||_2^2 + \mu * ||x - z^{i-1}||_2^2\end{aligned}\end{align} \]

by unrolling twice using the gradient descent algorithm and replacing \(R\) with a neural network. More specifically, for \(z_0, x_0 = \text{SENSE}(\tilde{y})\):

\[z^{i} = \alpha_{i-1} \times f_{\theta_{i-1}}\Big(\mu(z^{i-1} - x^{i-1}), z^{i-1}\Big), \quad i=1,\cdots,T_{reg}\]

where \(x^{i}\) is the output of

\[(x^{i})^{j} = (x^{i})^{j-1} - \beta_{j-1} \Big[ A^{*}\big( A( (x^{i})^{j-1} ) - \tilde{y} \big) + \mu ((x^{i})^{j-1} - z^{i}) \Big], \quad j=1,\cdots,T_{dc},\]

i.e. \(x^{i}=(x^{i}^{T_{reg}})\).

__init__(forward_operator, backward_operator, num_steps_reg, num_steps_dc, image_init=InitType.SENSE, no_parameter_sharing=True, image_model_architecture=ModelName.UNET, kspace_no_parameter_sharing=True, kspace_model_architecture=None, **kwargs)

Inits MRIVarSplitNet.

Parameters:

• forward_operator (Callable) – Forward Operator.

• backward_operator (Callable) – Backward Operator.

• num_steps_reg (int) – Number of regularization steps.

• num_steps_dc (int) – Number of data consistency steps.

• image_init (str) – Image initialization type. Default: InitType.SENSE.

• no_parameter_sharing (bool) – If False, parameters are shared across iterations. Default: True.

• image_model_architecture (ModelName) – Image model architecture. Default: ModelName.UNET.

• kspace_no_parameter_sharing (bool) – If False, kspace model parameters are shared. Default: True.

• kspace_model_architecture (Optional[ModelName]) – Kspace model architecture. Default: None.

• **kwargs – Additional keyword arguments.

forward(masked_kspace, sensitivity_map, sampling_mask, scaling_factor=None)

Computes forward pass of MRIVarSplitNet.

Parameters:

• masked_kspace (Tensor) – torch.Tensor

• shape (Sensitivity map of)

• sensitivity_map (Tensor) – torch.Tensor

• shape – None.

• sampling_mask (Tensor) – torch.Tensor

• scaling_factor (Optional[Tensor]) – torch.Tensor

Returns:

torch.Tensor Output image of shape (N, height, width, complex=2).

Return type:

image

compute_model_per_coil(model, data)

Performs forward pass of model per coil.

Parameters:

• model (Module) – nn.Module

• run. (Model to)

• data (Tensor) – torch.Tensor

• shape (Multi-coil data of)

Returns:

torch.Tensor Computed output per coil.

Return type:

output

direct.nn.varsplitnet.varsplitnet_engine module

class direct.nn.varsplitnet.varsplitnet_engine.MRIVarSplitNetEngine(cfg, model, device, forward_operator=None, backward_operator=None, mixed_precision=False, **models)

Bases: MRIModelEngine

MRIVarSplitNet Engine.

__init__(cfg, model, device, forward_operator=None, backward_operator=None, mixed_precision=False, **models)

Inits MRIVarSplitNetEngine.

Parameters:

• cfg (BaseConfig) – Configuration file.

• model (Module) – Model.

• device (str) – Device. Can be “cuda: {idx}” or “cpu”.

• forward_operator (Optional[Callable]) – The forward operator. Default: None.

• backward_operator (Optional[Callable]) – The backward operator. Default: None.

• mixed_precision (bool) – Use mixed precision. Default: False.

• **models (Module) – Additional models.

forward_function(data)

Forward function for MRIVarSplitNetEngine.

Parameters:

data (Dict[str, Any]) – Input data dictionary containing keys such as “masked_kspace”, “sampling_mask”, and “sensitivity_map”.

Return type:

Tuple[Tensor, None]

Returns:

Tuple containing output image and None for k-space.

Module contents