端到端线框解析

本仓库包含论文的官方 PyTorch 实现：Yichao Zhou, Haozhi Qi, Yi Ma. "End-to-End Wireframe Parsing." ICCV 2019。

简介

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

主要结果

定性评估

LSD	AFM	Wireframe	L-CNN	Ground Truth

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

定量评估

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

	上海科技大学 (sAP10)	上海科技大学 (APH)	上海科技大学 (FH)	上海科技大学 (mAPJ)
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 等人）存放目录
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 <预训练模型路径> <图像路径>

其中，-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 <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}
}

L-CNN 快速上手指南

L-CNN 是一个用于端到端线框解析（Wireframe Parsing）的 PyTorch 实现，能够高效地从图像中检测线框结构，性能优于传统的线检测器。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.6+
• GPU: 支持 CUDA 的 NVIDIA 显卡（训练推荐显存 12GB+，如 GTX 1080Ti/RTX 2080Ti；推理可酌情降低）
• 依赖软件: git, curl, wget

注意：若需评估 $AP^H$ 指标，必须安装 MATLAB 并将其加入系统 $PATH，同时建议安装 Parallel Computing Toolbox。

安装步骤

推荐使用 miniconda 管理环境以保证复现性。

1. 克隆仓库并创建环境

   git clone https://github.com/zhou13/lcnn
   cd lcnn
   conda create -y -n lcnn python=3.7
   source activate lcnn
   

2. 安装深度学习依赖 请根据您的 CUDA 版本调整 cudatoolkit 版本号（参考 PyTorch 官网）。以下示例适用于 CUDA 10.1：

   # 安装 PyTorch (国内用户可使用清华源加速：-c https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/)
   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
   

3. 创建必要目录

   mkdir data logs post
   

基本使用

1. 下载预训练模型

从 HuggingFace 下载官方提供的预训练权重（基于 ShanghaiTech 数据集训练 312k 次迭代）：

• 下载地址：HuggingFace Repo
• 将下载的 .pth 文件放置在任意方便访问的路径。

2. 检测自定义图像的线框

使用预训练模型对单张图片进行线框检测是最简单的用法。

命令格式：

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

示例： 假设您使用 GPU 0，预训练模型路径为 pretrained.pth，待检测图片为 test.jpg：

python ./demo.py -d 0 config/wireframe.yaml pretrained.pth test.jpg

• -d 0: 指定使用的 GPU ID。
• -d "": 强制使用 CPU 进行推理（速度较慢）。
• 输出结果通常包含可视化的线框图及对应的数据文件。

3. (可选) 使用处理好的数据集

如果您计划复现论文结果或进行微调，可下载已预处理好的 ShanghaiTech 数据集：

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