/pytorch-hed

Primary LanguageJupyter NotebookMIT LicenseMIT

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

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.