Skip to content

Machine learning models

The qim3d library aims to ease the creation of ML models for volumetric images

qim3d.ml.models

qim3d.ml.models.UNet

Bases: Module

2D UNet model for QIM imaging.

This class represents a 2D UNet model designed for imaging segmentation tasks.

Parameters:

Name Type Description Default
size small or medium or large

Size of the UNet model. Must be one of 'small', 'medium', or 'large'. Defaults to 'medium'.

 'medium'
dropout float

Dropout rate between 0 and 1. Defaults to 0.

 0
kernel_size int

Convolution kernel size. Defaults to 3.

 3
up_kernel_size int

Up-convolution kernel size. Defaults to 3.

 3
activation str

Activation function. Defaults to 'PReLU'.

 'PReLU'
bias bool

Whether to include bias in convolutions. Defaults to True.

 True
adn_order str

ADN (Activation, Dropout, Normalization) ordering. Defaults to 'NDA'.

 'NDA'

Raises:

Type Description
ValueError

If size is not one of 'small', 'medium', or 'large'.
Source code in qim3d/ml/models/_unet.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
class UNet(nn.Module):
    """
    2D UNet model for QIM imaging.

    This class represents a 2D UNet model designed for imaging segmentation tasks.

    Args:
        size ('small' or 'medium' or 'large', optional): Size of the UNet model. Must be one of 'small', 'medium', or 'large'. Defaults to 'medium'.
        dropout (float, optional): Dropout rate between 0 and 1. Defaults to 0.
        kernel_size (int, optional): Convolution kernel size. Defaults to 3.
        up_kernel_size (int, optional): Up-convolution kernel size. Defaults to 3.
        activation (str, optional): Activation function. Defaults to 'PReLU'.
        bias (bool, optional): Whether to include bias in convolutions. Defaults to True.
        adn_order (str, optional): ADN (Activation, Dropout, Normalization) ordering. Defaults to 'NDA'.

    Raises:
        ValueError: If `size` is not one of 'small', 'medium', or 'large'.
    """

    def __init__(
        self,
        size="medium",
        dropout=0,
        kernel_size=3,
        up_kernel_size=3,
        activation="PReLU",
        bias=True,
        adn_order="NDA",
    ):
        super().__init__()
        if size not in ["small", "medium", "large"]:
            raise ValueError(
                f"Invalid model size: {size}. Size must be one of the following: 'small', 'medium', 'large'."
            )

        self.size = size
        self.dropout = dropout
        self.kernel_size = kernel_size
        self.up_kernel_size = up_kernel_size
        self.activation = activation
        self.bias = bias
        self.adn_order = adn_order

        self.model = self._model_choice()

    def _model_choice(self):
        from monai.networks.nets import UNet as monai_UNet

        if self.size == "small":
            self.channels = (64, 128, 256)
        elif self.size == "medium":
            self.channels = (64, 128, 256, 512, 1024)
        elif self.size == "large":
            self.channels = (64, 128, 256, 512, 1024, 2048)

        model = monai_UNet(
            spatial_dims=2,
            in_channels=1,  # TODO: check if image has 1 or multiple input channels
            out_channels=1,
            channels=self.channels,
            strides=(2,) * (len(self.channels) - 1),
            kernel_size=self.kernel_size,
            up_kernel_size=self.up_kernel_size,
            act=self.activation,
            dropout=self.dropout,
            bias=self.bias,
            adn_ordering=self.adn_order,
        )
        return model

    def forward(self, x):
        x = self.model(x)
        return x