CubiCasa5K：一个用于平面图图像分析的数据集和改进的多任务模型

1^11阿尔托大学计算机科学系，芬兰埃斯波 {firstname.lastname}@aalto.fi
2^22CubiCasa Inc., 芬兰奥卢 {firstname.lastname}@cubicasa.com

摘要

更好地理解建筑内部结构并进行建模，以及更令人印象深刻的 AR/VR 技术的出现，使得自动解析平面图图像的需求日益增长。然而，目前明显缺乏具有代表性的数据集来进一步研究该问题。为了弥补这一不足，本文提出了一个名为 CubiCasa5K 的新型图像数据集，这是一个大规模的平面图图像数据集，包含 5000 个样本，标注了超过 80 种平面图对象类别。数据集的标注使用多边形来区分不同对象，具有密集且灵活的特点。与基于强启发式规则和低级像素操作的传统方法不同，我们提出了一种基于改进的多任务卷积神经网络的方法。通过发布这个新颖的数据集和我们的实现，本研究为自动平面图图像分析提供了更丰富的工具集，从而能够以更全面的方式研究该问题，显著推动了该领域的研究。

关键词：平面图图像 · 数据集 · 卷积神经网络 · 多任务学习。

数据和代码地址：https://github.com/CubiCasa/CubiCasa5k

1 引言

平面图图像分析或理解长期以来一直是计算机视觉分支——自动文档分析中的一个研究课题。平面图是按比例绘制的图纸，从上方展示建筑物或公寓的结构，其目的是向观看者传达这种结构信息及其相关语义。平面图中的常见关键元素包括房间、墙壁、门、窗户和固定家具，但它们也可以涵盖更多技术信息，如建筑材料、电线或管道。

虽然平面图最初通常是使用 CAD 软件绘制并生成矢量图形格式，但对于房地产经济学中的常见用例（例如为了营销目的进行数字媒体上的打印或发布，如出售或出租），它们在发布前通常会被栅格化。然而，对于当今的应用程序（如 3D 房地产虚拟漫游或基于平面图的 3D 模型创建），这一过程是致命的，因为它丢弃了所有结构化的几何和语义信息，使得进一步利用这些平面图变得困难。

从栅格化的平面图图像中恢复丢失的信息并非易事。当前自动平面图图像分析的最先进模型基于深度卷积神经网络 (CNNs)。一种有前途的训练方案是仅使用一个网络主干，配合多个多任务头 (multi-task heads)，以恢复丢失的对象、结构和语义。

虽然结果已经很有前景 [13]，但用于训练和基准测试的数据集仍然相当小（少于 1000 张图像 [13]），与其他主流计算机视觉领域（如图像分类（数百万张图像 [18,19,9]）或图像分割（数万张图像 [15,20]））常用的数据集相比差距明显。众所周知，深度学习模型需要大量数据才能发挥效用，增加数据集规模很可能总是会带来更好的结果 [15,19]。

在本文中，我们提出了一个包含 5000 张图像的新型平面图图像数据集，其密集且丰富的真实标签 (ground-truth) 注释全部编码为多边形。该数据集涵盖了三种不同的平面图图像类别，即高质量、高质量建筑和彩色。由人类专家生成的注释涵盖了 80 多种不同的平面图元素类别。与之前最大的数据集 [13] 相比，所提出的数据集规模大了五倍以上，且由于包含了精确的对象形状和方向，其注释更加准确。它在公寓类型和绘图风格方面也表现出更大的变异性。作为一个强有力的基线，我们提出了一种受近期文献报道启发的全自动多任务学习方案。具体而言，我们使用了近期的“多任务不确定性损失 (multi-task uncertainty loss)”，它能够自动推导网络各损失项的权重。我们的初步结果表明，该方法在实践中具有巨大价值，因为在权重范围完全未知的情况下，它可以节省超参数调优的时间。

我们将这种损失与编码器-解码器卷积架构相结合，并在之前的平面图分析基准数据集上展示了最先进的结果。我们发布了所提出的新颖基准测试数据集以及我们的代码和训练好的模型，以便于轻松复现本文的结果。

2 相关工作

与许多视觉识别问题一样，2D 平面图分析的研究重点已从精心的特征工程转向依赖从数据中学习的方法。这种转变是由于能够在合理的时间内训练更大、更强大的深度神经网络 [11]。在我们的研究背景下，突破性的工作是 [13]，它提出了一种依赖 CNN 的自动平面图解析方法。该方法没有应用一堆低级图像处理及随后的启发式规则，而是训练了一个常规的全卷积网络来标记对象（房间、图标和开口）并定位墙壁节点。提取的低级信息随后被送入后处理器，以将原始平面图对象恢复为 2D 多边形。在 [13] 中，模型针对分割和定位任务进行了联合优化。主要发现是，深度神经网络可以作为最终后处理启发式规则的有效前置步骤，以恢复平面图元素（包括其几何和语义）。该方法显著提高了最先进水平，并启发了该领域的近期研究。

与 [13] 并行，[6] 提出了一种使用分割、目标检测和字符识别来解析平面图图像的基于 CNN 的方法。与 [13] 的主要区别在于，给定的任务都是使用独立的网络执行的。[6] 在墙壁分割上进行的实验清楚地证明了基于 CNN 的方法优于使用标准浅层分类器（如支持向量机）的传统基于补丁 (patch-based) 的模型。总之，深度神经网络时代为 2D 平面图分析带来了显著更好的方法。根据 [13,6]，特别是全卷积 CNN 在提取准确的像素级几何和语义信息方面具有巨大潜力，这些信息可在后续的后处理步骤中进一步利用，以构建更有效的启发式规则来恢复丢失的平面图元素。

为平面图解析构建更好的 CNN 的问题归结为两个与网络架构和目标相关的设计选择。语义分割研究的突破发生在引入全卷积网络 (FCNs) [14] 时。通用密集像素级预测的改进架构是带有跳跃连接 (skip-connections) 的 U-net 架构 [17]。如 [13] 所示，通过将自顶向下路径中的普通卷积层更改为残差块 (residual blocks) [7]，可以进一步提升容量。这种模型（也称为沙漏架构 (hourglass architecture) [3]）已在语义分割 [16] 和通过热图回归进行的人体姿态估计 [3] 等密集问题中被证明是有效的。最后的任务是选择训练目标。对于普通的分割问题，这通常是单一的交叉熵损失，或者在热图回归中是单一的欧几里得损失层。然而，实践中的许多问题（如我们的问题）可以从训练期间激活的多个目标中受益，这被称为多任务学习 [4]。使用这种方法的成功高度依赖于额外的超参数，即每个任务损失之间的相对权重。Kendall 等人 [10] 提出了一种简单的解决方案，可以在由分割、深度估计和实例分割组成的多任务设置中训练这种权重。与 [13] 相比，我们应用了 [10] 的方法（在 [12] 中进行了修订）来自动调整任务之间的权重，从而减少了对广泛超参数调优的需求。与 [13] 中报告的结果相比，我们的结果产生了显著的性能提升。

总之，目前关于自动平面图转换的研究仍然缺乏具有代表性的大规模数据集。目前，公开可用的最大注释数据集包含少于 1K 的样本。对象（例如不同的房间和图标类型）的多样性以及其注释的一致性和准确性（例如墙壁的厚度）都受到限制。这反过来意味着，还有进一步研究的空间，以调查使用内容更丰富的大型数据集来训练深度 CNN 的好处。在本文中，我们提出了一个包含 5K 样本的数据集，据我们所知，这是目前可用的最大的注释平面图数据集。

3 CubiCasa5K：一个新型的平面图数据集

CubiCasa5K 数据集是一个在线、部分手动的平面图矢量化流水线3^33的副产品，主要处理来自芬兰地区的房地产营销材料转换。它的主要使命是为研究界提供手段，以开发更准确的自动化模型，用于房地产和其他用例。

该数据集包含 5000 个平面图（带有人工注释），这些平面图是从更大的 15000 个（主要是芬兰的）平面图图像集合中收集和审查的。这些图像分为三个子类别：高质量建筑、高质量和彩色，分别包含 3732、992 和 276 个平面图。为了训练强大的机器学习模型，数据集被随机拆分为训练集、验证集和测试集，每个集合分别包含 4200、400 和 400 个平面图。注释采用每图像 SVG 矢量图形格式，每个注释都包含相应平面图中出现的所有元素的语义和几何注释。

3^33http://cubitool.cubi.casa.s3-website-us-west-2.amazonaws.com/?config=customize&rl=2&loc=na&id=8000&color=000000

注释及其一致性：所提出数据集的所有样本都经过了一个注释流水线，生成了具有丰富注释的矢量化平面图图像。单个输入始终是源自原始平面图绘制的栅格扫描（通常是扫描件）。注释由受过该任务培训的人工注释员手动完成。单张图像的注释耗时从 5 到 120 分钟不等，具体取决于源的复杂性、清晰度以及楼层数量。

每个平面图都遵循描述元素注释顺序的注释协议进行注释。这是为了利用给定平面图中先前已注释元素的所有可用信息（例如，墙壁是房间的边界）。注释是使用专为绘制平面图而定制的特殊 CAD 工具完成的。为了确保注释的一致性，有一个包含两个阶段的质量保证 (QA) 流程。该流程应用于每个已注释的样本图像。具体而言，应用的 QA 流程旨在控制注释的放置准确性以及正确的标签。该流程的第一轮由注释员完成，他们检查已注释的平面图并审查所有注释，最后纠正所有可能的错误。第二轮由不同的 QA 人员完成，他们执行与初始注释员相同的检查程序，并纠正第一轮中遗漏的任何错误。

数据集统计：图 2-4 提供了关于 CubiCasa5K 数据集的统计信息，突出了类别分布以及与 [13] 数据集相比平面图样本复杂性的各个方面。图 3 分别显示了排名的房间和图标类别的分布。在图 4 中，我们比较了 CubiCasa5K 数据集和 [13] 数据集中包含固定数量的已注释图标、墙壁和房间的图像频率。在图 2 中，我们报告了整个数据集中图像分辨率的分布。最后，在表 1 中，我们进一步将一些关键统计数据与所有现有的已注释平面图数据集进行了比较。

根据所有这些信息，可以得出结论，CubiCasa5K 是目前公开可用的最大且最多功能的平面图数据集。

(注：此处省略原文中的图表图像描述，保留表格内容)
表 1. 可用数据集之间的指标比较。

a^aa该数据集包含更多图标标签，但没有位置或多边形信息。我们忽略了这些图标。

4 我们的多任务模型

我们的任务是解析输入 2D 平面图图像中的所有常见元素。遵循 [13]，我们依赖一个网络，该网络输出两个分割图（一个用于不同的房间类型，一个用于不同的图标类型）以及一组热图，以精确定位墙壁节点、图标角点和开口端点（从现在开始，这三者统称为兴趣点 (interest points)）。使用定位的兴趣点，随后应用一组启发式规则来推断几何形状，即可能出现在图像中的所有元素的位置和尺寸。最后，使用这两个分割图来获取语义，即房间和图标的类型。我们的主要贡献在于流水线的后一步，我们应用了一个可训练的模块 [10] 来调整多任务损失项之间的相对权重。

网络架构。我们利用了 [13] 中使用的网络架构，该架构基于使用 ImageNet [5] 预训练的 ResNet-152 [7]。层的组织如图 5 顶部所示，给出了其中每个层操作的细节。遵循 [13]，图 5 中描绘的大部分网络层首先通过在 ImageNet [18] 上训练，然后在 MPII 人体姿态数据集 [2] 上进行初始化。为了使其针对本文研究的问题进行定制，必须进行一些更改。具体而言，D1 针对输入通道进行了更改（从 19 改为 3），并且最后两层（即 D9 和 D10）都被替换，以实现两个分割图和 21 个热图所需的输出通道数。因此，这三个给定的层（D1、D9 和 D10）必须进行随机初始化。

训练目标。在 [13] 中，该方法依赖于共享表示来密集预测像素的语义标签并回归兴趣点的位置。这意味着在网络末端（如图 5 所示）应用了多任务损失。具体而言，总共有 21 个输出图用于不同的兴趣点（墙壁节点、图标角点和开口端点）。学习的内容是通过单独的热图回归任务（所有任务均基于均方误差 (MSE) 作为训练目标）来实现所有兴趣点的像素级精确定位。除此之外，网络还输出两个分割图。第一个用于分割背景、房间和墙壁；第二个用于分割不同的图标和开口（窗户和门）。这两个分割任务都通过应用标准的交叉熵损失进行训练。在 [13] 中，所有任务都用于以多任务方式训练给定的共享表示，其相对权重是手动固定的。

Kendall 等人 [10] 的一项近期研究表明，多任务损失之间的相对权重可以自动学习。这将开发者从困难、耗时且非常昂贵的手动调优权重的步骤中解放出来。具体而言，权重是通过所谓的同方差不确定性项 (homoscedastic uncertainty terms) 隐式学习的，这些项被预测为每个任务的额外输出。详细信息可以在 [10] 中找到，我们在此直接转到最终损失，在我们的例子中，该损失表示为Ltot=LH+LSL_{tot} = L_H + L_SLtot​=LH​+LS​，其中：

LH=∑i[12σi2∥yi−fWi(x)∥+log⁡(1+σi)],(1)L_H = \sum_i \left[ \frac{1}{2\sigma_i^2} \| y_i - f_W^i(x) \| + \log(1+ \sigma_i) \right], \quad (1)LH​=i∑​[2σi2​1​∥yi​−fWi​(x)∥+log(1+σi​)],(1)

以及

LS=−∑k∈{rooms,icons}1σkyk⋅log⁡softmax(fWk(x))+log⁡σk.(2)L_S = - \sum_{k \in \{rooms, icons\}} \frac{1}{\sigma_k} y_k \cdot \log \text{softmax}(f_W^k(x)) + \log \sigma_k. \quad (2)LS​=−k∈{rooms,icons}∑​σk​1​yk​⋅logsoftmax(fWk​(x))+logσk​.(2)

LHL_HLH​用于训练热图回归器，它由一堆项（与需要定位的特定兴趣点数量一样多）组成，这些项基于重新加权的 MSE 进行最小化。权重与在训练期间学习的所谓不确定性参数σi\sigma_iσi​成反比。项log⁡(1+σi)\log(1+\sigma_i)log(1+σi​)[12] 充当正则化器以避免平凡解。此外，通过在取对数前加一，我们强制其始终为正 [12]。LSL_SLS​则用于分割部分，它由两个交叉熵项组成，在本例中用于房间和图标分割任务，以进行最小化。在这种情况下，权重出现时没有平方 [10]。根据我们的实验发现，正则化项log⁡σk\log \sigma_klogσk​在整个训练期间始终保持为正。

后处理。为了生成输入栅格化平面图的最终矢量图形等效表示，多任务 CNN 的输出被分派到一个由四个步骤组成的后处理器。目标是精确编码其位置、尺寸和类别标签的格式，提取给定输入中存在的所有平面图元素（墙壁、房间、图标和开口）。

后处理器从推断墙壁多边形开始。具体而言，该过程从与 [13] 中相同的步骤开始：基于节点的朝向成对连接节点，即如果有两个节点在垂直/水平方向上对齐（可能存在几个像素的未对齐），并且两者都有朝向彼此的节点。该过程产生一个墙壁骨架，接下来基于墙壁分割对其进行修剪。最后，通过沿墙壁线采样并检查墙壁分割图的强度分布来推断墙壁的宽度。

房间的位置和尺寸部分基于墙壁节点进行推断。具体而言，我们搜索所有跨越不包含任何节点的矩形区域的节点三元组。这导致了平面图内部的单元格网格化。然后根据基于房间分割图的投票机制对生成的单元格进行标记。最后，如果且仅当它们之间没有完全分隔的墙壁并且它们共享相同的房间标签时，所有相邻的单元格才会被合并。恢复图标的过程与房间提取非常相似，但我们利用来自负责图标角点热图预测的图的三元组，而不是墙壁节点热图。

最后，推断门和窗户。这是通过使用来自相应热图的预测，连接两个垂直/水平对齐的开口端点来完成的。标签再次基于分割图得出。开口的宽度与墙壁多边形相同。所有未落入墙壁分割内的此类开口端点都将被拒绝。

5 结果

在本节中，我们介绍评估指标和获得的结果。在展示我们新颖的 CubiCasa5K 数据集的基线结果之前，我们在 [13] 使用的相同数据集上验证了我们的方法。

初步实验。遵循 [13]，网络使用 [3] 的人体姿态网络权重进行初始化（在 ImageNet 和 MPII 上训练）。那些必须被替换的层（见第 4 节）被随机初始化。

我们使用不确定性驱动的任务权重对网络进行了 400 个 epoch 的训练，批量大小为 20。数据增强包括 90 度旋转、颜色抖动，以及在裁剪和缩放至 256x256 并填充零之间随机选择。我们使用了 Adam 优化器，初始学习率为1×10−31 \times 10^{-3}1×10−3，ϵ=1×10−8\epsilon = 1 \times 10^{-8}ϵ=1×10−8，β\betaβ值为 0.9 和 0.999。我们使用了一个调度器，如果在前 20 个 epoch 中基于验证损失没有观察到改进，则将学习率降低 0.1 倍。下降后，训练从那时起产生最佳验证损失的阶段继续，直到结束。最后，基于验证损失选择最佳模型。

根据我们的实验，学习率只需下降一次，训练似乎在 epoch 数接近 300 时收敛。在 Nvidia GeForce GTX TitanX GPU 卡上，训练耗时三个小时。

为了评估我们的模型，我们使用了与 [13] 中相同的评估设置。如表 2 所示，我们显著改进了 [13] 中提出的结果。我们进一步应用了测试时增强 (TTA) 方案，其中最终预测基于将同一图像四次输入到同一网络中，每次旋转 90 度。最终预测基于这四次预测的平均值。可以看出，无论是否使用整数规划 (IP)，这种增强在两种情况下似乎都是有益的。

表 2. 在 [13] 提出的数据集上的评估结果。
(注：此处简化表格呈现，保留核心指标)

我们注意到 [13] 提出的数据集的原始注释4^44中存在错误。在使用我们的修复重新评估 [13] 的模型后，我们注意到 [13] 的性能（见“our eval”）实际上比 [13] 中最初报告的要好。我们进一步通过遵循原始论文中报告的细节训练了模型（“best reproduced from [13]”），结果或多或少相似。最后，我们将没有测试时增强的最佳模型（“Ours”）与 [13] 的“our eval”版本进行了比较，可以看出，我们的模型明显更好。

CubiCasa5K 实验。在利用 CubiCasa5K 数据集的当前实验中，一些原始的房间类型和图标类型被耦合，因此我们的目标总共涵盖 12 个房间类别和 11 个图标类别（见表 4 中选择的类别，以及项目网站上的更多详细信息）。至于其他细节，网络包含相同的热图回归层，并使用与前一个实验相同的目标进行训练。然而，对训练方案进行了以下调整：我们使用前一个实验在 ImageNet、MPII 和 Liu 等人 [13] 数据集上预训练的权重开始训练。我们在前 100 个 epoch 使用第 5.1 节中给出的相同增强进行训练。之后，我们继续使用当时的最佳权重（根据验证集上的损失），首先将优化器参数初始化为其起始值，然后放弃将图像调整大小到 256x256 的增强。然后我们训练网络 400 个 epoch，这导致了收敛。

遵循语义分割领域的常见做法 [14,20]，我们使用三个评估指标报告结果，即指示正确分类像素比例的整体准确率 (Overall Accuracy)，以及所有类别平均的正确分类像素比例的平均准确率 (Mean Accuracy)。最后，我们报告平均交并比 (Mean IoU)，它指示预测像素和真实像素之间重叠区域的面积，在所有类别上取平均。我们进一步报告了关于原始分割和多边形化 § 实例（即后处理步骤之后）的结果。我们采取了与 [13] 不同的模型评估方法，因为我们认为平面图解析问题非常接近语义分割问题。

我们在表 3 中报告了关于所述指标的性能。根据结果，原始分割测试得分明显优于基于多边形化分割实例的得分。主要原因是，如果遗漏了墙壁或图标节点或它们没有被正确定位，则无论分割质量如何，都无法创建多边形。在表 4 中，我们进一步报告了本研究中使用的关于所有房间和图标类别的类别特定 IoU 和准确率。图 1 展示了我们流水线的一个示例结果。

表 3. CubiCasa5K 数据集的评估结果。

6 结论

在本文中，我们提出了一个名为 CubiCasa5K 的新型平面图图像数据集。与其他现有的已注释平面图数据集相比，我们的数据集大了 5 倍以上，并且其注释更加多样化，涵盖了超过 80 个平面图对象类别。与该新颖数据集一起，我们提供了使用改进的多任务卷积神经网络的基线结果，该网络产生了最先进的性能。

对于未来的方向，我们计划将 [6] 中使用的目标检测器作为任务之一集成到我们的流水线中。尝试 [1] 的方法直接推断平面图元素为多边形也将是很有趣的。

参考文献

1. Acuna, D., Ling, H., Kar, A., Fidler, S.: Efficient interactive annotation of segmentation datasets with polygon-rnn++. In: Proc. CVPR (2018)
2. Andriluka, M., Pishchulin, L., Gehler, P., Schiele, B.: 2D human pose estimation: New benchmark and state of the art analysis. Proc. CVPR pp. 3686–3693 (2014)
3. Bulat, A., Tzimiropoulos, G.: Human pose estimation via convolutional part heatmap regression. Proc. ECCV pp. 717–732 (2016)
4. Caruana, R.: Multitask learning. Machine Learning 28(1), 41–75 (Jul 1997)
5. Deng, J., Dong, W., Socher, R., Li, L., Li, K., Fei-Fei, L.: Imagenet: A large-scale hierarchical image database. In: Proc. CVPR. pp. 248–255 (2009)
6. Dodge, S., Xu, J., Stenger, B.: Parsing floor plan images. In: MVA. pp. 358–361 (2017)
7. He, K., Zhang, X., Ren, S., Sun, J.: Deep Learning for Image Recognition. Proc. CVPR pp. 770–778 (2016)
8. de las Heras, L.P., Terrades, O., Robles, S., Sánchez, G.: Cvc-fp and sgt: a new database for structural floor plan analysis and its groundtruthing tool. IJDAR (2015)
9. Hinton, G., Vinyals, O., Dean, J.: Distilling the Knowledge in a Neural Network. NIPS (2015)
10. Kendall, A., Gal, Y., Cipolla, R.: Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In: Proc. CVPR (2018)
11. Lecun, Y., Bengio, Y., Hinton, G.: Deep learning. Nature 521(7553), 436–444 (2015)
12. Liebel, L., Körner, M.: Auxiliary tasks in multi-task learning. CoRR abs/1805.06334 (2018)
13. Liu, C., Wu, J., Kohli, P., Furukawa, Y.: Raster-to-vector: Revisiting floorplan transformation. In: The IEEE International Conference on Computer Vision (ICCV) (2017)
14. Long, J., Shelhamer, E., Darrell, T.: Fully convolutional networks for semantic segmentation. In: Proc. CVPR (2015)
15. Neuhold, G., Ollmann, T., Rota Bulo, S., Kontschieder, P.: The mapillary vistas dataset for semantic understanding of street scenes. In: Proc. ICCV (2017)
16. Pohlen, T., Hermans, A., Mathias, M., Leibe, B.: Full-resolution residual networks for semantic segmentation in street scenes. In: Proc. CVPR (2017)
17. Ronneberger, O., Fischer, P., Brox, T.: U-net: Convolutional networks for biomedical image segmentation. In: Navab, N., Hornegger, J., Wells, W.M., Frangi, A.F. (eds.) Proc. MICCAI. pp. 234–241 (2015)
18. Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M., Berg, A.C., Fei-Fei, L.: Imagenet large scale visual recognition challenge. IJCV 115(3), 211–252 (2015)
19. Sun, C., Shrivastava, A., Singh, S., Gupta, A.: Revisiting unreasonable effectiveness of data in deep learning era. In: Proc. ICCV (2017)
20. Zhou, B., Zhao, H., Puig, X., Fidler, S., Barriuso, A., Torralba, A.: Scene parsing through ade20k dataset. In: Proc. CVPR (2017)