SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We have implemented the system with PyTorch and used a machine with 4 NVIDIA RTX A5000 GPUs. We have borrowed the official codebase of Karras et al.[19] to implement the diffusion model framework. The guidance networks are implemented as set-to-set Transformer decoders, and the loss weights for the guidance training are  $\lambda_{1} = 1$ ,  $\lambda_{2} = 0.05$ ,  $\lambda_{3} = 0.1$ . The implementation of the denoising network refers to the competing state-of-the-art, RoomFormer [47] and MapTR [23] for floorplan and HD map reconstruction, respectively. Concretely, we employ the same ResNet [12] image encoder, DETR-style Transformer architectures (discussed in §4), learning rate settings, and optimizer settings as RoomFormer or MapTR, while increasing the number of training iterations ( $\times 2.5$  for floorplan and  $\times 1.5$  for HD map), as the model converges slower under a denoising formulation than a detection/regression formulation. Note that we tried the same training schedule for the competing methods but their performance dropped due to overfitting (see §5.3 for ablation study). For each task, we use the same trained model for different types of proposals (i.e., existing methods or rough annotations). Complete implementation details are in Appendix B, and supplementary experiments are in Appendix D
5.1 Floorplan reconstruction
Dataset and metrics: Structured3D dataset [48] contains 3500 indoor scenes (3000/250/250 for training/validation/test) with diverse house floorplans. The average/maximum numbers of polygons and vertex per polygon across the dataset are 6.29/17 and 5.72/38, respectively. Point clouds are converted to  $256 \times 256$  top-view point-density images as inputs. We use the same evaluation metrics as previous works [7, 42, 47], which consists of three levels of metrics with increasing difficulty: room, corner, and angle. The precision, recall, and F1 score are reported at each metric level.
Competing approaches: We compare PolyDiffuse with five approaches from the literature: Floor-SP [5], MonteFloor [42], LETR [46], HEAT [7], and RoomFormer [47]. The first two approaches design sophisticated optimization algorithms to solve for the optimal floorplan with learned metric functions, while the latter three use end-to-end Transformer-based neural networks, thus being much faster. For our proposal generator RoomFormer, we simplify the implementation by removing the Dice loss. As the dataset is very small, we add a random rotation data augmentation and train the modified RoomFormer for  $4 \times$  the original training iterations, which improves the overall performance.
Table 1: Quantitative evaluation on Structured3D test set [48]. PolyDiffuse outperforms the state-of-the-art methods by clear margins, and works well with rough annotations. Results of previous works are copied from [47]. * indicates our modified implementation. †: The running time is measured on a single Nvidia RTX A5000 GPU, and we only report those run by ourselves. When using RoomFormer to produce the initial proposals, the running time counts both the RoomFormer proposal generator and the GS-DM.
Evaluation Level → Method	Stages	Steps	\( FPS^† \)	Room	Corner	Angle
Prec.	Rec.	F1	Prec.	Rec.	F1	Prec.	Rec.	F1
Floor-SP [5]	2	-	-	89.	88.	88.	81.	73.	76.	80.	72.	75.
MonteFloor [42]	2	500	-	95.6	94.4	95.0	88.5	77.2	82.5	86.3	75.4	80.5
LETR [46]	1	1	-	94.5	90.0	92.2	79.7	78.2	78.9	72.5	71.3	71.9
HEAT [7]	2	3	-	96.9	94.0	95.4	81.7	83.2	82.5	77.6	79.0	78.3
RoomFormer [47]	1	1	-	97.9	96.7	97.3	89.1	85.3	87.2	83.0	79.5	81.2
RoomFormer*	1	1	29.9	96.3	96.2	96.2	89.7	86.7	88.2	85.4	82.5	83.9
+PolyDiffuse(Ours)	2	2	11.7	96.9	96.4	96.6	90.3	87.1	88.7	85.8	82.8	84.3
+PolyDiffuse(Ours)	2	5	7.1	98.5	97.9	98.2	92.5	89.0	90.7	90.3	86.9	88.6
+PolyDiffuse(Ours)	2	10	4.4	98.7	98.1	98.4	92.8	89.3	91.0	90.8	87.4	89.1
Rough annotations	1	-	-	17.9	18.2	18.0	1.3	1.4	1.3	0.1	0.1	0.1
+PolyDiffuse(Ours)	2	10	19.2	97.4	98.2	97.8	91.7	92.2	91.9	89.2	89.7	89.4
Quantitative & qualitative evaluation: Table 1 presents the quantitative evaluation results. We tried two proposal generators: 1) the improved RoomFormer and 2) rough annotations. The rough annotations are prepared by turning the ground-truth data into fixed-radius small circles centered
Figure 4: Qualitative results for floorplan reconstruction. Left: PolyDiffuse improves the geometry of RoomFormer and turns rough annotations into decent reconstructions. Right: PolyDiffuse enables search-based self-refinement by enumerating different vertex numbers of polygons and evaluating its own reconstruction likelihood (i.e., NLL per dimension).
at the instance centroid (see Figure 4), which simulates what a human annotator can provide easily (i.e., indicating instance center and the number of vertices by looking at images). When using RoomFormer as the proposal generator, PolyDiffuse outperforms all competing methods by clear margins across all the evaluation metrics. The performance gap increases with the difficulty level of the metric, indicating that PolyDiffuse is superior in high-level geometry reasoning. With the rough annotations as the initial proposal, PolyDiffuse also produces reasonable reconstructions. The recalls are especially high because of the correct number of vertices. Note that our training is not dependent on any specific proposal generator, but PolyDiffuse works well with different generators at test time. Figure 4 presents concrete qualitative examples of how PolyDiffuse improves the initial reconstruction and enforces better geometry correctness. We also provide more qualitative examples in Appendix D.3.
Running speed analysis: Table 1 also presents the speed-performance tradeoff of PolyDiffuse on the floorplan reconstruction task. During the denoising process, the image encoding parts of the denoising network only run once, while the transformer decoder part runs for multiple rounds. In our computation environment, the time of image encoding vs. the time of transformer decoder is 2:1. Since speed is not a crucial factor for floorplan reconstruction, we can just use more denoising steps in real applications for better reconstruction quality.
Likelihood-based refinement: Figure 4(right) demonstrates how we leverage the likelihood evaluation to refine reconstruction results with minimal user interaction. We manually indicate the imperfect polygons (inside dashed boxes at the left), then run our system with different numbers of vertices for these polygons while evaluating the reconstruction likelihood. The result with the best negative log-likelihood (NLL) per dimension is the refined reconstruction.
5.2 HD map reconstruction
Dataset and metrics: The nuScenes dataset [2] provides a standard benchmark for HD map reconstruction. The dataset contains 1000 ego-centric sequences collected by autonomous vehicles with rich annotations. The data is annotated at  $2\mathrm{Hz}$ . Each sample has 6 RGB images and LiDAR sweeps captured by onboard sensors, covering the  $360^{\circ}$  horizontal FOV. In the top-down Bird-eye-view (BEV) space, the perception range is  $[-15m, 15m]$  for X-axis and  $[-30m, 30m]$  for Y-axis. The map includes three categories of elements: pedestrian crossing, divider, and road boundary. For fair comparisons, we use the same dataset split and pre-processing setups as previous works [22, 23, 25], and also represent each polyline with 20 uniformly-interpolated vertices. Similar to MapTR [23], we only use the RGB images as the sensor inputs in our experiments.
We follow the common evaluation protocol from the literature [23, 25], which employs the average precision (AP) with the Chamfer distance as a matching criterion. The AP is averaged over three
Table 2: Quantitative evaluation results of HD map construction on nuScenes[2]. All methods in the table use RGB inputs only and employ ResNet50 as the image backbone.  $^+$ : Results are obtained by running the official code with the released model checkpoint.  $\dagger$ : The running time here is measured on a single Nvidia RTX A5000 GPU, and we only report for the methods run by ourselves. When using MapTR to produce the proposals, the running time counts both the MapTR proposal generator and the GS-DM.
Matching Criterion →	Chamfer distance	+ Ordered angle distance
Method	Stages	Steps	FPS†	APp	APd	APb	mAP	APp	APd	APb	mAP
VectorMapNet [25]	1	20	-	36.1	47.3	39.3	40.9	-	-	-	-
VectorMapNet [25]	2	20	-	42.5	51.4	44.1	46.0	-	-	-	-
VectorMapNet+	1	20	3.9	40.0	47.6	39.0	42.2	33.0	44.5	27.3	34.9
MapTR [23]	1	1	-	56.2	59.8	60.1	58.7	-	-	-	-
MapTR+	1	1	14.3	55.8	60.9	61.1	59.3	46.1	43.4	41.9	43.8
+PolyDiffuse(Ours)	2	2	6.3	56.8	59.8	60.9	59.2	50.3	48.2	44.3	47.6
+PolyDiffuse(Ours)	2	5	4.8	58.1	59.7	61.2	59.6	51.8	49.5	45.4	48.9
+PolyDiffuse(Ours)	2	10	3.4	58.2	59.7	61.3	59.7	52.0	49.5	45.4	49.0
Rough annotations	-	-	-	18.5	2.2	0.7	7.1	8.7	0.0	0.0	2.9
+PolyDiffuse(Ours)	2	10	11.3	55.3	60.4	55.3	57.0	48.3	49.5	38.1	45.3
distance thresholds:  $\{0.5m, 1.0m, 1.5m\}$ , and the mean average precision (mAP) is obtained by averaging across the three classes of map elements. However, the Chamfer distance does not consider the vertex order and vertex matching between prediction and G.T., thus failing to measure the structured correctness. As directional information is critical for autonomous vehicles, we propose to augment the matching criterion with an extra order-aware angle distance. Concretely, we find the optimal vertex matching between a prediction and a G.T. with the smallest coordinate distance, then trace along the two polylines/polygons to compute the average angle distance (in degrees). A prediction is considered a true positive only when satisfying both the Chamfer and angle distance thresholds. We augment the three-level thresholds to  $\{(0.5m, 5^\circ), (1.0m, 10^\circ), (1.5m, 15^\circ)\}$ . Results are reported in both the original and augmented metrics. We provide complete implementation details and visualizations about the augmented metric in the Appendix C.
Competing approaches: We compare our method with VectorMapNet [25] and MapTR [23]. VectorMapNet [25] employs an auto-regressive transformer decoder to generate the vertex of polygon/polyline one by one, conditioned on the image inputs. MapTR [23] is our proposal generator and is a DETR-style hierarchical detection method that directly estimates the polygon/polyline as a set of vertex sequences. MapTR and RoomFormer are technically very similar but focus on different tasks.
Quantitative & qualitative evaluation: Table 2 presents the quantitative evaluation results. PolyDiffuse does not significantly improve the Chamfer-distance mAP of MapTR, as it does not discover new instances. However, it shows clear advantages in the order-aware angle distance, a metric that is indicative of high-level structural regularities. Figure 5 qualitatively shows that PolyDiffuse accurately reconstructs curves with greater resemblance to the ground truth. We also experimented with the same rough annotations as in the floorplan task, and the quantitative results show that the mAP becomes lower compared to using MapTR as the proposal generator. By analyzing the results with rough annotations more carefully, we found the AP with the lowest threshold becomes much worse, while the one with the highest threshold improves. We believe this is because the guidance networks are trained with "G.T. proposals", where the rough annotations (i.e., fixed radius circles) have a very different curve style when serving as the proposals, thus producing inaccurate initial Gaussians and stepwise guidance. This is a potential limitation of the current method. We provide additional qualitative examples in Appendix D.3.
Running speed analysis: Table 2 also presents our speed-performance tradeoff on the HD map reconstruction task. In our computation environment, the time of image encoding vs. the time of transformer decoder is 4:1 when using MapTR's model architecture. As FPS is an important consideration for online applications, we need to pick the number of denoising steps carefully. Furthermore, since PolyDiffuse is not restricted to a specific task-specific method for the model
Surrounding views
Figure 5: Qualitative results for HD map reconstruction. PolyDiffuse improves the geometry correctness of MapTR (marked in red) and can turn rough annotations into reasonable reconstructions.
architecture and proposal generator, our performance/efficiency can keep improving when better task-specific base methods appear in the future.
We provide key ablation studies with the floorplan reconstruction benchmark. Please also find additional ablation studies in Appendix D.1.
Standard DM vs. GS-DM: To better understand the effectiveness of our GS-DM, we define a standard DM using a vanilla DDPM with images as the condition. The standard DM does not have the guidance network and proposal generator. The reverse process starts with a noise sampled from the standard Gaussian distribution and gradually denoises it into the final reconstruction using the same sampler as GS-DM. It also has the same denoising network architecture as the GS-DM of PolyDiffuse. As shown in Table 3, although standard DM produces reasonable results with 10 sampling steps, it does not outperform the baseline method RoomFormer. Also, the performance of standard DM drops significantly with fewer sampling steps, as the denoising ambiguity hurts the performance more seriously with a larger sampling step size. We also demonstrate in Appendix D.2 that the standard DM can fail in a toy experiment due to denoising ambiguity.
Sampling steps: Unlike the standard DM, Table 1, Table 2, and Table 3 all suggest that PolyDiffuse is robust to fewer sampling steps and can still slightly improve RoomFormer or MapTR with a very small number of sampling steps.
Training epochs: We observed that the model converges slower under the denoising formulation than the detection/regression formulation of previous approaches. In Table 4, PolyDiffuse keeps improving when training with longer schedules, while RoomFormer suffers from overfitting. However, even with the default schedule, PolyDiffuse significantly outperforms the baseline.
Structured reconstruction of polygonal shapes: Structured geometry reconstruction is an active research area in computer vision, focusing on converting raster sensor data into vectorized geometries such as room layouts [8, 13], floorplans [1, 24], planes [11, 37], wireframes [16], etc. Taking advantage of the advancements in deep learning, methods with end-to-end neural architecture [7, 23, 25, 47, 49, 51] gradually dominate the area with superior accuracy and efficiency. In this paper, we formulate the reconstruction of polygonal shapes as a conditioned generation process with diffusion models, one of the essential machinery of the recent exploding generative AI. Our
Table 3: Comparison between standard DM and GS-DM with different sampling steps. Results are reported in F1 scores.
Table 4: Ablation study on the training schedule of PolyDiffuse and RoomFormer. Results are reported in F1 scores.
Method	Steps	Room	Corner	Edge	Method	Epochs	Room	Corner	Edge
RoomFormer	1	96.2	88.2	83.9	RoomFormer	1×	96.2	88.2	83.9
2	12.5	6.5	5.4	2×	96.0	87.9	83.7
Standard DM	5	86.7	79.6	77.0	1×	98.0	90.6	88.6
10	90.4	83.6	81.1	PolyDiffuse	1.5×	97.7	90.6	88.7
2	96.6	88.7	84.3	2×	97.8	90.7	88.9
PolyDiffuse (w/ GS-DM)	5	98.1	90.7	88.6	2.5×	98.4	91.0	89.1
10	98.4	91.0	89.1
method, PolyDiffuse, outperforms the current state of the arts on floorplan [47] and HD map [23] reconstruction while enabling broader practical use cases.
Diffusion-based generative models: Diffusion Models (DM) or score-based generative models [28, 38, 40, 41] have made tremendous progress in the last few years and demonstrated promising performance in content generations [9, 31, 32, 33, 43] and likelihood estimations [19, 20, 28]. In this paper, we explore the potential of DM in the context of structured reconstruction and propose a Guided Set Diffusion Model (GS-DM) by extending the DDPM [14] formulation. GS-DM reconstructs complex polygonal shapes with a guided denoising (reverse) process subject to sensor data and can evaluate its reconstruction likelihood by leveraging the underlying probability flow ODE [41]. Our high-level formulation is relevant to PriorGrad [21], a diffusion model for acoustic data, but has essential differences - GS-DM learns the guidance networks to control the diffusion and direct the denoising in a per-element and stepwise manner, while PriorGrad pre-computes the mean and variance of the target Gaussian from the condition and directly shifts the diffusion/denoising trajectory.
This paper introduces PolyDiffuse, a system designed to reconstruct high-quality polygonal shapes from sensor data via a conditional generation procedure. At the heart of PolyDiffuse is a novel Guided Set Diffusion Model, that controls the noise injection in the diffusion process to avoid permutation ambiguity for denoising. PolyDiffuse achieves state-of-the-art performance on two challenging tasks: floorplan reconstruction from a point density image and HD map construction from onboard RGB images. To our knowledge, this paper is the first to demonstrate that Diffusion Models, generally regarded as a generation technique, are also powerful for reconstruction, potentially encouraging the community to further investigate the effectiveness of Diffusion Models in broader domains.
Limitations: PolyDiffuse suffers from two major limitations. First, it does not recover geometry instances (e.g., rooms for floorplan) missing in the initial reconstruction, while the likelihood-based refinement provides a possible solution by incorporating search or external inputs. Second, as shown in the HD mapping results with rough annotations, when the style of initial reconstructions (e.g., fixed radius small circles) is different from the curve style of ground truth, the guidance networks could produce bad initial Gaussians and guidance, leading to inaccurate locations of the final results. Training specific guidance and denoising networks for each type of initial reconstruction is a solution but induces more computation.
Broader impact: The paper benefits applications in architecture, construction, autonomous driving, and potentially human-in-the-loop annotations systems to reduce human workload. However, potential malicious or unintended applications might include military scenarios, for example, reconstructing digital models of important buildings or structures for targeted military operations.
Acknowledgements: This research is partially supported by NSERC Discovery Grants with Accelerator Supplements and DND/NSERC Discovery Grant Supplement, NSERC Alliance Grants, and John R. Evans Leaders Fund (JELF). We thank the Digital Research Alliance of Canada and BC DRI Group for providing computational resources.
[42] Sinisa Stekovic, Mahdi Rad, Friedrich Fraundorfer, and Vincent Lepetit. Montefloor: Extending mcts for reconstructing accurate large-scale floor plans. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021.
[43] Shitao Tang, Fuyang Zhang, Jiacheng Chen, Peng Wang, and Yasutaka Furukawa. Mvdiffusion: Enabling holistic multi-view image generation with correspondence-aware diffusion. ArXiv, abs/2307.01097, 2023.
[44] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems (NeurIPS), 2017.
[45] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 2011.
[46] Yifan Xu, Weijian Xu, David Cheung, and Zhuowen Tu. Line segment detection using transformers without edges. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2021.
[47] Yuanwen Yue, Theodora Kontogianni, Konrad Schindler, and Francis Engelmann. Connecting the dots: Floorplan reconstruction using two-level queries. CVPR, 2023.
[48] Jia Zheng, Junfei Zhang, Jing Li, Rui Tang, Shenghua Gao, and Zihan Zhou. Structured3d: A large photo-realistic dataset for structured 3d modeling. In ECCV, 2020.
[49] Yichao Zhou, Haozhi Qi, Yuexiang Zhai, Qi Sun, Zhili Chen, Li-Yi Wei, and Yi Ma. Learning to reconstruct 3d Manhattan wireframes from a single image. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2019.
[50] Xizhou Zhu, Weijie Su, Lewei Lu, Bin Li, Xiaogang Wang, and Jifeng Dai. Deformable detr: Deformable transformers for end-to-end object detection. arXiv preprint arXiv:2010.04159, 2020.
[51] Chuhang Zou, Alex Colburn, Qi Shan, and Derek Hoiem. Layoutnet: Reconstructing the 3d room layout from a single rgb image. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2051-2059, 2018.
- §A: The derivation of the Guided Set Diffusion Models based on the DDPM [14] formulation.
$\S \mathrm{B}$  Additional implementation details, including:
- Architectures and implementation details of the guidance networks.
- Architectures and implementation details of the denoising networks.
- Implementation details of the permutation loss  $L_{\mathrm{perm}}(\phi)$  for the guidance training.
- Implementation details of the diffusion models framework.
- §C: More details and analyses on the augmented matching criterion for the mean AP metric to evaluate HD map reconstruction.
$\S \mathrm{D}$  :Extra experimental results including:
- Additional ablation studies on the design choices of the denoising networks.
- Additional qualitative results for floorplan and HD map reconstruction.
A Derivation of Guided Set Diffusion Models (GS-DM)
In this section, we present more details of the derivation of the Guided Set Diffusion Models (GS-DM) proposed in §3 of the main paper.
A.1 Forward process
We derive Eq. 7, 8, and 9 in Sec. 3 of the main paper by induction. The trivial case of  $t = 1$  can be easily verified. Let  $\mathbf{x}_{t - 1}^i = \sqrt{\bar{\alpha}_{t - 1}}\mathbf{x}_0^i + \bar{\mu}_\phi(\mathbf{x}_0, t - 1, i) + \bar{\sigma}_\phi(\mathbf{x}_0, t - 1, i)\epsilon_{t - 1}^i$  for  $t > 1$  and  $\epsilon_{t - 1}^i \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ . By definition, we have
$$
\begin{array}{l} \mathbf {x} _ {t} ^ {i} = \sqrt {\alpha_ {t}} \mathbf {x} _ {t - 1} ^ {i} + \boldsymbol {\mu} _ {\phi} (\mathbf {x} _ {0}, t, i) + \sqrt {1 - \alpha_ {t}} \sigma_ {\phi} (\mathbf {x} _ {0}, t, i) \boldsymbol {\epsilon} _ {t} ^ {i} \\ = \sqrt {\alpha_ {t}} (\sqrt {\bar {\alpha} _ {t - 1}} \mathbf {x} _ {0} ^ {i} + \bar {\boldsymbol {\mu}} _ {\phi} (\mathbf {x} _ {0}, t - 1, i) + \bar {\sigma} _ {\phi} (\mathbf {x} _ {0}, t - 1, i) \boldsymbol {\epsilon} _ {t - 1} ^ {i}) + \\ \boldsymbol {\mu} _ {\phi} \left(\mathbf {x} _ {0}, t, i\right) + \sqrt {1 - \alpha_ {t}} \boldsymbol {\sigma} _ {\phi} \left(\mathbf {x} _ {0}, t, i\right) \boldsymbol {\epsilon} _ {t} ^ {i} \tag {15} \\ = \sqrt {\bar {\alpha} _ {t}} \mathbf {x} _ {0} ^ {i} + \sqrt {\alpha_ {t}} \bar {\boldsymbol {\mu}} _ {\phi} (\mathbf {x} _ {0}, t - 1, i) + \boldsymbol {\mu} _ {\phi} (\mathbf {x} _ {0}, t, i) + \\ \sqrt {\alpha_ {t}} \bar {\sigma} _ {\phi} (\mathbf {x} _ {0}, t - 1, i) \boldsymbol {\epsilon} _ {t} ^ {i} + \sqrt {1 - \alpha_ {t}} \boldsymbol {\sigma} _ {\phi} (\mathbf {x} _ {0}, t, i) \boldsymbol {\epsilon} _ {t} ^ {i} \\ \end{array}
$$
where  $\epsilon_t^i$  is a standard Gaussian variable independent of  $\epsilon_{t - 1}^{i}$ . Since  $\epsilon_{t - 1}^{i}$  and  $\epsilon_{t}^{i}$  are independent,  $\sqrt{\alpha_t}\bar{\sigma}_{\phi}(\mathbf{x}_0,t - 1,i)\epsilon_{t - 1}^i +\sqrt{1 - \alpha_t}\pmb {\sigma}_{\phi}(\mathbf{x}_0,t,i)\pmb {\epsilon}_t^i$  is a Gaussian random variable of zero mean and a variance of  $\alpha_{t}\bar{\sigma}_{\phi}^{2}(\mathbf{x}_{0},t - 1,i) + (1 - \alpha_{t})\sigma_{\phi}^{2}(\mathbf{x}_{0},t,i)$ . Thus we have
$$
\bar {\boldsymbol {\mu}} _ {\phi} \left(\mathbf {x} _ {0}, t, i\right) := \sqrt {\alpha_ {t}} \bar {\boldsymbol {\mu}} _ {\phi} \left(\mathbf {x} _ {0}, t - 1, i\right) + \boldsymbol {\mu} _ {\phi} \left(\mathbf {x} _ {0}, t, i\right) \text {a n d} \tag {16}
$$