SegmentationModels

UNet(in_channels::Integer=3, num_classes::Integer=1; init_channels::Integer=16, stages::Integer=4, final_activation=sigmoid)

Instantiate a UNet with a given number of input channels and output classes. The inital number of convolution channels can be set and also the number of pooling operations (stages). The final activation of the segmentation head can also be supplied. The model applies two convolutions with relu activation and BatchNorm at each encoder or decoder stage. It doubles the number of convolution channels at each down-stage. Upsampling is done bilinear.