PyTorch-UNet

PyTorch Implementation for Foreground Segmentation and Saliency Prediction

Prerequisites

• Linux or macOS
• Python 2
• CPU or NVIDIA GPU + CUDA CuDNN

Getting Started

Installation

• Install PyTorch and dependencies from http://pytorch.org
• Install Torch vision from the source.

git clone https://github.com/pytorch/vision
cd vision
python setup.py install

• Clone this repo:

git clone https://github.com/LongLong-Jing/PyTorch-UNet
cd PyTorch-UNet

Unet train/test

• Train a model:

python train_D.py

• Test the model:

python test.py