unet

Module containing UNet models. Most code from https://github.com/milesial/Pytorch-UNet

class DoubleConv(in_channels: int, out_channels: int, mid_channels: int | None = None)

Applies a double convolution to the input (convolution => [BN] => ReLU) * 2

Adapted from https://github.com/milesial/Pytorch-UNet for minerva.

double_conv

Double convolutions.

Type:: Module

Parameters:

in_channels (int) – Number of input channels.
out_channels (int) – Number of output channels.
mid_channels (int) – Optional; Intermediate number of channels between convolutions. If None, set to out_channels.

forward(x: Tensor) → Tensor

Applies the double convolutions to the input.

Parameters:: x (Tensor) – Input tensor.
Returns:: Input passed through the double convolutions.
Return type:: Tensor

class Down(in_channels: int, out_channels: int)

Downscaling with maxpool then double convolution.

Adapted from https://github.com/milesial/Pytorch-UNet for minerva.

maxpool_conv

Sequential of MaxPool2d then DoubleConv.

Type:: Module

Parameters:

in_channels (int) – Number of input channels.
out_channels (int) – Number of output channels.

forward(x: Tensor) → Tensor

Applies a maxpool then double convolution to the input.

Parameters:: x (Tensor) – Input tensor.
Returns:: Input tensor passed through maxpooling then double convolutions.
Return type:: Tensor

class OutConv(in_channels: int, out_channels: int)

1x1 convolution to change the number of channels down of the input.

Adapted from https://github.com/milesial/Pytorch-UNet for minerva.

conv

1x1 convolutional layer.

Type:: Module

Parameters:

in_channels (int) – Number of input channels.
out_channels (int) – Number of output channels.

forward(x: Tensor) → Tensor

Passes the input tensor through the 1x1 convolutional layer.

Parameters:: x (Tensor) – Input tensor.
Returns:: Input passed reduced from in_channels to out_channels.
Return type:: Tensor

class UNet(criterion: Any, input_size: tuple[int, ...] = (4, 256, 256), n_classes: int = 8, bilinear: bool = False, scaler: GradScaler | None = None)

UNet model. Good for segmentation problems.

Adapted from https://github.com/milesial/Pytorch-UNet for use in minerva.

bilinear

Type:: bool

inc

Double convolutional layers as input to the network to 64 channels.

Type:: DoubleConv

down1

Downscale then double convolution from 64 channels to 128.

Type:: Down

down2

Downscale then double convolution from 128 channels to 256.

Type:: Down

down3

Downscale then double convolution from 256 channels to 512.

Type:: Down

down4

Downscale then double convolution from 512 channels to 1024 and the latent space.

Type:: Down

up1

First upsample then concatenated input double de-convolutional layer.

Type:: Up

up2

Second upsample then concatenated input double de-convolutional layer.

Type:: Up

up3

Third upsample then concatenated input double de-convolutional layer.

Type:: Up

up4

Fourth upsample then concatenated input double de-convolutional layer.

Type:: Up

outc

1x1 output convolutional layer.

Type:: OutConv

Parameters:

criterion – torch loss function model will use.
input_size (tuple[int, ...]) – Optional; Defines the shape of the input data in order of number of channels, image width, image height.
n_classes (int) – Optional; Number of classes in data to be classified.
bilinear (bool) – Optional;

forward(x: Tensor) → Tensor

Performs a forward pass of the UNet.

Adapted from https://github.com/milesial/Pytorch-UNet for minerva.

Parameters:: x (Tensor) – Input tensor to the UNet.
Returns:: Output from the UNet.
Return type:: Tensor

class UNetR(criterion: Any, input_size: tuple[int, ...] = (4, 256, 256), n_classes: int = 8, bilinear: bool = False, scaler: GradScaler | None = None, backbone_weight_path: str | None = None, freeze_backbone: bool = False, backbone_kwargs: dict[str, Any] = {})

UNet model which incorporates a ResNet as the encoder.

backbone_name

Name of the backbone class.

Type:: str

backbone

Backbone of the FCN that takes the imagery input and extracts learned representations.

Type:: Module

up1

First upsample then concatenated input double de-convolutional layer.

Type:: Up

up2

Second upsample then concatenated input double de-convolutional layer.

Type:: Up

up3

Third upsample then concatenated input double de-convolutional layer.

Type:: Up

upsample1

First upsample from output of up3.

Type:: Module

upsample2

Second upsample from output of up3 to match input spatial size.

Type:: Module

outc

1x1 output convolutional layer.

Type:: OutConv

Parameters:

criterion – torch loss function model will use.
input_size (tuple[int, ...]) – Optional; Defines the shape of the input data in order of number of channels, image width, image height.
n_classes (int) – Optional; Number of classes in data to be classified.
bilinear (bool) – Optional;
backbone_name (str) – Optional; Name of the backbone within this module to use for the UNet.
backbone_weight_path (str) – Optional; Path to pre-trained weights for the backbone to be loaded.
freeze_backbone (bool) – Freezes the weights on the backbone to prevent end-to-end training if using a pre-trained backbone.
backbone_kwargs (dict[str, Any]) – Optional; Keyword arguments for the backbone packed up into a dict.

forward(x: Tensor) → Tensor

Performs a forward pass of the UNet using backbone.

Passes the input tensor to the backbone. Then passes the output tensors from the resnet residual blocks to corresponding stages of the decoder, upsampling to create an output with the same spatial size as the input size.

Parameters:: x (Tensor) – Input tensor to the UNetR.
Returns:: Output from the UNetR.
Return type:: Tensor

class UNetR101(criterion: Any, input_size: tuple[int, ...] = (4, 256, 256), n_classes: int = 8, bilinear: bool = False, scaler: GradScaler | None = None, backbone_weight_path: str | None = None, freeze_backbone: bool = False, backbone_kwargs: dict[str, Any] = {}): UNet with a ResNet101 as the backbone.

class UNetR152(criterion: Any, input_size: tuple[int, ...] = (4, 256, 256), n_classes: int = 8, bilinear: bool = False, scaler: GradScaler | None = None, backbone_weight_path: str | None = None, freeze_backbone: bool = False, backbone_kwargs: dict[str, Any] = {}): UNet with a ResNet152 as the backbone.

class UNetR18(criterion: Any, input_size: tuple[int, ...] = (4, 256, 256), n_classes: int = 8, bilinear: bool = False, scaler: GradScaler | None = None, backbone_weight_path: str | None = None, freeze_backbone: bool = False, backbone_kwargs: dict[str, Any] = {}): UNet with a ResNet18 as the backbone.

class UNetR34(criterion: Any, input_size: tuple[int, ...] = (4, 256, 256), n_classes: int = 8, bilinear: bool = False, scaler: GradScaler | None = None, backbone_weight_path: str | None = None, freeze_backbone: bool = False, backbone_kwargs: dict[str, Any] = {}): UNet with a ResNet34 as the backbone.

class UNetR50(criterion: Any, input_size: tuple[int, ...] = (4, 256, 256), n_classes: int = 8, bilinear: bool = False, scaler: GradScaler | None = None, backbone_weight_path: str | None = None, freeze_backbone: bool = False, backbone_kwargs: dict[str, Any] = {}): UNet with a ResNet50 as the backbone.

class Up(in_channels: int, out_channels: int, bilinear: bool = True)

Upscaling then double convolution.

Adapted from https://github.com/milesial/Pytorch-UNet for use in minerva.

up

Upsampling if bilinear==True, else transpose convolutional layer.

Type:: Module

conv

Double convolutional layers.

Type:: DoubleConv

Parameters:

in_channels (int) – Number of input channels.
out_channels (int) – Number of output channels.
bilinear (bool) – Optional;

forward(x1: Tensor, x2: Tensor) → Tensor

Applies upscaling to x1, concats x1 with x2 then applies DoubleConv to the result.

Parameters:

x1 (Tensor) – Input tensor 1 to be upscaled to match x2.
x2 (Tensor) – Input tensor 2 to be concated with upscaled x1 and passed through DoubleConv.

Returns:

Output tensor of the the upscaling and double convolutions.

Return type:

Tensor