briancuestas/pytorch-hed

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.