This repository contains the official PyTorch implementation of the paper: Yichao Zhou, Haozhi Qi, Yi Ma. "End-to-End Wireframe Parsing." ICCV 2019.

L-CNN is a conceptually simple yet effective neural network for detecting the wireframe from a given image. It outperforms the previous state-of-the-art wireframe and line detectors by a large margin. We hope that this repository serves as an easily reproducible baseline for future researches in this area.

LSD	AFM	Wireframe	L-CNN	Ground Truth

More random sampled results can be found in the supplementary material of the paper.

The following table reports the performance metrics of several wireframe and line detectors on the ShanghaiTech dataset.

Below is a quick overview of the function of each file.

########################### Data ########################### figs/ data/ # default folder for placing the data wireframe/ # folder for ShanghaiTech dataset (Huang et al.) logs/ # default folder for storing the output during training ########################### Code ########################### config/ # neural network hyper-parameters and configurations wireframe.yaml # default parameter for ShanghaiTech dataset dataset/ # all scripts related to data generation wireframe.py # script for pre-processing the ShanghaiTech dataset to npz misc/ # misc scripts that are not important draw-wireframe.py # script for generating figure grids lsd.py # script for generating npz files for LSD plot-sAP.py # script for plotting sAP10 for all algorithms lcnn/ # lcnn module so you can "import lcnn" in other scripts models/ # neural network structure hourglass_pose.py # backbone network (stacked hourglass) line_vectorizer.py # sampler and line verification network multitask_learner.py # network for multi-task learning datasets.py # reading the training data metrics.py # functions for evaluation metrics trainer.py # trainer config.py # global variables for configuration utils.py # misc functions demo.py # script for detecting wireframes for an image eval-sAP.py # script for sAP evaluation eval-APH.py # script for APH evaluation eval-mAPJ.py # script for mAPJ evaluation train.py # script for training the neural network post.py # script for post-processing process.py # script for processing a dataset from a checkpoint

For the ease of reproducibility, you are suggested to install miniconda before following executing the following commands.

git clone https://github.com/zhou13/lcnn cd lcnn conda create -y -n lcnn source activate lcnn # Modify the command with your CUDA version: https://pytorch.org/ conda install -y pytorch cudatoolkit=10.1 -c pytorch conda install -y tensorboardx -c conda-forge conda install -y pyyaml docopt matplotlib scikit-image opencv mkdir data logs post

You can download our reference pre-trained models from our HuggingFace Repo. Those models were trained with config/wireframe.yaml for 312k iterations. Use demo.py, process.py, and eval-*.py to evaluate the pre-trained models.

To test LCNN on your own images, you need download the pre-trained models and execute

python ./demo.py -d 0 config/wireframe.yaml <path-to-pretrained-pth> <path-to-image>

Here, -d 0 is specifying the GPU ID used for evaluation, and you can specify -d "" to force CPU inference.

Make sure curl is installed on your system and execute

cd data wget https://huggingface.co/yichaozhou/lcnn/resolve/main/Data/wireframe.tar.xz tar xf wireframe.tar.xz rm wireframe.tar.xz cd ..

Alternatively, you can download the pre-processed dataset wireframe.tar.xz manually from our HuggingFace Repo and proceed accordingly.

Optionally, you can pre-process (e.g., generate heat maps, do data augmentation) the dataset from scratch rather than downloading the processed one. Skip this section if you just want to use the pre-processed dataset wireframe.tar.xz.

cd data wget https://huggingface.co/yichaozhou/lcnn/resolve/main/Data/wireframe_raw.tar.xz tar xf wireframe_raw.tar.xz rm wireframe_raw.tar.xz cd .. dataset/wireframe.py data/wireframe_raw data/wireframe

The default batch size assumes your have a graphics card with 12GB video memory, e.g., GTX 1080Ti or RTX 2080Ti. You may reduce the batch size if you have less video memory.

To train the neural network on GPU 0 (specified by -d 0) with the default parameters, execute

python ./train.py -d 0 --identifier baseline config/wireframe.yaml

To generate wireframes on the validation dataset with the pretrained model, execute

./process.py config/wireframe.yaml <path-to-checkpoint.pth> data/wireframe logs/pretrained-model/npz/000312000

To post process the outputs from neural network (only necessary if you are going to evaluate AP^H), execute

python ./post.py --plot --thresholds="0.010,0.015" logs/RUN/npz/ITERATION post/RUN-ITERATION

where --plot is an optional argument to control whether the program should also generate images for visualization in addition to the npz files that contain the line information, and --thresholds controls how aggressive the post processing is. Multiple values in --thresholds is convenient for hyper-parameter search. You should replace RUN and ITERATION to the desired value of your training instance.

To evaluate the sAP (recommended) of all your checkpoints under logs/, execute

python eval-sAP.py logs/*/npz/*

To evaluate the mAP^J, execute

python eval-mAPJ.py logs/*/npz/*

To evaluate AP^H, you first need to post process your result (see the previous section). In addition, MATLAB is required for AP^H evaluation and matlab should be under your $PATH. The parallel computing toolbox is highly suggested due to the usage of parfor. After post processing, execute

python eval-APH.py post/RUN-ITERATION/0_010 post/RUN-ITERATION/0_010-APH

to get the plot, where 0_010 is the threshold used in the post processing, and post/RUN-ITERATION-APH is the temporary directory storing intermediate files. Due to the usage of pixel-wise matching, the evaluation of AP^H may take up to an hour depending on your CPUs.

See the source code of eval-sAP.py, eval-mAPJ.py, eval-APH.py, and misc/*.py for more details on evaluation.

If you find L-CNN useful in your research, please consider citing:

@inproceedings{zhou2019end, author={Zhou, Yichao and Qi, Haozhi and Ma, Yi}, title={End-to-End Wireframe Parsing}, booktitle={ICCV 2019}, year={2019} }