unet

pytorch models wrapped for DSM/LabIA projects.

class mfai.pytorch.models.unet.CustomUNet(in_channels, out_channels, input_shape, settings=CustomUNetSettings(encoder_name='resnet18', encoder_depth=5, encoder_weights=True, autopad_enabled=False))

Bases: BaseModel, AutoPaddingModel

CustomUNet is a model that allows the user to define a specific configuration, from pretrained weights or not (from ResNet encoders).

Parameters:

• in_channels (int) – Number of input channels.

• out_channels (int) – Number of channels for output mask (or you can think of it as the number of classes).

• input_shape (tuple[int, ...]) – Shape of the input.

• settings (CustomUNetSettings) – A CustomUNetSettings instance.

features_last = False

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Tensor)

Return type:

Tensor

model_type = 2

num_spatial_dims: int = 2

onnx_supported = True

register: bool = True

property settings: CustomUNetSettings

Returns the settings instance used to configure for this model.

settings_kls

alias of CustomUNetSettings

supported_num_spatial_dims = (2,)

validate_input_shape(input_shape)

Given an input shape, verifies whether the inputs fit with the

calling model’s specifications.

Parameters:

input_shape (Size) – The shape of the input data, excluding any batch dimension and channel dimension. For example, for a batch of 2D tensors of shape [B,C,W,H], [W,H] should be passed. For 3D data instead of shape [B,C,W,H,D], instead, [W,H,D] should be passed.

Returns:

Returns a tuple where the first element is a boolean signaling whether the given input shape

already fits the model’s requirements. If that value is False, the second element contains the closest shape that fits the model, otherwise it will be None.

Return type:

tuple[bool, Size]

class mfai.pytorch.models.unet.CustomUNetSettings(encoder_name='resnet18', encoder_depth=5, encoder_weights=True, autopad_enabled=False)

Bases: object

Parameters:

• encoder_name (Literal['resnet18', 'resnet34', 'resnet50']) – Name of the encoder used for the UNet. Defaults to ‘resnet18’.

• encoder_depth (int) – Number of layers to use of the initial encoder. Defaults to 5.

• encoder_weights (bool) – If true, uses pretrained weights of the encoder. Defaults to True.

• autopad_enabled (bool) – If true, allows autopadding of input image. Defaults to False.

autopad_enabled: bool

encoder_depth: int

encoder_name: Literal['resnet18', 'resnet34', 'resnet50']

encoder_weights: bool

classmethod from_dict(kvs, *, infer_missing=False)

Return type:

TypeVar(A, bound= DataClassJsonMixin)

Parameters:

kvs (dict | list | str | int | float | bool | None)

classmethod from_json(s, *, parse_float=None, parse_int=None, parse_constant=None, infer_missing=False, **kw)

Return type:

TypeVar(A, bound= DataClassJsonMixin)

Parameters:

s (str | bytes | bytearray)

classmethod schema(*, infer_missing=False, only=None, exclude=(), many=False, context=None, load_only=(), dump_only=(), partial=False, unknown=None)

Return type:

SchemaF[TypeVar(A, bound= DataClassJsonMixin)]

Parameters:

• infer_missing (bool)

• many (bool)

• partial (bool)

to_dict(encode_json=False)

Return type:

Dict[str, Union[dict, list, str, int, float, bool, None]]

to_json(*, skipkeys=False, ensure_ascii=True, check_circular=True, allow_nan=True, indent=None, separators=None, default=None, sort_keys=False, **kw)

Return type:

str

Parameters:

• skipkeys (bool)

• ensure_ascii (bool)

• check_circular (bool)

• allow_nan (bool)

• indent (int | str | None)

• separators (Tuple[str, str] | None)

• default (Callable | None)

• sort_keys (bool)

class mfai.pytorch.models.unet.DoubleConv(in_channels, out_channels, name)

Bases: Module

Parameters:

• in_channels (int)

• out_channels (int)

• name (str)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Tensor)

Return type:

Tensor

class mfai.pytorch.models.unet.UNet(in_channels, out_channels, input_shape, settings=UNetSettings(init_features=64, autopad_enabled=False))

Bases: BaseModel, AutoPaddingModel

Returns a UNet architecture, with uninitialised weights, matching desired numbers of input and output channels.

Implementation from the original paper: https://arxiv.org/pdf/1505.04597.pdf.

Parameters:

• in_channels (int) – Number of input channels.

• out_channels (int) – Number of channels for output mask (or you can think of it as the number of classes)

• input_shape (tuple[int, int]) – Shape of the input.

• settings (UNetSettings) – A UNetSettings instance.

Returns:

A UNet instance.

Return type:

torch.nn.Module

features_last = False

forward(x)

Description of the architecture from the original paper (https://arxiv.org/pdf/1505.04597.pdf): The network architecture is illustrated in Figure 1. It consists of a contracting path (left side) and an expansive path (right side). The contracting path follows the typical architecture of a convolutional network. It consists of the repeated application of two 3x3 convolutions (unpadded convolutions), each followed by a rectified linear unit (ReLU) and a 2x2 max pooling operation with stride 2 for downsampling. At each downsampling step we double the number of feature channels. Every step in the expansive path consists of an upsampling of the feature map followed by a 2x2 convolution (“up-convolution”) that halves the number of feature channels, a concatenation with the correspondingly cropped feature map from the contracting path, and two 3x3 convolutions, each fol- lowed by a ReLU. The cropping is necessary due to the loss of border pixels in every convolution. At the final layer a 1x1 convolution is used to map each 64- component feature vector to the desired number of classes. In total the network has 23 convolutional layers. To allow a seamless tiling of the output segmentation map (see Figure 2), it is important to select the input tile size such that all 2x2 max-pooling operations are applied to a layer with an even x- and y-size.

Return type:

Tensor

Parameters:

x (Tensor)

model_type = 2

num_spatial_dims: int = 2

onnx_supported = True

register: bool = True

property settings: UNetSettings

Returns the settings instance used to configure for this model.

settings_kls

alias of UNetSettings

supported_num_spatial_dims = (2,)

validate_input_shape(input_shape)

Given an input shape, verifies whether the inputs fit with the

calling model’s specifications.

Parameters:

input_shape (Size) – The shape of the input data, excluding any batch dimension and channel dimension. For example, for a batch of 2D tensors of shape [B,C,W,H], [W,H] should be passed. For 3D data instead of shape [B,C,W,H,D], instead, [W,H,D] should be passed.

Returns:

Returns a tuple where the first element is a boolean signaling whether the given input shape

already fits the model’s requirements. If that value is False, the second element contains the closest shape that fits the model, otherwise it will be None.

Return type:

tuple[bool, Size]

class mfai.pytorch.models.unet.UNetSettings(init_features=64, autopad_enabled=False)

Bases: object

init_features: number of features of the first layer. This number will be used for each following layer. Default is 64. autopad_enabled: whether to allow autopadding of input image. Default is False.

Parameters:

• init_features (int)

• autopad_enabled (bool)

autopad_enabled: bool

classmethod from_dict(kvs, *, infer_missing=False)

Return type:

TypeVar(A, bound= DataClassJsonMixin)

Parameters:

kvs (dict | list | str | int | float | bool | None)

classmethod from_json(s, *, parse_float=None, parse_int=None, parse_constant=None, infer_missing=False, **kw)

Return type:

TypeVar(A, bound= DataClassJsonMixin)

Parameters:

s (str | bytes | bytearray)

init_features: int

classmethod schema(*, infer_missing=False, only=None, exclude=(), many=False, context=None, load_only=(), dump_only=(), partial=False, unknown=None)

Return type:

SchemaF[TypeVar(A, bound= DataClassJsonMixin)]

Parameters:

• infer_missing (bool)

• many (bool)

• partial (bool)

to_dict(encode_json=False)

Return type:

Dict[str, Union[dict, list, str, int, float, bool, None]]

to_json(*, skipkeys=False, ensure_ascii=True, check_circular=True, allow_nan=True, indent=None, separators=None, default=None, sort_keys=False, **kw)

Return type:

str

Parameters:

• skipkeys (bool)

• ensure_ascii (bool)

• check_circular (bool)

• allow_nan (bool)

• indent (int | str | None)

• separators (Tuple[str, str] | None)

• default (Callable | None)

• sort_keys (bool)