model.py

"""
defines our our model based on 3D CNN architecture
"""
import torch
import torch.nn as nn

import config


class Simple3DCNN(nn.Module):
    def __init__(self, num_classes=1):
        super(Simple3DCNN, self).__init__()
        # Input shape: (batch_size, 1, 48, 48, 48) -> (N, C, D, H, W)

        self.conv_block1 = nn.Sequential(
            nn.Conv3d(1, 32, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.BatchNorm3d(32),
            nn.MaxPool3d(kernel_size=2, stride=2)  # Output: 32 x 24 x 24 x 24
        )

        self.conv_block2 = nn.Sequential(
            nn.Conv3d(32, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.BatchNorm3d(64),
            nn.MaxPool3d(kernel_size=2, stride=2)  # Output: 64 x 12 x 12 x 12
        )

        self.conv_block3 = nn.Sequential(
            nn.Conv3d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.BatchNorm3d(128),
            nn.MaxPool3d(kernel_size=2, stride=2)  # Output: 128 x 6 x 6 x 6
        )

        # Flatten the output for the fully connected layer
        self.flatten = nn.Flatten()

        # Calculate the input size for the linear layer
        # Patch size: 48 -> 24 -> 12 -> 6
        fc_input_size = 128 * \
            (config.PATCH_SIZE[0] // 8) * \
            (config.PATCH_SIZE[1] // 8) * (config.PATCH_SIZE[2] // 8)

        self.fc_block = nn.Sequential(
            nn.Linear(fc_input_size, 256),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes)
        )

    def forward(self, x):
        x = self.conv_block1(x)
        x = self.conv_block2(x)
        x = self.conv_block3(x)
        x = self.flatten(x)
        x = self.fc_block(x)
        return x


if __name__ == '__main__':
    model = Simple3DCNN().to(config.DEVICE)
    dummy_input = torch.randn(config.BATCH_SIZE, 1,
                              *config.PATCH_SIZE).to(config.DEVICE)
    output = model(dummy_input)
    print(f"Model Architecture:\n{model}")
    print(f"Input shape: {dummy_input.shape}")
    print(f"Output shape: {output.shape}")