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
- 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).
- decoder
Final Layer
- 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 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.
- encoder
Decoder
- decoder_3plus
Final Layer
- 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.scale_image(scale: tuple, image: ndarray | None = None, mask: ndarray | None = None, nn=False, device='cpu') Tuple[ndarray, ndarray, int] | Tuple[ndarray, int] | Tuple[None, int]
Scale image module using torch GPU interpolation
Expect 2D (XY/YX), 3D (ZYX)
- Parameters:
image (np.ndarray, Optional) – image data
mask (np.ndarray, Optional) – Optional binary mask image data
scale (tuple) – scale value for image
nn (bool)
device (str)
- tardis_em.cnn.utils.utils.nn_scaling(img: ndarray, scale: tuple, dtype: dtype, device='cpu') ndarray
- tardis_em.cnn.utils.utils.pil_LANCZOS(img: ndarray, scale: tuple, dtype: dtype) ndarray
Scale a 3D image by first down-sampling in the XY direction and then in the Z direction.
- Parameters:
img – image array.
scale – Scale array size.
dtype – Output dtype for scale array.
- Returns:
Up or Down scale 3D array.
- Return type:
no.ndarray
- tardis_em.cnn.utils.utils.linear_scaling(img: ndarray, scale: tuple, dtype: dtype, device='cpu') ndarray
Scaling of 2D/3D array using trilinear method from pytorch
- Parameters:
img – image array.
scale – Scale array size.
dtype – Output dtype for scale array.
- Returns:
Up or Down scale 3D array.
- Return type:
no.ndarray
- tardis_em.cnn.utils.utils.area_scaling(img: ndarray, scale: tuple, dtype: dtype, device='cpu') ndarray
Scaling of 3D array using area method from pytorch
- Parameters:
img – 3D array.
scale – Scale array size.
dtype – Output dtype for scale array.
- Returns:
Up or Down scale 3D array.
- Return type:
no.ndarray
- 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