Awesome
hed_pytorch
Pytorch implementation of Holistically-nested Edge Detection (HED). This repo can be used to train a HED model and perform test inference. The implementation was evaluated on the BSDS dataset.
Prerequisites
- Compute environment with GPU (optional).
- Python environment: Anaconda is recommended.
- SciPy stack: Install instructions.
- [Optional] CUDA and NVidia Drivers: Install instructions.
- Pytorch: Install instructions.
I use docker to avoid dependency problems. Installation instructions for my setup are available here.
Usage
- Download repository:
$ git clone https://github.com/buntyke/pytorch-hed.git
- Create
data
folder, download and extract BSDS dataset into folder:$ cd pytorch-hed/ $ mkdir data; cd data $ wget http://vcl.ucsd.edu/hed/HED-BSDS.tar $ tar -xvf HED-BSDS.tar $ rm HED-BSDS.tar $ cd HED-BSDS/ $ head -n 10 train_pair.lst > val_pair.lst $ cd ../../
- Download the VGG pretrained model to initialize training
$ mkdir model; cd model/ $ wget https://download.pytorch.org/models/vgg16-397923af.pth $ mv vgg16-397923af.pth vgg16.pth $ cd ..
- Train HED model by running
train.py
ortrain.ipynb
notebook following the instructions:
The trained model along with validation results are stored in the train folder.$ python train.py
Performance
Files
- train.ipynb: Notebook to train HED model.
- trainer.py: Helper class to train model and perform validation.
- model.py: HED model definition given through several class implementations.
- dataproc.py: Dataset class implementation used in Trainer class.
Acknowledgement
The source code is derived from three different implementations available online. Thanks to @s9xie for original caffe implementation. Thanks to @EliasVansteenkiste, @xlliu, @BinWang-shu for the pytorch implementations.