PyTorch Implementation for Foreground Segmentation and Saliency Prediction
- Linux or macOS
- Python 2
- CPU or NVIDIA GPU + CUDA CuDNN
- 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- Train a model:
python train_D.py- Test the model:
python test.py