xwjabc / hed

A PyTorch reimplementation of Holistically-Nested Edge Detection
170 stars 42 forks source link

A PyTorch Reimplementation of HED

Introduction

This is a PyTorch reimplementation of Holistically-nested Edge Detection (HED). The code is evaluated on Python 3.6 with PyTorch 1.0 (CUDA9, CUDNN7) and MATLAB R2018b.

Instructions

Prepare
  1. Clone the repository:

    git clone https://github.com/xwjabc/hed.git
  2. Download and extract the data:

    cd hed
    wget https://cseweb.ucsd.edu/~weijian/static/datasets/hed/hed-data.tar
    tar xvf ./hed-data.tar
Train and Evaluate
  1. Train:

    python hed.py --vgg16_caffe ./data/5stage-vgg.py36pickle

    The results are in output folder. In the default settings, the HED model is trained for 40 epochs, which takes ~27hrs with one NVIDIA Geforce GTX Titan X (Maxwell).

  2. Evaluate:

    cd eval
    (echo "data_dir = '../output/epoch-39-test'"; cat eval_edge.m)|matlab -nodisplay -nodesktop -nosplash

    The evaluation process takes ~7hrs with Intel Core i7-5930K CPU @ 3.50GHz.

Besides, based on my observation, the evaluated performance is somewhat stable after 5 epochs (5 epochs: ODS=0.788 OIS=0.808 vs. 40 epochs: ODS=0.787 OIS=0.807).

Evaluate the Pre-trained Models
  1. Evaluate the my pre-trained version:

    python hed.py --checkpoint ./data/hed_checkpoint.pt --output ./output-mypretrain --test
    cd eval
    (echo "data_dir = '../output-mypretrain/test'"; cat eval_edge.m)|matlab -nodisplay -nodesktop -nosplash

    The result should be similar to ODS=0.787 OIS=0.807.

  2. Evaluate the official pre-trained version:

    python hed.py --caffe_model ./data/hed_pretrained_bsds.py36pickle --output ./output-officialpretrain --test
    cd eval
    (echo "data_dir = '../output-officialpretrain/test'"; cat eval_edge.m)|matlab -nodisplay -nodesktop -nosplash

    The result should be similar to ODS=0.788 OIS=0.806.

Acknowledgement

This reimplementation is based on lots of prior works. Thanks to Saining for the original Caffe implementation. Thanks to @meteorshowers for a PyTorch implementation where I adopt most of the code from. Thanks to @jmbuena for a fixed version of Piotr's Toolbox. Thanks to Berkeley Institute for Data Science which provides a mirror of BSDS500 dataset (the original link to the dataset seems broken).