SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jia-Xing Zhong, Ta-Ying Cheng, Yuhang He, Kai Lu, Kaichen Zhou $\boxtimes$ , Andrew Markham, Niki Trigoni
Department of Computer Science, University of Oxford
{jiaxing.zhong, ta-ying.cheng, yuhang.he, kai.lu, rui.zhou}@cs.ox.ac.uk
{andrew.trigoni}@cs.ox.ac.uk
A truly generalizable approach to rigid segmentation and motion estimation is fundamental to 3D understanding of articulated objects and moving scenes. In view of the closely intertwined relationship between segmentation and motion estimates, we present an SE(3) equivariant architecture and a training strategy to tackle this task in an unsupervised manner. Our architecture is composed of two interconnected, lightweight heads. These heads predict segmentation masks using point-level invariant features and estimate motion from SE(3) equivariant features, all without the need for category information. Our training strategy is unified and can be implemented online, which jointly optimizes the predicted segmentation and motion by leveraging the interrelationships among scene flow, segmentation mask, and rigid transformations. We conduct experiments on four datasets to demonstrate the superiority of our method. The results show that our method excels in both model performance and computational efficiency, with only 0.25M parameters and 0.92G FLOPs. To the best of our knowledge, this is the first work designed for category-agnostic part-level SE(3) equivariance in dynamic point clouds.
Comprehending point cloud motion is critical for various 3D vision tasks in dynamic scenarios, e.g., tracking, animation, simulation, and manipulation. Many types of 3D motion can be described as a composition of multi-body rigid movements, such as articulated objects [68], and vehicular traffic scenes [25]. Specifically, the setting of multi-body rigid motion requires all moving parts (i.e., bodies) to undergo only translation and rotation, without any type of deformation. The process of modeling these multi-body rigid movements typically involves two primary portions [32]: 1) the identification of distinct moving bodies (i.e., rigid segmentation), and 2) the calculation of individual movements for each identified body (i.e., motion estimation).
Figure 1: Comparison of FLOPs, parameter number, and AP on SAPIEN. Supervised methods are marked in  $\bigcirc$
Being the first modeling portion, rigid segmentation
significantly differs from traditional semantic segmentation tasks that concentrate on category information [16, 3, 30, 31, 70]. Based on multi-body motion instead of category-dependent semantics, rigid segmentation becomes category-agnostic about moving parts as rigid bodies are not limited to a specific set of shapes and naturally cannot be assigned certain category labels. As regards motion
estimation, individual pose variations would occur arbitrarily, including transformations unseen in the training data. This constitutes an open set of pose changes.
Considering the relative inaccessibility and the constrained generalizability of manual annotations, Song and Yang [61] have proposed a seminal work for unsupervised rigid segmentation. Nevertheless, the unsupervised paradigm remains a formidable challenge because of the generalizable requirement in multi-body movements and the interdependent nature between rigid segmentation and motion estimation. Due to the aforementioned category-agnostic rigid prior and open-set pose changes in multi-body motion, a model is required to generalize well to transformations of appearance, location, and orientation. Moreover, rigid segmentation and multi-body motion estimation are highly coupled — the calculation of movements for each body is based on the output of a rigid partition, while an estimate of multi-body motion facilitates further segmentation. In absence of precise supervision, it is difficult to independently obtain effective training signals for either rigid segmentation or motion estimation.
To tackle the difficulties of generalizability and interdependence, this paper presents a three-fold contribution in the form of architecture, training strategy, and experimental evaluation:
1. Architecture. We design a part-level SE(3)-equivariant network that demonstrates strong generalization to open-set motion and category-agnostic moving bodies. The SE(3)-equivariance allows features of a model to keep coherence with the rigid transformations of point cloud inputs, thereby enhancing the model's robustness to such spatial transformations (i.e., rigid motion) [13, 17, 9, 14, 67] and therefore is more generalizable for open-set pose changes of global rigid targets [63, 43, 8, 53]. Tailing the SE(3)-equivariance is two lightweight heads for segmentation and motion estimation. The former, unlike previous segmentation networks using global SE(3) equivariance, exhibits point-level local flexibility to SE(3)-invariance. The latter, robustly employing the probability of part consistency, utilizes part-level "softly" matching operations to handle SE(3)-equivariance without requiring knowledge of part categories. As a result of integrating two lightweight heads into a unified network, rather than utilizing large independent networks, our model is characterized by a small number of parameters  $(0.25\mathrm{M})$  and low computational complexity  $(0.92\mathrm{G}$  FLOPs). To the best of our knowledge, this is the first work designed for category-agnostic part-level SE(3)-equivariance in dynamic point clouds.
2. Training Strategy. We leverage the interdependence between rigid segmentation and motion estimation and present a unified training strategy to jointly optimize their outputs in an online fashion. Based on our proposed two-head network, we simultaneously filter out the noisy flow predictions and refine the estimates of rigid motion by exploiting the interrelation among scene flow, segmentation mask, and rigid transformation. Superior to previous arts, our online training process is free from complicated components and manual intervention, e.g., the alternation of Markov Chain Monte Carlo proposals and Gibbs sampling updates [28], the offline repetition of training multiple segmentation networks from scratch [61], or the continual optimization of an accessory neural network of flow estimation [32].
3. Experiments. Experiments on four datasets demonstrate the efficacy of our model in performance, as well as its efficiency in parameters and computational complexity. We conduct comprehensive experiments on four datasets (SAPIEN [68], OGC-DR [61], OGC-DRSV [61], and KITTI-SF [48]) across three application scenarios (articulated objects, furniture arrangements, and vehicular traffic). Noticeably, as shown in Figure 1, our performance on the SAPIEN dataset of articulated objects surpasses state-of-the-art results w.r.t. all evaluation metrics, achieving the relative gain of at least  $14.7\%$  AP in rigid segmentation and at least  $23.3\%$  in motion estimation, with only  $0.25\mathrm{M}$  parameters and  $0.92\mathrm{G}$  FLOPs. The code is available at https://github.com/jx-zhong-for-academic-purpose/Multibody_SE3.
Deep Learning on Dynamic Point Clouds. Deep neural networks for point cloud video modeling aim to understand our dynamic 3D surroundings. These networks enable the resolution of several downstream tasks, such as video-level classification [20, 22, 44, 23, 21, 74, 29], frame-level prediction [57, 36, 50], object-level detection [6, 40, 54, 58] or tracking [26, 56, 72], part-level mobility parsing [60, 59, 71], point-level segmentation [22, 23, 7, 20], and scene flow estimation [18, 64, 34, 11, 2]. In contrast to the above dynamic tasks that usually assume contiguous
sequential input, multi-body rigid motion may be captured between discrete frames [32], presenting a unique challenge.
3D Motion Segmentation & Multi-body Rigid Motion. Although the systematic formulation of multi-body rigid motion modeling is a relatively recent development [32], a related problem, i.e., motion segmentation, has been long sought after. Motion segmentation aims to group points with similar motion patterns from the input of scene flow, using such techniques as factorization [10, 42, 69], clustering[33], graph optimization [47, 35, 4], and deep learning [39, 28, 71, 1]. However, motion segmentation does not explicitly take into account multi-frame multi-body consistency. To address this issue, Huang et al. [32] present the seminal fully-supervised work for modeling multi-body rigid motion. Song and Yang [61] further attempt to predict the mask of rigid segmentation via three object geometry losses in an unsupervised setting. Following the same unsupervised paradigm, our approach has distinct motivations: 1) to achieve high generalizability in the presence of low-quality training signals, and 2) to simplify the training process through online optimization.
SE(3) Equivariant Networks for Point Clouds. Equivariant networks have superior discriminative and expressive ability for various data structures (e.g., graphs [5, 65], images [15, 49]) under the transformation of some symmetry groups. Among them, the equivariance of a Special Euclidean group (3) (SE(3)) has drawn increasing attention to 3D point cloud processing [13, 17, 9, 14, 67, 63, 43, 8, 53]. A vast majority of these researches focus on global, while the part-level local equivariance is under-explored. Recently, [73] has utilized part-level SE(3) equivariance for supervised bounding box detection. Concurrently with our work, Lei et al. [41] and Feng et al. [24] apply part-level equivariance to category-specific tasks of object segmentation and human body parsing, respectively. Differently, we leverage part-level equivariance to the unsupervised category-agnostic problem.
3.1 Background: SE(3)-equivariance/invariance & Discretization
Given a point cloud  $X \in \mathbb{R}^{n \times 3}$  of  $n$  points and a rigid transformation  $\mathbf{T} \in \mathrm{SE}(3): \mathbb{R}^{n \times 3} \to \mathbb{R}^{n \times 3}$ , a neural network mapping inputs to a feature domain  $\phi: \mathbb{R}^{n \times 3} \to \mathcal{F}$  is defined as  $SE(3)$ -equivariant if  $\phi(\mathbf{T} \circ X) = \mathbf{T} \circ \phi(X), \forall \mathbf{T} \in \mathrm{SE}(3)$ , where  $\circ$  means performing an SE(3) transformation. Likewise,  $SE(3)$ -invariance is expressed as  $\phi(\mathbf{T} \circ X) = \phi(X)$ . To reduce the computational cost of equivariant networks, [12] discretizes the rotation space into an icosahedral group  $\mathcal{G}$  of 60 rotational angles ( $|\mathcal{G}| = 60$ ). As a further extension, Equivariant Point Network (EPN) [9] operates the SE(3) discretization on point clouds. For a point  $x \in \mathbb{R}^3$  in the input  $X$ , the feature extractor of EPN outputs the per-point representation  $f(x) \in \mathbb{R}^{|G| \times D}$ , where  $D$  is the feature dimension. By definition,  $f(x)$  is SE(3)-equivariant to its neighbors of the convolution kernel's receptive field:
1) Rotation equivariance within the icosahedral group:  $f(g\circ x) = g\circ f(x),\forall g\in \mathcal{G}$
2) Translation invariance w.r.t. arbitrary translation  $t\colon f(t\circ x) = f(x)$
By relaxing the strictly rotational equivariance into the equivariance w.r.t. the 60 icosahedral angles in  $\mathcal{G}$ , EPN is proven to be robust to many downstream tasks, such as 6D pose estimation [43], and place recognition [45]. Due to its high robustness, we choose EPN as our SE(3) equivariant backbone of feature extraction. In principle, our training strategy is versatile, and other networks with per-point  $SE(3)$ -equivariant representations can also serve as its feature extractor.
We define a set of point clouds  $P$  from  $K$  frames that are not necessarily consecutive as  $P = \{P_1, P_2, \dots, P_{K-1}, P_K\}$ , where each frame  $P_k = \{p_k^i \in \mathbb{R}^3 | i = 1, 2, \dots, N-1, N\}$  contains  $N$  three-dimensional points. All frames of  $P$  represent the same target, which is segmented into  $S$  independently moving rigid parts. For a given frame  $P_k$ , the point-part rigid mask is defined as  $\mathbf{M}_k \in \{0, 1\}^{N \times S}$ .  $\mathbf{M}_k^{is} = 1$  indicates that the point  $p_k^i$  belongs to the  $s^{th}$  rigid part;  $\mathbf{M}_k^{is} = 0$  indicates that it does not. The rigid motion of the  $s^{th}$  part between two frames  $(P_k, P_l)$  is denoted
<sup>1</sup>The term "target" refers to objects, environments, or scenes associated with multi-body rigid motion, such as articulated objects, arrangements of indoor furniture, or scenes of vehicular traffic.
Figure 2: An overview of network structure and training strategy. We feed a pair of point clouds  $P_{k}, P_{l}$  into the SE(3) equivariant networks to obtain a pair of SE(3) icosahedronic features. These features are then fed into the two proposed heads for motion segmentation and invariant segmentation. The two outputs, combined with scene flow, are used for a joint optimization that can be done in an unsupervised manner.
as  $\mathbf{T}_{kl}^{s}\in \mathrm{SE}(3)$ . This rigid transformation is specified by a rotation matrix  $\mathbf{R}_{kl}^{s}\in \mathrm{SO}(3)$  and a translation vector  $\mathbf{t}_{kl}^{s}\in \mathbb{R}^{3}$ . Note that the targets in the dataset belong to various categories, including classes possibly unseen in the training data.
Given the point cloud set  $P$ , we aim to segment each frame  $P_{k}$  with rigid-part masks  $\mathbf{M}_k$ , and obtain part-level rigid transformation  $\mathbf{T}_{kl}^{s}$  between frames. The former task is called rigid segmentation, while the latter is named motion estimation. In the supervised setting, ground-truth segmentation masks and rigid motions (or scene flow) are provided in the training data. In the unsupervised formulation, neither segmentation nor motion information is available, and our training input solely consists of point cloud frames.
As shown in Figure 2, our network takes a pair of point cloud frames as the input and outputs the predictions of rigid segmentation and motion estimation. There are three main components in the proposed framework, i.e., a feature extractor, an invariant segmentation head, and an equivariant motion estimation head.
Per-point Feature Extractor. For an input frame  $P_{k}$ , the feature extractor of EPN outputs per-point SE(3)-equivariant representations  $F_{k} \in \mathbb{R}^{N \times |\mathcal{G}| \times D}$ , following the notations in Section 3.1 and 3.2. The corresponding feature  $f_{k}^{i} \in \mathbb{R}^{| \mathcal{G}| \times D}$  of a point  $p_{k}^{i}$  can be viewed as a concatenation of different representations w.r.t.  $g_{j} \in \mathcal{G}$  over the rotation group dimension:
$$
f _ {k} ^ {i} = \left[ \theta \left(p _ {k} ^ {i}, g _ {1}\right), \theta \left(p _ {k} ^ {i}, g _ {2}\right), \dots , \theta \left(p _ {k} ^ {i}, g _ {| \mathcal {G} | - 1}\right), \theta \left(p _ {k} ^ {i}, g _ {| \mathcal {G} |}\right) \right], \tag {1}
$$
where  $\theta$  is a  $D$ -dimensional encoder based on a stack of convolution kernels. The prediction module of vanilla EPN is designed for global SE(3)-equivariance for all input points. Differently, our unsupervised multi-body task requires the model's ability to handle part-level local equivariance, especially under low-quality training signals. For this purpose, we further devise two heads for rigid segmentation and motion estimation.
Point-level Invariant Segmentation Head. Rigid segmentation is an SE(3)-invariant task, as the predicted mask should remain consistent for the same points across various poses and positions. Traditional SE(3)-equivariant structures assume that all input points undergo the same rigid transformation, which is not in alignment with the multi-body setting. To encode distinct transformations for individual input points, the equivariant features  $\theta(p_k^i, g_j)$  are aggregated into an invariant representation  $u_k^i$  across the dimension of the rotation group, as depicted in Figure 3:
$$
u _ {k} ^ {i} = \sum_ {j} ^ {| \mathcal {G} |} w \left(p _ {k} ^ {i}, g _ {j}\right) \theta \left(p _ {k} ^ {i}, g _ {j}\right), \tag {2}
$$
where  $w(p_k^i, g_j) \in [0, 1]$  is a selection probability of the discrete rotation  $g_j$  in  $\mathcal{G}$ , derived through a  $1 \times 1$  convolution. The weighted sum  $u_k^i$  is invariant to rigid motion given a point with its neighbors
Figure 3: An overview of Segmentation and Motion heads. The invariant segmentation head comprises an invariant module that sums  $\theta(\cdot)$  and  $w(\cdot)$  to obtain an invariant representation which is then aggregated through multiple layers to obtain a segmentation mask. Pairs of segmentation masks  $\hat{M}_k, \hat{M}_l$  are fed into the equivariant motion head along with per-point equivariant features  $F_k$  to obtain correlations features for predicting the final transformations of points.
in the convolution receptive field. As Figure 3 illustrates, by fusing such invariant representations  $u_{k(1)}^i \in \mathbb{R}^D, \dots, u_{k(h)}^i \in \mathbb{R}^{D'}$  from  $h$  layers in the feature extractor, the segmentation head outputs a soft prediction  $\hat{M}_k \in [0,1]^{N \times S}$  of the rigid mask.
Part-level Equivariant Motion Estimation Head. Part-level SE(3)-equivariance is desirable for motion analysis, especially rotation estimation. Based on the noisy predictions  $(\hat{M}_k,\hat{M}_l)$  of the frames  $(k,l)$  from the head of rigid segmentation, the motion head is supposed to handle these uncertain category-agnostic parts. Figure 3 demonstrates the operational scheme of our motion estimation head. First of all, the part-level SE(3) feature  $V_{k,j}\in \mathbb{R}^{N\times D\times S}$  w.r.t. the rotation  $g_{j}\in \mathcal{G}$  of a single frame  $k$  is obtained from the per-point equivariant representations  $F_{k}$  and the predicted rigid mask  $\hat{M}_k$ :
$$
V _ {k, j} = \left\{\hat {m} _ {k} ^ {1} \cdot \theta \left(p _ {k} ^ {1}, g _ {j}\right), \hat {m} _ {k} ^ {2} \cdot \theta \left(p _ {k} ^ {2}, g _ {j}\right), \dots , \hat {m} _ {k} ^ {N - 1} \cdot \theta \left(p _ {k} ^ {N - 1}, g _ {j}\right), \hat {m} _ {k} ^ {N} \cdot \theta \left(p _ {k} ^ {N}, g _ {j}\right) \right\}, \tag {3}
$$
where  $\hat{m}_k^i$  is the element corresponding to a point  $p_i^k$  in  $\hat{M}_k$ , and  $\cdot$  is the broadcast operation of Hardamard product. Afterward,  $V_{k,j}$  over the rotation group  $\mathcal{G}$  is concatenated as  $V_{k} \in \mathbb{R}^{N \times |\mathcal{G}| \times D \times S}$ , followed by a permutation-invariant operation  $\sigma: \mathbb{R}^{N \times |\mathcal{G}| \times D \times S} \to \mathbb{R}^{|\mathcal{G}| \times D \times S}$  (e.g., max pooling) over all the points to produce part-level equivariant features  $\tilde{V}_k = \sigma(V_k)$ . Between two frames  $(k,l)$ , the part-level rotation correlated feature  $C_{kl} \in \mathbb{R}^{|\mathcal{G}| \times |\mathcal{G}| \times S}$  is defined as:
$$
C _ {k l} = \tilde {V} _ {k} \tilde {V} _ {l} ^ {T}, \tag {4}
$$
where  $T$  is matrix transposition.  $C_{kl}$  is calculated upon "softly matching" within each consistent rigid part, while the specific category labels can be agnostic to the model. Based on the correlated feature  $C_{kl}$ , the motion head estimates rotation  $\hat{\mathbf{R}}_{kl}^{s}$  and translation  $\hat{\mathbf{t}}_{kl}^{s}$  of each rigid part  $s$ . More implementation details can be found in Supp.
3.4 Unsupervised Training Strategy
Ideally, the relation among scene flow  $\delta_{kl} \in \mathbb{R}^{N \times 3}$ , rigid segmentation  $M_k^s$  of the  $s^{th}$  part, and multi-body transformation  $\mathbf{T}_{kl}^s$  is as follows:
$$
P _ {l} = P _ {k} + \delta_ {k l} = \bigcup_ {s} ^ {S} \left\{\mathbf {T} _ {k l} ^ {s} \circ \left(M _ {k} ^ {s} \star P _ {k}\right) \right\}, \tag {5}
$$
wherein the operation of union set is denoted as  $\bigcup$ , and  $\star$  signifies the removal of a point  $p_i^k$  if it does not in the rigid  $s$ . Nevertheless, their predictions  $\hat{\mathbf{T}}_{kl}^s$ ,  $\hat{M}_k$  and  $\hat{\delta}_{kl}$  may exhibit a considerable amount of interdependent noise during the unsupervised training process. To mitigate the noise in one prediction by leveraging the other two, we propose an online training strategy without interrupting the end-to-end optimization of a network, as shown in Figure 2.
Cold Start & Scene Flow  $\hat{\delta}_{kl}$  Updating. Scene flow creates a relation between motion and segmentation. As a dense vector field that maps points to points, it can directly assist point-level segmentation. At the same time, scene flow also contains movement information for motion estimation. We initialize by collecting noisy scene flow from an unsupervised flow estimator (e.g., FlowStep3D [38]), and calculate rudimentary estimates of rigid masks and transformation by minimizing the discrepancy between their derived displacement and scene flow. As the accuracy of their estimates  $\hat{M}_k^s$  and  $\hat{\mathbf{T}}_{kl}^{s}$  improves, the scene flow is online corrected during a training epoch:
$$
\hat {\delta} _ {k l} = \alpha \hat {\delta} _ {k l} + (1 - \alpha) \left(\bigcup_ {s} ^ {S} \left\{\hat {\mathbf {T}} _ {k l} ^ {s} \circ \left(\hat {M} _ {k} ^ {s} \star P _ {k}\right) \right\} - P _ {k}\right), \tag {6}
$$
where  $\alpha \in [0,1]$  is a decay factor to control the updating rate. In this manner, improved scene flow is capable of providing enhanced supervision to learn segmentation masks and motion estimates.
Motion & Flow  $\rightarrow$  Segmentation  $\hat{M}_k^s$ . Although scene flow encodes point-level supervisory signals for segmentation, accurately computing flow between non-adjacent frames is challenging, as pointed out by Song and Yang [61]. To alleviate the influence of miscalculated scene flow, we optimize our segmentation head based on the consensus between the motion head and updated scene flow. Given the outputs  $\hat{\mathbf{R}}_{kl}^{s}$ ,  $\hat{\mathbf{t}}_{kl}^{s}$  of our motion head, and the interpolation point  $p_l^{\prime i} = p_k^i +\delta_{kl}^i$  derived from scene flow, the consensus score  $\beta_{kl}^{s(i)}$  of the point  $p_k^i$  w.r.t. the rigid  $s$  between the frames  $(k,l)$  is defined as:
$$
\beta_ {k l} ^ {s (i)} = e x p \left(- \tau | | \hat {\mathbf {R}} _ {k l} ^ {s} p _ {k} ^ {i} + \hat {\mathbf {t}} _ {k l} ^ {s} - p _ {l} ^ {\prime i} |\right), \tag {7}
$$
where  $\tau$  is a temperature coefficient to set the sharpness of  $\beta_{kl}^{s(i)}$ , and  $||\cdot ||$  is  $\ell 2$ -norm. Intuitively, a low  $\beta_{kl}^{s(i)}$  indicates a large disagreement between scene flow and the motion head, of which the mask should have a small learning weight in the segmentation loss:
$$
l _ {s e g} = \frac {1}{N S} \sum_ {i} ^ {N} \sum_ {s} ^ {S} | | \beta_ {k l} ^ {s} \hat {M} _ {k} ^ {s} \left(p _ {l} ^ {\prime i} - \hat {\mathbf {T}} _ {k l} ^ {s} \circ p _ {k} ^ {i}\right) | |, \tag {8}
$$
where  $\hat{\mathbf{T}}_{kl}^{s}$  is computed based on the scene flow and segmentation, as described in the next paragraph.
Segmentation & Flow  $\rightarrow$  Motion  $\hat{\mathbf{T}}_{kl}^{s}$ . Following previous works [32, 61], we employ the weighted-Kabsch algorithm [37, 27] to determine part-level rigid transformation  $\hat{\mathbf{T}}_{kl}^{s}$  given the estimates of scene flow and rigid masks. Further, the motion head is optimized by the rotation component of  $\hat{\mathbf{T}}_{kl}^{s}$ , and estimates corresponding translation by minimizing our probability-based part-level distance. The implementation details of this distance and our motion loss can be found in Supp.
4.1 Datasets & Metrics