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
Usage
Performance
Files
Acknoledgement
References

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 or train.ipynb notebook following the instructions:
```
$ python train.py 
```
The trained model along with validation results are stored in the train folder.

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.