U-Net

Architecture
The architecture in UNET consists of two paths. Contracting path (encoder network) followed by an expansive path (decoder network).  U-Net gets its name from the architecture, like the letter U, as depicted in the image below.
Original paper by Olaf Ronneberger, Philipp Fischer, Thomas Brox
The encoder (left side) typically consists of a pre-trained classification network, such as ResNet, using convolution blocks followed by maxpool downsampling to encode the input image into feature representations at various levels. In this architecture, there is repeated application of two 3x3 convolutions. Each convolution is followed by a ReLU and batch normalization. Then a 2x2 max pooling operation is applied to reduce dimensions by half. Again, at each downsampling step, we double the number of feature channels, while we cut in half the spatial dimensions.
The decoder path (right side) consists of upsampling of the feature map followed by a 2x2 transpose convolution, which cuts the number of feature channels in half. The concatenation with the matching feature map from the contracting path, as well as a 3x3 convolutional filter (each followed by a ReLU) is also implemented. A 1x1 convolution is employed in the final layer to map the channels to the required number of classes.
U-Net differs from autoencoder architecture by overcoming the problem of compressed low dimensional representation of the input.
Code Implementation example
 U-Net with batch normalization for biomedical image segmentation with pretrained weights for abnormality segmentation in brain MRI
1
import torch
2
#Loading a U-Net model pre-trained for abnormality segmentation
3
model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet',
4
    in_channels=3, out_channels=1, init_features=32, pretrained=True)
5
​
6
# Download an example input image
7
import urllib
8
url, filename = ("https://github.com/mateuszbuda/brain-segmentation-pytorch/raw/master/assets/TCGA_CS_4944.png", "TCGA_CS_4944.png")
9
try: urllib.URLopener().retrieve(url, filename)
10
except: urllib.request.urlretrieve(url, filename)
11
​
12
import numpy as np
13
from PIL import Image
14
from torchvision import transforms
15
​
16
input_image = Image.open(filename)
17
m, s = np.mean(input_image, axis=(0, 1)), np.std(input_image, axis=(0, 1))
18
preprocess = transforms.Compose([
19
    transforms.ToTensor(),
20
    transforms.Normalize(mean=m, std=s),
21
])
22
input_tensor = preprocess(input_image)
23
input_batch = input_tensor.unsqueeze(0)
24
​
25
if torch.cuda.is_available():
26
    input_batch = input_batch.to('cuda')
27
    model = model.to('cuda')
28
​
29
with torch.no_grad():
30
    output = model(input_batch)
31
​
32
print(torch.round(output[0]))
Copied!
U-Net Segmentation results (IOU) on PhC-U373 and DIC-HeLa datasets.
Parameters in model playground
Encoder network
In Model Playground, we can select feature extraction (encoding) network to use as either ResNet or EffiecientNet. The depth of the ResNet model/ EfficientNet sub type has to be specified as well.
Weights
After selecting the desired encoder network, the weights to be used for model initialization can be chosen.
Further Resources
​U-Net: Convolutional Networks for Biomedical Image Segmentation​
​Stripping down U-Net for Segmentation of Biomedical Images on Platforms with Low Computational Budgets​
​Wiki entry for DeepLabv3+​
​Wiki entry for U-Net++​

Model architectures 🏛 - Previous

DeepLabv3+

Next - Model architectures 🏛

FBNetV3

Last modified 1mo ago

Copy link

Contents

Architecture

Code Implementation example

U-Net Segmentation results (IOU) on PhC-U373 and DIC-HeLa datasets.

Parameters in model playground

Encoder network

Weights

Further Resources