unet

class src.unet.UNet(*args: Any, **kwargs: Any)

Simple PyTorch U-Net

Leaving all default parameters will create a U-Net as described in the original paper: “U-Net: Convolutional Networks for Biomedical Image Segmentation” (Ronneberger, et al.)

Parameters

in_channels: int, optional

Number C if input channels in a batch of (N x C x H x W). Defaults to 1 corresponding to a standard black-and-white image.

out_channels: int, optional

Number C of output channels in a batch of (N x C x H x W). Corresponds to the amount of classes that should be differentiated, e.g. C = 10 if you want to differentiate between 10 different segmentation classes. Defaults to 2.

depth: int, optional

Depth of the U-Net. A depth of 5 corresponds to the U-Net described in the original paper, i.e. 5 convolutional layers and 4 up-convolutional layers. Default is 5.

start_features: int, optional

Number of initial features in the convolutional blocks. Features are doubled (or halved when up-convoluting) after every block. Default is 64.

up: str, optional

Mode of the up-convolution. May be either 'upsample' or 'upconv'. Default is 'upsample'.

padding: int, optional

Padding to be applied when convoluting. User must ensure validity. Default is 0.

verbose: bool, optional

Verbose mode. Whether or not to print debug information about the batch size after every convolution step. Useful for fixing image size and padding problems. Default is False.

Methods

__call__(*args, **kwargs)

Call self as a function.

forward