biapy.models.simple_cnn

This module implements a simple Convolutional Neural Network (CNN) for image classification tasks. It is designed to be a straightforward and adaptable model for both 2D and 3D image inputs.

The simple_CNN class constructs a network composed of two main convolutional blocks, each followed by batch normalization, activation, pooling, and dropout. A final dense layer with Softmax activation is used for classification.

The architecture is flexible, automatically adapting to 2D or 3D input based on the provided image_shape.

Classes:

  • simple_CNN: The main class for creating the simple CNN model.

This module uses a helper function get_activation from biapy.models.blocks to dynamically select the activation function.

class biapy.models.simple_cnn.simple_CNN(image_shape, activation='ReLU', n_classes=2)[source]

Bases: Module

Create a simple Convolutional Neural Network (CNN) model.

This CNN architecture is designed for classification tasks and can handle both 2D and 3D image inputs. It consists of two main convolutional blocks followed by pooling and dropout, culminating in a fully connected layer for classification.

Parameters:

  • image_shape (Tuple[int, …]) – Dimensions of the input image.

    • For 2D: (height, width, channels)

    • For 3D: (depth, height, width, channels)

    The last element image_shape[-1] should be the number of input channels.

  • activation (str, optional) – Name of the activation layer to use within the convolutional blocks. Defaults to β€œReLU”.

  • n_classes (int, optional) – Number of output classes for the classification task. Defaults to 2.

Returns:

model – The constructed simple CNN model.

Return type:

nn.Module

forward(x) → Dict[source]

Perform the forward pass of the simple CNN model.

The input x passes sequentially through block1, block2, and then the last_block which flattens the features and applies a linear layer with Softmax for classification.

Parameters:

x (torch.Tensor) – The input image tensor. Expected shape for 2D: (batch_size, channels, height, width). Expected shape for 3D: (batch_size, channels, depth, height, width).

Returns:

A dictionary containing the classification probabilities. The key is typically β€˜out’ or similar, mapping to a torch.Tensor of shape (batch_size, n_classes).

Return type:

Dict