SlowGuess/NeurIPS_2023_OCR huggingface dataset:
PuzzleFusion: Unleashing the Power of Diffusion Models for Spatial Puzzle Solving
Sepidehsadat Hosseini, Mohammad Amin Shabani, Saghar Irandoust*, Yasutaka Furukawa
{sepidh, mshabani, sirandou, furukawa} @sfu.ca
This paper presents an end-to-end neural architecture based on Diffusion Models for spatial puzzle solving, particularly jigsaw puzzle and room arrangement tasks. In the latter task, for instance, the proposed system takes a set of room layouts as polygonal curves in the top-down view and aligns the room layout pieces by estimating their 2D translations and rotations, akin to solving the jigsaw puzzle of room layouts. A surprising discovery of the paper is that the simple use of a Diffusion Model effectively solves these challenging spatial puzzle tasks as a conditional generation process. To enable learning of an end-to-end neural system, the paper introduces new datasets with ground-truth arrangements: 1) 2D Voronoi jigsaw dataset, a synthetic one where pieces are generated by Voronoi diagram of 2D pointset; and 2) MagicPlan dataset, a real one offered by MagicPlan from its production pipeline, where pieces are room layouts constructed by augmented reality App by real-estate consumers. The qualitative and quantitative evaluations demonstrate that our approach outperforms the competing methods by significant margins in all the tasks. We have provided code and data here.
Spatial puzzle solving demands meticulous reasoning and arrangement of elements within a given space. The classic example of jigsaw puzzles, which many of us have enjoyed as a recreational activity, showcases the challenge and satisfaction of such tasks. Applications of jigsaw puzzles extend beyond mere entertainment, encompassing areas such as the restoration of shattered 2D artwork and documents Das et al. (2017); McBride and Kimia (2003), image stitching in computer vision Hammoudeh and Pollett (2017), and even DNA sequence assembly in genomics Pop (2009); Marande and Burger (2007). In real estate, room layout arrangement has emerged as a compelling spatial puzzle task, where consumers leverage mobile devices to scan individual rooms that are to be arranged into a floorplan.
Spatial puzzle solving poses a considerable challenge even for humans, making it an engaging mental exercise. Addressing this challenge for computational approaches is even more demanding, despite the emergence of deep learning. Current state-of-the-art techniques typically enumerate pairs of aligned pieces, evaluate their compatibility by learning, and employ optimization or search methods to identify the most likely global arrangement Harel and Ben-Shahar (2021); Hoff and Olver (2014a); Shih and Lu (2018). However, these approaches struggle to scale as the complexity of global arrangement increases exponentially with the number of pieces.
This paper makes a breakthrough in spatial puzzle solving with an end-to-end neural architecture based on Diffusion Models Ho et al. (2020), named PuzzleFusion. The surprising discovery of this paper is that a Diffusion Model, typically regarded as a powerful generative model, is an effective
*Work done while being student at SFU
Figure 1: Room layout arrangement is the task of taking a set of room layouts and their corresponding room types as the input and predicting the position and the orientation of each room. The biggest discovery and surprise of this paper is that conditional generation by a Diffusion Model solves this challenging problem.
spatial puzzle solver. PuzzleFusion formulates the spatial puzzle as a conditional generation process where the piece information is the condition.
Concretely, each puzzle piece is a polygonal shape, represented by a sequence of 2D corner coordinates. A piece is associated with task-dependent properties, such as a texture image for pictorial jigsaw puzzles or the room types and door locations for Room Layout Arrangement. The output is a 2D position and an orientation for each piece. The forward diffusion process adds noise to a feature vector initialized with the ground truth position and orientation. The reverse denoising process learns to infer position and orientation subject to the piece shape and property information as a condition.
Qualitative and quantitative evaluations over the three tasks (Cross-cut jigsaw, Voronoi jigsaw, and room layout arrangement) demonstrate that PuzzleFusion outperforms the current state-of-the-art methods by significant margins. The paper also makes dataset contributions by introducing large-scale datasets with ground-truth labels for Voronoi jigsaw and room layout arrangement tasks. Specifically, the room layout arrangement dataset comes from a production pipeline by MagicPlan (https://www.magicplan.app/) consisting of room-loads and floorplans for 98,780 houses, which we obtained permission to share with the research community.
In summary, this paper makes the following three key contributions: 1) An end-to-end neural architecture based on Diffusion Models for spatial puzzle solving; 2) State-of-the-art performance across three spatial puzzle tasks in terms of accuracy and speed; and 3) The new spatial puzzle datasets including room-loads and floorplans for 98,780 houses from a production pipeline. We will make all our code and data public.
Spatial puzzle solving is closely related to Structure from Motion (SfM), pose estimation, arrangement learning, and more. The section discusses the related techniques.
Feature matching has been successful for the SfM problem Snavely et al. (2006); Li et al. (2020); Lin et al. (2016). The rise of deep neural networks enables more robust feature matching by learning Yi et al. (2018); Sun et al. (2021); Sarlin et al. (2020). However, these techniques require visual overlaps. Our task has little to no visual overlaps between adjacent images.
Geometry inference estimates a relative pose between images or partial scans with minimal overlaps by registering inferred or hallucinated geometry. A popular approach learns the priors of room shapes and alternates pairwise alignment and scene completion Lin et al. (2019); Yang et al. (2019). Yang et al. (2020) combines global relative pose estimation and local pose refinement
with panoramas. These techniques learn priors of a single room, while our approach learns the arrangements of multiple rooms on a house scale.
Arrangement learning is the current state-of-the-art for indoor room layout arrangement. An early work uses windows to align indoor and outdoor reconstructions Cohen et al. (2016). Shabani et al. Shabani et al. (2021) use doors to enumerate room arrangements and learn to score each candidate. Their approach is exponential in the number of rooms with many heuristics. Lambert et al. Lambert et al. (2022a) uses doors, windows, and openings to create room alignment hypotheses. They utilize depth maps to create top-down views and learn to verify the correctness, improving a run-time from exponential to polynomial. Our end-to-end approach does not enumerate arrangement candidates and makes significant performance improvements. Lastly, an annotated site map and SfM reconstructions are aligned to solve a challenging structure from motion problem MartinBrualla et al. (2014); Hosseini and Furukawa (2022), which they coined as a "3D jigsaw puzzle".
Puzzle Solving has been an engaging research area for a long time Freeman and Garder (1964); Radack and Badler (1982); Markaki and Panagiotakis (2022), ranging from pictorial puzzles with image information Shih and Lu (2018); Le and Li (2019); Toler-Franklin et al. (2010) to apictorial puzzles with only geometry information Goldberg et al. (2002); Harel and Ben-Shahar (2020); Hoff and Olver (2014b, 2013); Harel and Ben-Shahar (2021). In previous studies, heuristic-based methods, such as edge and color matching Wolfson et al. (1988); Nielsen et al. (2008); Chung et al. (1998), have been predominantly used for solving both methods. Recently, deep learning-based methods seek to learn high-level pictorial features Noroozi and Favaro (2016); Li et al. (2021a). However, they suffer from poor geometric reasoning and are limited to simple square puzzles. Consequently, researchers combine learning-based models with heuristics to handle more complex puzzles Le and Li (2019). Nonetheless, these methods still require explicit pairwise comparison of pieces and lack data-driven high-level reasoning. Our method overcomes these limitations through an innovative end-to-end use of Diffusion Models which are usually regarded as powerful generative models.
Diffusion models (DMs) are emerging generative models, which slowly corrupt a sample by adding noise Ho et al. (2020); Dhariwal and Nichol (2021); Nichol and Dhariwal (2021a); Shabani et al. (2022), learn to invert the process, and generate a diverse set of samples from noise signals. DMs have established SOTA performances in numerous tasks such as image colorization/inpainting Song et al. (2020); Nichol et al. (2021), image to image translation Sasaki et al. (2021); Zhao et al. (2022), text to image Ramesh et al. (2022), super-resolution Rombach et al. (2021); Saharia et al. (2021); Li et al. (2021b), image and semantic editing Avrahami et al. (2022); Meng et al. (2022), and denoising Kawar et al. (2022). Recent works use DMs as representation learners for discriminative tasks such as image segmentation Baranchuk et al. (2021); Wolleb et al. (2021). Diffusion inspired models have been used for human pose estimation Shan et al. (2023); Gong et al. (2023), and object placement Wei et al. (2023), however, both tasks do not require sophisticated shape comparisons. While in our case a more advanced and intricate approach is necessary to capture and exploit all details of the shapes.
Figure 2: Spatial puzzle datasets for Voronoi and Cross-cut puzzle solving (Left) and room layout arrangement (Right). We consider both pictorial and apictorial versions of the Cross-cut jigsaw puzzle. Note that colors in puzzle solving problems are random and just indicate different pieces.
3 Spatial Puzzle Solving Tasks
Spatial puzzle solving involves arranging puzzle pieces, each with its geometry and optional features such as images or categories. In real scenarios, puzzle pieces are affected by erosion, duplication, or loss. This paper investigates three distinct puzzle-solving tasks.
- Cross-cut Jigsaw Puzzle (CJP) generates pieces by making random straight cuts through a larger polygonal shape Harel and Ben-Shahar (2021). Puzzle pieces are convex polygons, each with an arbitrary number of neighboring pieces. The sum of the two adjacent angles at any corner equals  $180^{\circ}$ . The application of the technique could be the restoration of shattered artwork.
- Voronoi Jigsaw Puzzle (VJP) generates pieces by randomly sampling points within a predetermined bounding box and extracting the cells of the Voronoi diagram as the pieces. With the lack of the "180° constraint", VJP is significantly more challenging, where no effective solution has been presented in the literature to our knowledge. Voronoi diagrams play a role in biological systems like cell arrangements Bock et al. (2010), potentially useful in studying natural phenomena.
- Room Layout Arrangement (RLA) determines the room arrangement and the corresponding floor-plan, offering a key application in real estate. Contrary to the prior tasks, a corner may align along the edge of another piece, expanding the solution space to be explored. Pieces come from the room layout estimation algorithms Shabani et al. (2021); Lambert et al. (2022a) or interactive augmented reality apps used by consumers.
Task input/output: The input is a set of  $N$  polygons (puzzle pieces in CJP/VJP and room layouts in RLA), each of which is a sequence of corner coordinates forming a 1D polygonal loop. In RLA, a door piece is also given as a line segment with two corners, and a room piece is associated with a room type as a 20D one-hot vector. For a pictorial version of CJP, a piece is associated with a 128D image feature vector obtained from a pretrained auto-encoder. For simplicity, we mix room-corners and door-corners, and use  $C_i^r$  to denote the 2D coordinate of the  $i$ th corner in the  $r$ th polygon.  $\mathcal{T}^r$  denotes the image feature or the room type vector. Please refer to supplementary for details. The output is the position of the piece/room center and the rotation around it (an angle between 0 and  $2\pi$  for CJP/VJP and a 4-fold Manhattan rotation for RLA). The center is the average of the corners.
Metrics: For CJP and VJP, we adopt metrics used in previous work Harel and Ben-Shahar (2020), namely Overlap, Precision, and Recall. The overlap score is the average IoU of pieces with the ground truth. Precision and Recall are on the connectivity of neighboring pieces. For RLA, we consider two metrics. The first metric is the Mean Positional Error in pixels (MPE) over the rooms Shabani et al. (2021). The second metric evaluates the correctness of the room connectivity in the reconstruction. We borrow a Graph Edit Distance (GED) by Nauata et al. Nauata et al. (2021), which counts the number of user edits necessary to fix the connectivity graph. We declare that two rooms are connected if the door centers are within 5 pixels from the two rooms. Note that we have designed task-specific metrics, respecting the methodologies of existing literature while providing a thorough evaluation of our performance.
4 Spatial Puzzle Solving as Conditional Generation
Our idea is simple, using a Diffusion Model to "conditionally generate" the correct arrangement, where the input is the center and rotation of room layouts/puzzle pieces and their types and shapes are the conditions. This section explains the forward and the reverse processes.
4.1 Forward process
The forward process adds a Gaussian noise to an arrangement. A compact representation would be per-polygon positions and rotations. Instead, we will use a redundant representation, where a center position and a rotation are stored at each corner.
$^{2}$ Shabani et al. Shabani et al. (2021) used a permissive metric (the availability of the "correct" solution in the k results with a certain error-tolerance) as the task was challenging. We make great improvements and use a standard metric.
Figure 3: Illustration of our diffusion model architecture employed for the RLA task. Given the arrangement estimation  $x_{t} = \{x_{i,t}^{r}\}$ , the reverse process infers the noise  $\{\delta_{i,t}^{r}\}$  and recovers  $x_{t-1} = \{x_{i,t-1}^{r}\}$ , while injecting the original room shapes  $\{C_{i}^{r}\}$  and types  $\{\mathcal{T}^r\}$  as the condition. Each room corner holds the room position and rotation estimation. The reverse process starts from  $x_{T}$  and denoises towards  $x_{0}$ .
There are a few reasons. Our reverse process is based on a Transformer architecture where each position/rotation estimation becomes a node. Our approach 1) enriches the capacity of the arrangement representation (also an adaptive capacity, where a complex polygon with more corners is given more capacity); 2) allows direct communications between corners/doors for which we will have a specific loss; and 3) makes it straightforward to combine with the condition (i.e., original corner coordinates and room types).
Concretely, we use  $x_{i,t}^{r}$  to denote the position/rotation of the  $r$ th room/piece stored at the  $i$ th corner at time  $t$  of the diffusion process, where  $t$  varies from 0 to 1,000 in our experiments:
$$
x _ {i, t} ^ {r} = \left(p _ {i, t} ^ {r}, o _ {i, t} ^ {r}\right). \tag {1}
$$
$p_{i,t}^{r}$  and  $o_{i,t}^{r}$  denote the polygon-center position (a 2D vector) and the rotation. We consider rotation as a 2D vector obtained from rotation matrix including  $\cos(o_{i,t}^{r})$  and  $-\sin(o_{i,t}^{r})$ . The forward process adds a noise by sampling  $\delta_{i,t}^{r} \in \mathcal{N}(\mathbf{0}, \mathbf{I})$  with a standard cosine noise scheduling with variance  $(1 - \alpha_{t})$  Nichol and Dhariwal (2021b):
$$
x _ {i, t} ^ {r} = \sqrt {\bar {\alpha} _ {t}} x _ {i, 0} ^ {r} + \sqrt {1 - \bar {\alpha} _ {t}} \delta_ {i, t} ^ {r}, \quad \bar {\alpha} _ {t} = \frac {f (t)}{f (0)}, \quad f (t) = \cos \left(\frac {t / T + 0 . 0 0 8}{1 + 0 . 0 0 8} \cdot \frac {\pi}{2}\right) ^ {2} \tag {2}
$$
4.2 Reverse process
Figure 3 illustrates our reverse process for room layout arrangement, which takes the arrangement  $\{x_{i,t}^r\}$  at time  $t$  and infers the noise  $\delta_{i,t}^{\tilde{r}}$  under the condition of the original room shapes  $\{C_i^r\}$  (i.e., a corner position with respect to the room center) and the room types as a one-hot vector  $\{\mathcal{T}^r\}$ .
Feature embedding: The reverse process is based on a Transformer architecture where every corner is a node. We initialize its feature embedding as
$$
\hat {x} _ {i, t} ^ {r} \leftarrow \operatorname {L i n} \left(x _ {i, t} ^ {r}\right) + \operatorname {L i n} \left(\left[ C _ {i} ^ {r}, r, i \right]\right) + \operatorname {M L P} (t) + \operatorname {L i n} \left(\mathcal {T} ^ {r}\right). \tag {3}
$$
The first term uses a linear layer to convert a 4D vector (i.e., 2 for the room center coordinate and 2 for the rotation vector) to a 256D embedding vector. The second term also uses a linear layer to convert a condition vector (i.e., 2 for the corner coordinate respecting to the center of the polygon
in addition to 32 for polygon index which shows which polygon each corner belongs  $r$  and 32 to show the corner index in each polygon  $i$ ). The third term uses a 2-layer MLP to convert a time step  $t$  to a 256D vector. The fourth term is extra information such as room/door type (20 types) for RLA which we extract it by applying a linear layer on 20D one-hot vector or 128D for image features in pictorial CJP)
Attention modules: Feature embeddings  $\{\hat{x}_{i,t}^r\}$  go through six blocks of self-attention modules that have two different attention mechanisms: Polygon Self Attention (P-SA) and Global Self Attention (G-SA). P-SA limits pairwise interactions between corners in each polygon. P-SA is akin to a sparse self attention family Guo et al. (2019); Child et al. (2019); Li et al. (2019), which helps to generate consistent positions and rotations at different corners of a polygon. G-SA is a standard self-attention between all corners of a puzzle or a house. After the attention blocks, a linear layer converts 256D embedding back to a 4D representation  $\tilde{\delta}_{i,t}^r$ , which is used for the following denoising formula Ho et al. (2020):
$$
x _ {i, t - 1} ^ {r} = \frac {1}{\sqrt {\alpha_ {t}}} \left(x _ {i, t} ^ {r} - \frac {1 - \alpha_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \tilde {\delta} _ {i, t} ^ {r}\right) + \sqrt {1 - \alpha_ {t}} z. \tag {4}
$$
$z \sim \mathcal{N}(0, I)$  for  $t > 1$  and otherwise 0. For the final result at time  $t = 0$ , we take the average polygon center position and the polygon rotation.
Loss functions: There are two loss functions. We first follow Ho et al. (2020); Dhariwal and Nichol (2021) and use a standard noise regression loss on  $\delta$ :
$$
L _ {\text {s i m p l e}} = E _ {t, x _ {i, 0} ^ {r}, \delta_ {i, t} ^ {r}} \left[ \left\| \delta_ {i, t} ^ {r} - \tilde {\delta} _ {i, t} ^ {r} \right\| ^ {2} \right]. \tag {5}
$$
To enhance the quality of supervision, we propose a "matching" loss, specifically aimed at the vertices where incident edges meet. These vertices are shared by two polygons. We denote the indices of these two polygons as  $r_1$  and  $r_2$ , then the corresponding vertex indices within  $r_1$  and  $r_2$  as  $i_1$  and  $i_2$ , respectively. The corresponding vertices must be at the same position, and as such, we calculate the loss as the Euclidean distance between their coordinates:
$$
L _ {\text {m a t c h}} = E _ {t, x _ {i, 0} ^ {r}, \delta_ {i, t} ^ {r}} \left[ \left\| \hat {C} _ {i _ {1}} ^ {r _ {1}} - \hat {C} _ {i _ {2}} ^ {r _ {2}} \right\| ^ {2} \right], \tag {6}
$$
$$
\hat {C} _ {i} ^ {r} = \tilde {p} _ {i, 0} ^ {r} + R _ {\delta_ {i, 0} ^ {r}} C _ {i} ^ {r}, \tag {7}
$$
$$
\left(\tilde {p} _ {i, 0} ^ {r}, \tilde {\sigma} _ {i, 0} ^ {r}\right) = \tilde {x} _ {i, 0} ^ {r} = \left(x _ {i, t} ^ {r} - \sqrt {1 - \bar {\alpha} _ {t}} \tilde {\delta} _ {i, t} ^ {r}\right) / \sqrt {\bar {\alpha} _ {t}}. \tag {8}
$$
$R_{\tilde{\sigma}_{i,0}^r}$  denotes the rotation matrix corresponding to  $\tilde{\sigma}_{i,0}^r$ . For the room layout arrangement task, we add the loss only to doors without walls, which yields superior results in our experiments. The total loss is defined as the sum of the above loss functions,  $L_{\mathrm{total}} = L_{\mathrm{simple}} + L_{\mathrm{match}}$ . In practice, we found that adding  $L_{\mathrm{match}}$  only for  $t < 500$  results in better performance.
5 Experimental Results and Discussions
We have implemented the system with PyTorch Paszke et al. (2019), using a workstation with a 3.70GHz Intel i9-10900X CPU (20 cores) and two NVIDIA RTX A6000 GPUs. We use the AdamW Loshchilov and Hutter (2017); Kingma and Ba (2014) optimizer with  $\beta_{1} = 0.9$ ,  $\beta_{2} = 0.999$ , weight decay equal to 0.05, and the batch size of 512. The learning rate is initialized to 0.0005. It takes roughly 24 hours to train a model and 3 seconds to estimate the arrangement for one sample.
5.1 Preprocessing and Datasets
We have carefully pre-processed and prepared datasets for fair evaluation, whose details are referred to supplementary. We here provide summary points for the three tasks (See Fig. 2).
Cross-cut Jigsaw Puzzle (CJP): We have used code provided by Harel et al. Harel and Ben-Shahar (2021) to generate  $100\mathrm{k} / 1\mathrm{k}$  puzzles for training/testing. The code generates a convex polygon and
cuts it by 3 to 5 lines. For the pictorial version, we have used images from COCO 2017 dataset Lin et al. (2014), that is, randomly selecting a training/testing image of the dataset for each training/testing CJP sample. Following Harel and Ben-Shahar (2021), to simulate real-world scenario, we perturb corner coordinates with noise with three different levels to create datasets under three different noise levels with thresholds equal to 0, 2, and 4.
Voronoi Jigsaw Puzzle (VJP): We have generated  $200\mathrm{k} / 1\mathrm{k}$  puzzles for training/testing by randomly sampling 3 to 15 points inside a random rectangle and extracting their Voronoi cells as the pieces. We have created datasets with three different levels of noise. For each corner, we added a random number from a Gaussian distribution with  $\sigma^2$  equal to 0, 2, and 4 for noise levels 0, 1, and 2, respectively.
Room Layout Arrangement (RLA): MagicPlan (https://www.magicplan.app), a mobile software company for real estate and construction, agrees to share production data with us and the research community, where the paper introduces the MagicPlan dataset, containing room shapes and their ground-truth arrangement for 98,780 single-story houses/apartments. We split the data into 93,780/5,000 training/testing samples. MagicPlan software reconstructs room shapes one by one by asking users to click room corners through an augmented reality app. Room shapes are Manhattan-rectified as an enforcement of the app, then manually arranged to form a floorplan, which we seek to automate. Each room is associated with a room type. The number of rooms (resp. corners) in a house ranges from 3 to 10 (resp. 12 to 182). We also use RPLAN Wu et al. (2019) for evaluation, containing 60k floorplans. We divide into  $55\mathrm{k} / 5\mathrm{k}$  training/testing, where the number of rooms (resp. corners) in a house ranges from 3 to 8 (resp. 14 to 98). Note that RICOH dataset Sahara et al. (2021) and ZIND dataset Lambert et al. (2022b) are too small for network training and are not used in our experiments.
5.2 Competing methods