pytorch>> Image-Segmentation-PyTorch>> 返回
项目作者： hiyouga

项目描述 ：
U-Net for image segmentation, PyTorch implementation.
高级语言： Python
项目地址: git://github.com/hiyouga/Image-Segmentation-PyTorch.git
创建时间： 2019-08-13T09:46:01Z
项目社区：https://github.com/hiyouga/Image-Segmentation-PyTorch
开源协议：MIT License
下载


Image-Segmentation-PyTorch

PyTorch implementation of the U-Net for image segmentation.

LICENSE

Requirement

  • Python 3
  • numpy
  • matplotlib
  • Pillow
  • torch
  • torchvision
  • pydensecrf

Dataset

Based on the shoe dataset [Google Drive] provided by our teacher.

example

Usage

Training

  1. python train.py

Inference

  1. python train.py --inference True --checkpoint [*.pt]

The checkpoint files can be found in the state_dict folder.

Show help message

  1. python train.py -h

Implemented model

Ronneberger, O., Fischer, P., and Brox, T. 2015. U-Net: Convolutional Networks for Biomedical Image Segmentation. In Medical Image Computing and Computer-Assisted Intervention (MICCAI). [pdf]

unet

Notes

The model was trained from scratch on Tesla V100 32GB*4. Training the model takes 3.6GB of memory and predicting images takes 1.3GB. If you don’t have enough GPU memory, consider using bilinear up-sampling rather than transposed convolution in the model.

Acknowledgements

  • Some of the code is borrowed from milesial.
  • This is a personal homework for “Machine Learning Theory and Application” in BUAA Summer School.

License

MIT