U-Net-in-PyTorch

This is an implementation of the U-Net model from the paper, U-Net: Convolutional Networks for Biomedical Image Segmentation.

How does it work?

U-Net is made up of a contracting path and an expanding path. The contracting route is a series of convolutional layers and pooling layers that gradually diminish the resolution of the feature map. The expansive route is a succession of up-sampling and convolutional layers that gradually enhance the resolution of the feature map.

The relevant feature map from the contracting path is concatenated with the current feature map at each step in the expansive path.

Architecture

Example output

Image -> Mask -> Predicted Mask

Generated using this test script.

Implementation

Check the attached notebook for implementation details.