端到端线框解析

本仓库包含论文《端到端线框解析》的官方 PyTorch 实现,作者为 周一超、齐浩志、马毅。该论文发表于 ICCV 2019。

简介

L-CNN 是一个概念简单但有效的神经网络,用于从给定图像中检测线框。它的性能大幅超过了之前最先进的线框和线条检测器。我们希望这个仓库能作为该领域未来研究的一个易于复现的基准。

主要结果

定性评估

LSD	AFM	Wireframe	L-CNN	真实标注

更多随机采样的结果可以在论文的补充材料中找到。

定量评估

下表报告了几种线框和线条检测器在上海科技大学数据集上的性能指标。

	上海科技 (sAP<sup>10</sup>)	上海科技 (AP<sup>H</sup>)	上海科技 (F<sup>H</sup>)	上海科技 (mAP<sup>J</sup>)
LSD	/	52.0	61.0	/
AFM	24.4	69.5	77.2	23.3
Wireframe	5.1	67.8	72.6	40.9
L-CNN	62.9	82.8	81.2	59.3

精确率-召回率曲线

代码结构

以下是每个文件功能的快速概览。

########################### 数据 ###########################
figs/
data/                           # 存放数据的默认文件夹
    wireframe/                  # 上海科技数据集(Huang et al.)的文件夹
logs/                           # 存储训练输出的默认文件夹
########################### 代码 ###########################
config/                         # 神经网络超参数和配置
    wireframe.yaml              # 上海科技数据集的默认参数
dataset/                        # 所有与数据生成相关的脚本
    wireframe.py                # 预处理上海科技数据集为npz的脚本
misc/                           # 不太重要的其他脚本
    draw-wireframe.py           # 生成图表网格的脚本
    lsd.py                      # 为LSD生成npz文件的脚本
    plot-sAP.py                 # 为所有算法绘制sAP10的脚本
lcnn/                           # lcnn模块,可以在其他脚本中"import lcnn"
    models/                     # 神经网络结构
        hourglass_pose.py       # 骨干网络(堆叠沙漏)
        line_vectorizer.py      # 采样器和线条验证网络
        multitask_learner.py    # 多任务学习网络
    datasets.py                 # 读取训练数据
    metrics.py                  # 评估指标函数
    trainer.py                  # 训练器
    config.py                   # 配置的全局变量
    utils.py                    # 其他函数
demo.py                         # 为图像检测线框的脚本
eval-sAP.py                     # sAP评估脚本
eval-APH.py                     # APH评估脚本
eval-mAPJ.py                    # mAPJ评估脚本
train.py                        # 训练神经网络的脚本
post.py                         # 后处理脚本
process.py                      # 从检查点处理数据集的脚本

复现结果

安装

为了便于复现,建议您在执行以下命令之前安装miniconda。

git clone https://github.com/zhou13/lcnn
cd lcnn
conda create -y -n lcnn
source activate lcnn
# 根据您的CUDA版本修改命令: 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

预训练模型

您可以从我们的HuggingFace仓库下载参考预训练模型。这些模型使用config/wireframe.yaml训练了312k次迭代。使用demo.py、process.py和eval-*.py来评估预训练模型。

为您自己的图像检测线框

要在您自己的图像上测试LCNN,您需要下载预训练模型并执行

python ./demo.py -d 0 config/wireframe.yaml <预训练pth路径> <图像路径>

这里,-d 0指定用于评估的GPU ID,您可以指定-d ""来强制CPU推理。

下载处理后的数据集

确保您的系统上安装了curl,然后执行

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

或者，你可以从我们的HuggingFace 仓库手动下载预处理好的数据集wireframe.tar.xz，然后按照相应步骤操作。

处理数据集

可选，你可以从头开始预处理（例如，生成热图，进行数据增强）数据集，而不是下载已处理好的数据集。如果你只想使用预处理好的数据集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

训练

默认的批量大小假设你有一张具有12GB显存的显卡，例如GTX 1080Ti或RTX 2080Ti。如果你的显存较少，可以减小批量大小。

要使用默认参数在GPU 0（由-d 0指定）上训练神经网络，执行以下命令：

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

测试预训练模型

要使用预训练模型在验证数据集上生成线框，执行以下命令：

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

后处理

要对神经网络的输出进行后处理（仅在你需要评估AP^H时必要），执行以下命令：

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

其中--plot是一个_可选_参数，用于控制程序是否除了生成包含线条信息的npz文件外，还生成可视化图像，--thresholds控制后处理的激进程度。--thresholds中的多个值便于进行超参数搜索。你应该将RUN和ITERATION替换为你的训练实例的所需值。

评估

要评估logs/下所有检查点的sAP（推荐），执行以下命令：

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

要评估mAP^J，执行以下命令：

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

要评估AP^H，你首先需要对结果进行后处理（参见上一节）。此外，AP^H评估需要MATLAB，且matlab应该在你的$PATH中。由于使用了parfor，强烈建议使用并行计算工具箱。后处理完成后，执行以下命令：

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

以获取图表，其中0_010是后处理中使用的阈值，post/RUN-ITERATION-APH是存储中间文件的临时目录。由于使用了像素级匹配，AP^H的评估可能需要长达一小时，具体取决于你的CPU。

有关评估的更多详细信息，请参阅eval-sAP.py、eval-mAPJ.py、eval-APH.py和misc/*.py的源代码。

引用端到端线框解析

如果你在研究中发现L-CNN有用，请考虑引用：

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