SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jiaxin Lu * Yifan Sun * Qixing Huang
{lujiaxin, yifansun12}@utexas.edu huangqx@cs.utexas.edu
Automated assembly of 3D fractures is essential in orthopedics, archaeology, and our daily life. This paper presents Jigsaw, a novel framework for assembling physically broken 3D objects from multiple pieces. Our approach leverages hierarchical features of global and local geometry to match and align the fracture surfaces. Our framework consists of four components: (1) front-end point feature extractor with attention layers, (2) surface segmentation to separate fracture and original parts, (3) multi-parts matching to find correspondences among fracture surface points, and (4) robust global alignment to recover the global poses of the pieces. We show how to jointly learn segmentation and matching and seamlessly integrate feature matching and rigidity constraints. We evaluate Jigsaw on the Breaking Bad dataset and achieve superior performance compared to state-of-the-art methods. Our method also generalizes well to diverse fracture modes, objects, and unseen instances. To the best of our knowledge, this is the first learning-based method designed specifically for 3D fracture assembly over multiple pieces. Our code is available at https://jiaxin-lu.github.io/Jigsaw/.
The task of assembling 3D fractures has extensive applications across numerous fields. For instance, orthopedic doctors need to realign dislocated bone fragments, and subsequently create bone plates and screws to heal compound fractures. Archaeologists, on the other hand, need to recreate the original shape and functionality of unearthed artifacts by assembling the fractures. These procedures demand significant expertise, are prone to errors, and can be tedious. Even in the context of daily life, furniture assembly can be a challenging and exhausting task that requires an understanding of mechanical structures and component matching. In the past two decades, many efforts attempted to address the challenge of automatic assembly. Traditional methods apply hand-crafted geometric features to detect the fracture surfaces and optimize pairwise matching among these surfaces [1, 2]. Recently the availability of large scale 3D datasets [3, 4, 5, 6] have boosted learning based frameworks for solving 3D assembly tasks. Semantic-aware methods [7, 8, 9, 10, 11, 12] target at assembly from semantically segmented parts and predict semantic labels as matching priors. Geometry based methods [13, 14, 15] leverage fracture shapes and continuity of textures in the procedure of piece matching.
However, in the context of fracture assembly for restoring broken objects, there is no guarantee that individual pieces will retain semantic meanings. Texture information may also be either non-accessible or lacking, such as on glass bottles or worn artifacts. This calls for a general 3D assembly solver that utilizes hierarchical features of both the global piece surfaces and local geometry shapes of fractures. In this work, we introduce a Joint Learning of Segmentation and Alignment Framework
(Jigsaw), for assembling objects damaged or shattered due to physical impact or force. Given a set of fractured pieces represented as point clouds without texture, our method recovers the global pose of each piece to restore the underlying object.
Our Jigsaw framework consists of four parts: (1) A front-end feature extractor with self-attention and cross-attention layers for local geometric feature extraction. (2) Categorize the surface of each fractured piece into two segments: the indiscernible fracture surface and the visible original surface. (3) A novel formulation for multi-piece assembly that learns a bipartite matching among points on the fracture surface from all pieces. (4) Recover pairwise pose based on the learned correspondences and perform robust global alignment to compute the global poses of all pieces. Key features of Jigsaw are that 1) it learns two correlated tasks, i.e., segmentation and matching among all involving pieces, jointly, and 2) it seamlessly integrates feature matching and the rigidity constraint among consistent features. We evaluate the effectiveness of our framework on Breaking Bad [4]. Experimental results demonstrate the effectiveness of our method in surface segmentation, fracture point matching, and significantly outperforming state-of-the-art methods[8, 16]. The main contribution includes the following:
- We propose Jigsaw, a novel joint learning framework tailored for multi-part fracture assembly. Our approach embodies an attention-based feature extractor network to accurately capture local geometry features of each point. Additionally, we introduce a primal-dual descriptor that effectively captures viewpoint-dependent characteristics for surface matching.
- Jigsaw incorporates fracture point segmentation to capture intrinsic features, employs a novel multi-part matching formulation to establish automatic piece positioning within one object, and utilizes global alignment for accurate global pose alignment.
- Experimental evaluation on the multi-part assembly Breaking Bad dataset demonstrates the superior performance of Jigsaw compared to baseline models, showcasing its strong generalization ability to unseen objects. We further highlight the limitations of baseline models that rely on global features, which we argue are too abstract for this task and lack generalizability.
- To the best of our knowledge, Jigsaw is the first learning-based method specifically designed for the assembly of multiple pieces from physically broken 3D objects.
Feature matching. Early works apply hand-crafted features over fracture surfaces for fracture matching between different pieces [17, 1, 2, 18]. These features are also used to identify the fractured part from the entire surface [17, 1, 2]. However, hand-crafted features are in lack of robustness for assembly tasks over large datasets due to different materials and fracture patterns of objects. In recent years, deep learning methods have gained significant traction in matching problems. Methods utilizing CNN, GNN, and attentions have found successful applications in various domains, such as image registration [19, 20] and graph or multi-graph matching on images [21, 22, 23, 24, 25]. While their efficacy in these scenarios is noteworthy, the majority of these methods predominantly focus on simpler settings and have not adequately addressed the challenges in multi-part 3D fracture assembly.
Part assembly. Semantic-aware learning methods have been highly successful in the task of part assembly. [9, 11] are designed for assembling specific CAD mechanics. For categorical everyday objects, [26, 10, 27, 12, 28] generate the missing parts based on the accumulated shape prior to completing the entire object, which can result in shape distortion from the input parts. [8, 7] apply graph learning to predict part labels and assembly orders. All of these methods require the input objects decomposed in a semantically consistent way and need specific training for each category of objects. Fracture assembly poses unique challenges due to the variety of objects and the lack of semantic meanings associated with individual pieces, adding to the complexity of the task.
Geometry based learning methods. Recent approaches have aimed to capture geometry information using deep features for piece matching. [13, 15] combine local geometry with textures for feature generation, but may suffer when texture information is not accessible, as in most point cloud representations. [14] apply a transformer for local shape encoding and an adversarial network to help generate a plausible assembly of two pieces. However, for multiple fracture assembly, objects can break into 5 pieces or more, with the largest piece smaller than half of the original object. Under this setting, training an adversarial network to evaluate the quality of assembly becomes less effective.
Figure 1: Overall pipeline for Jigsaw (mesh used only for visualization). The method consists of four parts: front-end feature extractor, surface segmentation, multi-part matching, and global fracture alignment. In the front-end feature extraction, we use a multi-scale grouping PointNet++ [31], and one self-attention layer followed by one cross-attention layer to extract features for each point. Surface segmentation is used to locate all fracture points in one broken object, and multi-part matching finds correspondences of fracture points among multiple pieces. The matching results will be used for pairwise pose recovery and joint alignment to retrieve an assembled object. More detail about each part will be discussed in Section 3.
Low-overlap 3D registration. Another relevant research field is low-overlap 3D registration. Recent works have shown the potential of data-driven methods in registration tasks with approximately  $30\%$  overlap [29, 16, 30]. [16, 30] learn a classifier to determine the overlapped sections and predict inter-piece point matching inside the sections, while [29] using object semantic to guide registration. The fracture assembly task can be viewed as an even more extreme case of registration, where the overlap between pieces can fall below  $4\%$ .
3 Joint learning framework for 3D fracture assembly
Given a set of fractured pieces  $\mathcal{P} = \{P_1, P_2, \dots, P_n\}$  represented as point clouds uniformly distributed on the surface, our goal is to recover the 6-DoF pose  $\{T_1, T_2, \dots, T_n\}$  in  $SE(3)$  for each piece and restore the underlying object  $O = T_1(P_1) \cup T_2(P_2) \cup \dots \cup T_n(P_n)$ , where  $T_i(\cdot), 1 \leq i \leq n$  is the operator to recover  $P_i$  to its original position by transformation  $T_i$ . In the fracture assembly setting, the object  $O$  is a rigid body, and the pieces are the result of physical cracking or breakage without any deformation. Also, no pieces are lost, which ensures that the restored pieces can approximately reconstruct the entire object  $O$ .
To handle this challenge, we propose Jigsaw, a learning-based framework that jointly optimizes surface segmentation and fracture matching among all pieces of the object. The entire assembly network consists of a segmentation module and a matching module that shares a front-end feature extractor. The segmentation module uses only the intrinsic shape information of each piece to separate the fracture surface from the original surface. The matching module exploits the primal dual descriptor of the affinity metric to propagate mutual information among all pieces and establish the matching between fracture points (that is, points on the fracture surface) of different pieces via Sinkhorn [32]. With the fracture point matching predicted by the network, we recover the pairwise transformations and perform global alignment with standard approaches. The complete pipeline of our framework is shown in Fig. 1. Note that this design ensures that segmentation and matching are performed jointly. In the subsequent part of this section, we will first introduce the front-end extractor, then delve into the details of the segmentation module, followed by the multi-part matching module and post-processing for pose recovery.
3.1 Front-end Feature Extractor
In contrast to previous methods [33, 4] that utilize global piece-wise descriptors for assembly, we focus on geometric features. As illustrated in Figure 7 of the PointNet paper [34], a global descriptor remains the same as long as the piece lies between the learned critical point set and the upper-bound
point set. This indicates that global descriptors are coarse and insufficient to represent the intricate geometry of the fracture surface, which is crucial for our task. To address this limitation, we employ a multi-scale grouping PointNet++ [31] as the backbone of our feature extractor to capture local geometry features, denoted as  $\mathbf{f}_p\in \mathbb{R}^D$  from each point  $p\in P_{i}$
Furthermore, We emphasize the significance of relative positional information between points and pieces for identifying the non-smooth fracture surfaces, enforcing rigidity during matching, and accurately placing the pieces. To facilitate the integration of such intra-piece and inter-piece information, we leverage transformer layers. Specifically, we introduce both self-attention and cross-attention layers as tools to reason about the relative information between points.
For the self-attention within a piece, we employ a single point transformer layer [35]. The self-attention mechanism is defined as follows:
$$
\mathbf {f} _ {p} ^ {s} = \sum_ {\mathbf {f} _ {q} \in \mathcal {N} (p)} \operatorname {s o f t m a x} \left(\operatorname {M L P} _ {-} \mathbf {s} \left(\left(W _ {Q} ^ {s} \mathbf {f} _ {p} - W _ {K} ^ {s} \mathbf {f} _ {q} + \mathbf {p} ^ {\mathrm {e n c}}\right)\right)\right) \odot \left(W _ {V} ^ {s} \mathbf {f} _ {q} + \mathbf {p} ^ {\mathrm {e n c}}\right) \tag {1}
$$
Here,  $W_{Q}^{s}, W_{K}^{s}, W_{V}^{s} \in \mathbb{R}^{D \times D}$  are weights for query, key, value in the attention layer,  $\mathbf{p}^{\mathrm{enc}}$  represents the positional encoding for point  $p_i$ , obtained using a Multi-Layer Perceptron (MLP). Unlike the standard dot-product attention layer, it employs vector weights: MLP_s is a mapping function that produces the weight vector,  $\text{softmax}(\cdot)$  normalizes the weights, and  $\odot$  denotes element-wise multiplication.  $\mathcal{N}(p) \in P_i$  represents the neighborhood of point  $p$ , calculated using  $k$ -nearest neighbors for local feature aggregation.
For cross-attention, we use standard multi-head attention with position-wise feed forward over the entire object to facilitate the communication of local features among different pieces. Let  $\mathbf{F}^s\in \mathbb{R}^{N\times D}$  pack all the point features of one object after the self-attention, and  $\mathbf{F}^c\in \mathbb{R}^{N\times D}$  denote the point features produced by cross-attention,
$\mathbf{F}^c = \mathrm{FFN}(W_O^m (\mathrm{cat}(\mathrm{head}_1,\dots ,\mathrm{head}_h))),$  where head  $i =$  Attention  $(W_{Q_i}^m\mathbf{F}^s,W_{K_i}^m\mathbf{F}^s,W_{Q_i}^m\mathbf{F}^s)$
$$
\text {a n d} \operatorname {A t t e n t i o n} (Q, K, V) = \operatorname {s o f t m a x} \left(\frac {Q K ^ {T}}{\sqrt {d _ {k}}}\right) V \tag {2}
$$
Here,  $W_{Q_i}^m, W_{K_i}^m, W_{V_i}^m \in \mathbb{R}^{D \times d_h}, W_O^m \in \mathbb{R}^{hd_h \times D}$  denote the weights for projection. The function  $\mathrm{FFN}(\cdot)$  consists of two linear functions with ReLU activation in between, followed by a LayerNorm as normalization.
3.2 Surface segmentation
The fracture assembly problem can be viewed as a special case of the 3D registration problem, where the overlapping between adjacent pieces only lies in the fracture surface between them. Let  $P_i^f \subset P_i$  be the subset of points on fracture surfaces of  $P_i$  and let  $P_{ij} \subset P_i^f$  denote the subset of points on the fracture surface between  $P_i$  and  $P_j$  for arbitrary two pieces  $P_i, P_j$ . For pieces  $P_i$  and  $P_j$  adjacent in their original pose  $T_i(P_i)$  and  $T_j(p_j)$ , we have  $P_{ij} \neq \emptyset, P_{ji} \neq \emptyset$ . Under the assumption that the point clouds are uniformly distributed, we can approximate the relative pose  $T_{ij}^* = T_j^{-1}T_i$  between  $P_i, P_j$  with
$$
T _ {i j} ^ {\star} = \underset {T _ {i j} \in S E (3)} {\operatorname {a r g m i n}} d \left(P _ {j i}, T _ {i j} \left(P _ {i j}\right)\right), \tag {3}
$$
where  $d(\cdot, \cdot)$  is some distance function between two point clouds. (3) becomes the standard formulation for point cloud registration with perfect prior for the overlapping. Before we can optimize (3) directly, we need to figure out the overlapping part  $P_{ij}$  and  $P_{ji}$  for each pair  $(P_i, P_j)$ . Although  $P_{ij}$  is dependent on both pieces, previous studies [17, 1, 2] have demonstrated that the segmentation of fracture surfaces  $P_i^f$  of  $P_i$  is an intrinsic property that can be inferred from the local shapes of each  $P_i$  individually. However, these methods require a continuous smooth surface to compute hand-crafted shape features for accurate classification. To capture local geometric properties under the discrete point cloud setting where surface normal and curvature are no more available, we introduce our deep surface segmentation module to handle the surface segmentation task.
The surface segmentation module can be viewed as a binary classifier of each point. Let  $c_{p}$  be the indicator function that determines whether a point  $p$  of a point cloud  $P \in \mathcal{P}$  is a fracture point. For each 3D model uniformly sampled as  $N = \sum_{i=1}^{n} N_{i}$  points, the surface segmentation module
takes the position of the points  $P_{i} \in \mathbb{R}^{N_{i} \times 3}$  of a piece  $P_{i} \in \mathcal{P}$  and predicts a confidence score  $c_{i}$  to segment a point  $p \in P_{i}$  into the fracture surface.
Our front-end feature extractor introduced in Section 3.1 has extracted local geometry feature  $\mathbf{f}_p\in \mathbb{R}^D$  for each point. Next, we use two MLP layers to reduce the number of channels to 1, followed by a sigmoid function to predict the confidence  $c_{p}$  for labeling  $p$  as a fracture point. We set the segmentation loss  $L_{\mathrm{seg}}$  to be a negative log-likelihood loss, which is supervised by the ground truth label  $c_{p}$  for every point  $p$  as
$$
\mathcal {L} _ {\mathrm {s e g}} = - \frac {1}{N} \sum_ {p \in O} c _ {p} \log \tilde {c} _ {p} + (1 - c _ {p}) \log (1 - \tilde {c} _ {p}). \tag {4}
$$
The ground truth label  $c_{p}$  for each point  $p \in P_{i}$  is constructed by examining its distance to its nearest neighbor in other parts:
$$
c _ {p} = \left\{ \begin{array}{l l} 1 & D i s (p, \mathbb {N N} (p, \mathcal {P} \backslash P _ {i})) \leq \eta , \\ 0 & \text {o t h e r w i s e .} \end{array} \right. \tag {5}
$$
where  $Dis(\cdot)$  is a distance function and  $\mathbb{N}\mathbb{N}(\cdot, \cdot)$  is to find the nearest neighbour of one point in a point set.
Although the fraction of the fracture surface area may differ between objects, and the positive and negative samples in (4) may be unbalanced, we have found that adding two MLP layers after feature encoding help to yield good performance for the segmentation task. The segmentation module is able to accurately predict fracture points even in cases with a large number of pieces or very small fracture surfaces, and we leave the evaluation details of this module in the appendix A.
3.3 Multi-part Matching
One key aspect of Jigsaw is its ability to combine multiple pieces. Traditional approaches to assembling pieces rely on pairwise matching, which involves assessing the compatibility of two parts, identifying corresponding points within them, and aligning those points. However, this method is prone to cumulative errors, and a single mistake can ruin the entire assembly task. In addition, matching small pieces with limited geometry information solely based on pairwise information is extremely challenging. Consequently, the pairwise information in the multipart fracture assembly dataset may not be comprehensive enough to capture the true matching possibilities between all the pieces involved in the assembly process. Even advanced low-overlap 3D registration methods like PREDATOR [16], which share a similar structure, have failed to handle assembly tasks effectively, as shown in Section 4.2. Thus, alternative approaches are needed to address the limitations of pairwise matching, and we introduce our multipart matching module.
Facilitated by the global fracture points of a segmented object in Section 3.2, we present a global view to match all its pieces simultaneously. We observe that each fracture surface should have one precise match within a broken object, which is a rule that applies specifically to the fracture piece assembly task. Therefore, we learn the matching among all pieces without specifying their pairwise relationship. This allows each fracture point to automatically find a similar counterpart on a different piece solely based on local geometry. Moreover, the correspondences highlight the assembly regions, avoiding the need for costly predictions of fine-grained pairwise overlapping areas.
Primal-dual Descriptor. Intuitively, when an object has been broken into two pieces by physical force, the two pieces should exhibit complementary geometry. Conventional feature representation is used to find similar geometry and, therefore, would match two identical surfaces, which
Figure 2: Illustration of Primal-dual descriptor: The red region and points represent the target fragment piece and the points, respectively. The purple feature (primal descriptor) on the left learns the local convex shape of the piece, while the blue feature (dual descriptor) on the right captures the concavity observed from the outside. The primary and dual descriptors on both sides will be matched in the matching module.
becomes undesirable in the context of fracture assembly. With this concern, we propose the primal-dual descriptor, designed to capture the essence of complementary geometry (see Fig. 2).
Let  $\{\hat{\mathbf{f}}_i^c\}_{i = 1}^{n_c}$  be the selected features of the fracture points extracted from the backbone network. These features encode the local geometry of the fracture surfaces. We apply MLP layers and a normalization layer to  $\hat{\mathbf{f}}_i^c$  and get the primal descriptor  $\hat{\mathbf{f}}_i^p$  and dual descriptor  $\hat{\mathbf{f}}_i^d$ :
$$
\hat {\mathbf {f}} _ {i} ^ {p}, \hat {\mathbf {f}} _ {i} ^ {d} = \operatorname {n o r m} \left(\mathrm {M L P} _ {-} \text {a f f} \left(\hat {\mathbf {f}} _ {i} ^ {c}\right)\right). \tag {6}
$$
The primal-dual descriptor learns surface features that capture the characteristics of a local surface from both directions. It promotes the robustness and reliability of surface matching and avoids learning biased features based on a single viewing direction. We leave more detailed analysis to the appendix B.
Affinity Metric. Let  $\hat{\mathbf{F}}^p, \hat{\mathbf{F}}^d \in \mathbb{R}^{\hat{N} \times d}$  collect all the fracture points features extracted from descriptors, we compute the global affinity matrix  $M \in \mathbb{R}^{\hat{N} \times \hat{N}}$  and the respective doubly-stochastic matrix  $\tilde{X} \in [0,1]^{\hat{N} \times \hat{N}}$  as follows:
$$
M = \exp \left(\frac {\hat {\mathbf {F}} ^ {p} {} ^ {\top} A \hat {\mathbf {F}} ^ {d}}{\tau}\right), \quad \tilde {X} = \operatorname {S i n k h o r n} (M), \tag {7}
$$
Here,  $A \in \mathbb{R}^{d \times d}$  consists of learnable affinity weights, and  $\tau$  denotes the temperature parameter [36]. We then apply the Sinkhorn layer [32], which is a differentiable operation, to obtain a doubly stochastic matrix  $\tilde{X} \in [0,1]^{\hat{N} \times \hat{N}}$ . This matrix represents the soft matching predicted by the multipart matching module.
Matching Loss. Given the ground truth position of each part, we construct the ground truth matching matrix  $X^{gt} \in \{0,1\}^{\hat{N} \times \hat{N}}$ . Here,  $x_{ij}^{gt} = 1$  if and only if the  $j$ -th fracture point is the nearest neighbor of the  $i$ -th fracture point and both points belong to different pieces; otherwise,  $x_{ij}^{gt} = 0$ . To compute the matching loss, we employ the cross entropy between  $\tilde{X}$  and  $X^{gt}$ :
$$
\mathcal {L} _ {\mathrm {m a t}} = - \frac {1}{\hat {N}} \sum_ {1 \leq i j \leq \hat {N}} x _ {i j} ^ {g t} \log \tilde {x} _ {i j} + \left(1 - x _ {i j} ^ {g t}\right) \log \left(1 - \tilde {x} _ {i j}\right) \tag {8}
$$
With the ground truth matrix  $X^{gt}$  generated based on finding the nearest neighbor of each point, the matching loss will also enforce a global rigidity guidance that each node should be matched to its neighbor.