Implement the U-net architecture for cell image data segmentation using PyTorch.
Figure 1: U-Net architecture[1]
- a 3 × 3 un-padded convolution layer which takes a feature map with
input_channel
number of channels and outputs a feature map withoutput_channelnumber of channels - a ReLU activation function
- another 3 × 3 un-padded convolution layer which keeps the number of channels unchanged
- a Batch normalization layer
- a ReLU activation function
Since the size of the data is too small for training a neural network with a huge number of parameters. Under this situation, the code has the following data augmentation applied:
- Horizontal/Vertical flip
- Zooming
- Rotation
Left is label, right is the corresponding prediction.
[1]Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. InInternational Conference on Medical image computing and computer-assisted intervention, pages 234–241. Springer, 2015.