CNN -> Utils

CNN Builder

class tardis_em.cnn.utils.build_cnn.BasicCNN(model='CNN', in_channels=1, out_channels=1, sigmoid=True, num_conv_layer=5, conv_layer_scaler=64, conv_kernel=3, padding=1, pool_kernel=2, img_patch_size=64, layer_components='3gcl', dropout=None, num_group=8, prediction=False)

Basic CNN MODEL

Back-compatible with old CNN builder method. New functionality allows packaging the model in onnx format and rapidly re-use/deploy.

Parameters:

in_channels (int) – Number of input channels for the first convolution.
out_channels (int) – Number of output channels for the last deconvolution.
sigmoid (bool) – If True, use nn.Sigmoid or nn.Softmax if False. Use True if nn.BCELoss is used as a loss function for (two-class segmentation).
num_conv_layer (int) – Number of convolution and deconvolution steps. A number of input channels for convolution is calculated as a linear progression. E.g. [64, 128, 256, 512].
conv_layer_scaler (int) – Scaler for the output feature channels.
conv_kernel (int) – Kernel size for the convolution.
padding (int) – Padding size for convolution.
pool_kernel (int) – kernel size for max_pooling.
img_patch_size (int) – Image patch size used for calculation network structure.
layer_components (str) – Convolution module used for building network.
dropout (float, optional) – If float, the dropout layer is built with a given drop-out rate.
num_group (int) – Number of groups for nn.GroupNorm.
prediction (bool) – If True, prediction mode is on.

Returns if Training:

torch.tensor: 5D torch without final activation.

Returns if Prediction:

If sigmoid is True:: torch.tensor: 5D torch with final activation from nn.Sigmoid().
If sigmoid is False:: torch.tensor: 5D torch with final activation from nn.Softmax(dim=1).

update_patch_size(img_patch_size, sigmoid)

build_cnn_model(): Encoder

class tardis_em.cnn.utils.build_cnn.UNet(model='CNN', **kwargs)

2D/3D UNET MODEL

“3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation” <https://arxiv.org/pdf/1606.06650.pdf>.

forward(x: Tensor) → Tensor

Forward for an Unet model.

Args:
x (torch.Tensor): Input image features.

Returns:
torch.Tensor: Probability mask of predicted image.

class tardis_em.cnn.utils.build_cnn.ResUNet(model='RCNN', **kwargs)

2D/3D RESNET MODEL

modified of <10.1016/j.isprsjprs.2020.01.013>

forward(x: Tensor) → Tensor

Forward for ResNet model.

Args:
x (torch.Tensor): Input image features.

Returns:
torch.Tensor: Probability mask of predicted image.

class tardis_em.cnn.utils.build_cnn.UNet3Plus(in_channels=1, out_channels=1, sigmoid=True, num_conv_layer=5, conv_layer_scaler=64, conv_kernel=3, padding=1, pool_kernel=2, img_patch_size=64, layer_components='3gcl', dropout=None, num_group=8, prediction=False, classifies=False, decoder_features=None)

3D FULLY CONNECTED UNET

modified of <https://arxiv.org/abs/2004.08790>

encoder: UNet3Plus classifier

decoder: Final Layer

static dot_product(x: Tensor, x_cls: Tensor) → Tensor

Dot product for two tensors.

Parameters:

x (torch.Tensor) – Image tensor.
x_cls (torch.Tensor) – Classified image tensor.

Returns:

Dot product of two tensors.

Return type:

torch.Tensor

forward(x: Tensor)

Forward for Unet3Plus model.

Args:
x (torch.Tensor): Input image features.

Returns:
torch.Tensor: Probability mask of predicted image.

class tardis_em.cnn.utils.build_cnn.FNet(in_channels=1, out_channels=1, sigmoid=True, num_conv_layer=5, conv_layer_scaler=64, conv_kernel=3, padding=1, pool_kernel=2, dropout=None, img_patch_size=64, layer_components='3gcl', num_group=8, attn_features=False, prediction=False)

New Unet model combining Unet and Unet3Plus The model shares encoder path which is split for decoding patch Unet and Unet3Plus style. The final layers from each are summed and sigmoid

Parameters:

in_channels – Number of input channels for the first convolution.
out_channels – Number of output channels for the last deconvolution.
sigmoid – If True, use nn.Sigmoid or nn.Softmax if False. Use True if nn.BCELoss is used as a loss function for (two-class segmentation).
num_conv_layer – Number of convolution and deconvolution steps. A number of input channels for convolution is calculated as a linear progression. E.g. [64, 128, 256, 512].
conv_layer_scaler – Feature output of the first layer.
conv_kernel – Kernel size for the convolution.
padding – Padding size for convolution.
pool_kernel – kernel size for max_pooling.
img_patch_size – Image patch size used for calculation of network structure.
layer_components – Convolution module used for building network.
num_group – Number of groups for nn.GroupNorm.
prediction – If True, prediction mode is on.

update_patch_size(img_patch_size, prediction)

build_cnn_model(): Encoder

forward(x: Tensor)

Forward for FNet model.

Args:
x (torch.Tensor): Input image features.

Returns:
torch.Tensor: Probability mask of predicted image.

CNN Utils

tardis_em.cnn.utils.utils.number_of_features_per_level(channel_scaler: int, num_levels: int) → list

Compute list of output channels for CNN.

Features = channel_scaler * 2^k: where: - k is layer number

Parameters:

channel_scaler (int) – Number of initial input channels for CNN.
num_levels (int) – Number of channels from max_number_of_conv_layer().

Returns:

List of output channels.

Return type:

list

tardis_em.cnn.utils.utils.max_number_of_conv_layer(img=None, input_volume=64, max_out=8, kernel_size=3, padding=1, stride=1, pool_size=2, pool_stride=2, first_max_pool=False) → int

Calculate maximum possible number of layers given image size.

Based on the torch input automatically select number of convolution blocks, based on a standard settings.

I = [(W - K + 2*P) / S] + 1 Out_size = [(I - F) / S] + 1 - W is the input volume - K is the Kernel size - P is the padding - S is the stride - F is the max pooling kernel size

Parameters:

img (np.ndarray, Optional) – Tensor data from which size of is calculated.
input_volume (int) – Size of the multiplayer for the convolution.
max_out (int) – Maximal output dimension after max_pooling.
kernel_size (int) – Kernel size for the convolution.
padding (int) – Padding size for the convolution blocks.
stride – (int) Stride for the convolution blocks.
pool_size (int)
pool_stride (int) – Max Pooling stride size.
first_max_pool (bool) – If first CNN block has max pooling.

Returns:

Maximum number of CNN layers.

Return type:

int

tardis_em.cnn.utils.utils.normalize_image(image: ndarray) → ndarray

Simple image data normalizer between 0,1.

Parameters:: image – Image data set.
Returns:: Normalized image between 0 and 1 values.
Return type:: np.ndarray

tardis_em.cnn.utils.utils.check_model_dict(model_dict: dict) → dict

Check and rebuild model structure dictionary to ensure back-compatibility.

Parameters:: model_dict (dict) – Model structure dictionary.
Returns:: Standardize model structure dictionary.
Return type:: dict