mmBaT: A multi-task framework for mmwave-based human body reconstruction and translation prediction

To aggregate the spatial and temporal information, this module takes the processed point cloud sequence ${\mathbf{P}}\in\mathbb{R}^{\textit{T}\times\textit{N}\times\textit{C}}$ as input, gotten from the body translation predictor module, where N denotes the fixed number of points in every frame and C represents the number of channels for each point.

We make use of PointNet++ [17] as the backbone to extract spatial features $\mathbf{F}$ of every point cloud frame ${\mathbf{P}}$:

$$\mathbf{F}=\text{PointNet++}(\mathbf{P})$$

(1)

The original PointNet++ adopts a max-pooling operation to aggregate local features from the sample points obtained by furthest point sampling, which may lose information from other sampled points[14]. As a result, we replace it with a linear layer inspired by the attention mechanism. Due to the inconsistency of the point clouds where some points may vanish in specific frames, the spatial features are fed into a bi-directional Gate Recurrent Unit (GRU) to obtain a 1024-dimensional global feature ${\mathbf{G}}$:

$$\mathbf{G}=\text{GRU}(\mathbf{F})$$

(2)

Raw radar point clouds often contain a significant amount of noise points, which can greatly affect the accuracy of body reconstruction. To extract the body-related points without any specific preprocessing or prior knowledge, an intuitive approach is to select points close to the subject, but the exact locations remain elusive during the given time interval. Besides spatial and temporal information, we also use the radial velocity information inherent in the mmWave point clouds. This information serves as an indicator for the motion trend of the subject, enabling us to make predictions regarding short-term translations. Consequently, there is an iterative interaction between the translation predictor and the global feature extractor.

Given a raw point cloud sequence $\mathbf{P}_{\text{in}}^{[t,t+T]}$ of length $T$, where each frame contains an arbitrary number of points. Our objective is to select $N$ points based on the body translation $\boldsymbol{\gamma}_{\boldsymbol{p}}^{[t,t+T]}$ predicted in the previous time window. Within this module, we take the global feature $\mathbf{G}^{[t-T,t]}$ as input, and a straightforward yet effective three-layer MLP decoder is adopted here to predict $\boldsymbol{\gamma}_{\boldsymbol{p}}^{[t,t+T]}$ for the current time window. The corresponding loss for this prediction is denoted as:

$$\mathcal{L}_{\text{pred}}=\sum\left\|\boldsymbol{\gamma}_{GT}^{[t,t+T]}-\boldsymbol{\gamma}_{\boldsymbol{p}}^{[t,t+T]}\right\|_{1}$$

(3)

Given $\boldsymbol{\gamma_{p}}$ as the central point, we generate a 3D bounding box with dimensions of 1m, 1m, and 3m along the x, y, and z axes, respectively. The points contained within the bounding box will be selected, and sampling or repeating is performed as needed to ensure a uniform amount of points across frames. The generated output ${\mathbf{P}}\in\mathbb{R}^{\textit{T}\times\textit{N}\times\textit{C}}$ is the point cloud sequence, as previously mentioned in Section 2.1, which is further utilized for global feature extractor.

Since the SMPL model parameters are the implicit representations of the human body’s pose and shape, and these parameters are highly non-linear and complex [14], direct regression from sparse and inconsistent point clouds is difficult. To address this challenge, we design the skeleton-aware body estimator that comprises two steps. In the first step, the coarse skeleton $\hat{\mathbf{J}}\in\mathbb{R}^{T\times N_{J}\times 3}$ is estimated by using an MLP decoder, where $N_{J}$ is the number of skeleton joints. The loss function for this step is defined as the L1 norm between the estimated coarse joint locations and the corresponding ground-truth locations:

$$\mathcal{L}_{\mathcal{{\hat{\mathbf{J}}}}}=\sum\left\|\mathbf{J}_{GT}-\hat{\mathbf{J}}\right\|_{1}$$

(4)

In the second step, we aim to regress the body parameters considering the awareness of skeleton joint locations. To fuse the global information with the skeleton spatial features, we concatenate the global feature ${\mathbf{G}}$ with the joint locations $\hat{\mathbf{J}}$ and get the fused joint features $\mathbf{I}\in\ \mathbb{R}^{T\times N_{J}\times(3+1024)}$. For a better feature representation and computational efficiency, we further put $\mathbf{I}$ through a linear layer to obtain enhanced joint features $\mathbf{I^{\prime}}\in\ \mathbb{R}^{T\times N_{J}\times 1024}$. Then, we employ self-attention mechanism [18] on $N_{J}$ to learn the implicit relationships among different joints:

$\displaystyle\text{attention}(\mathbf{Q},\mathbf{K},\mathbf{V})$

$\displaystyle=\mathbf{V}\cdot\operatorname{softmax}\left(\frac{\mathbf{Q}^{T}\cdot\mathbf{K}}{\sqrt{\mathbf{C^{k}}}}\right)$

(5)

where $\mathbf{Q}=\mathbf{W}_{Q}\cdot\mathbf{I^{\prime}}$, $\mathbf{K}=\mathbf{W}_{K}\cdot\mathbf{I^{\prime}}$ and $\mathbf{V}=\mathbf{W}_{V}\cdot\mathbf{I^{\prime}}$ denote the obtained queries, keys and values, respectively. Then, the softmax function is performed, and the attention output is calculated as the weighted sum of the value $\mathbf{V}$.

To increase the subspaces and model expressiveness, we apply multi-head attention instead of single self-attention mechanism:

	$\displaystyle\mathbf{h_{i}}$	$\displaystyle=\operatorname{attention}\left(\mathbf{Q_{i}},\mathbf{K_{i}},\mathbf{V_{i}}\right)$		(6)
	$\displaystyle\mathbf{O}$	$\displaystyle=\operatorname{\mathbf{W_{M}}\cdot concat}\left(\mathbf{h_{1}},...,\mathbf{h_{n}}\right)$		(7)

Multiple self-attention modules are performed in parallel with independent weights for $\mathbf{Q}$, $\mathbf{K}$ and $\mathbf{V}$. The results of $\boldsymbol{n}$ self-attention modules are concatenated to obtain the output $\mathbf{O}\in\ \mathbb{R}^{T\times N_{J}\times 1024}$ by a projection technique.

At the end of this module, the output $\mathbf{O}$ is fed into an MLP layer to regress the SMPL body parameters, including the pose parameter $\boldsymbol{\theta}\in\ \mathbb{R}^{T\times 6N_{J}}$ in a 6-dimensional representation, the shape parameter $\boldsymbol{\beta}\in\ \mathbb{R}^{T\times N_{\beta}}$ and the body translation $\boldsymbol{\gamma}\in\ \mathbb{R}^{T\times 3}$. These parameters are next fed into the off-the-shelf SMPL model to compute the final joint positions $\mathbf{J}$ and mesh vertices $\mathbf{M}$. The losses associated with these parameters are defined as:

$$\mathcal{L}_{\text{SMPL}}=\mathcal{L}_{\mathcal{\theta}}+\mathcal{L}_{\mathcal{\beta}}+\mathcal{L}_{\mathcal{\gamma}}+\mathcal{L}_{\mathcal{J}}+\mathcal{L}_{\mathcal{M}}$$

(8)

where $\mathcal{L}_{\mathcal{\beta}}$, $\mathcal{L}_{\mathcal{\gamma}}$, $\mathcal{L}_{\mathcal{J}}$ and $\mathcal{L}_{\mathcal{M}}$ are the L1 loss with the same definition in Eq.4, while $\mathcal{L}_{\mathcal{\theta}}$ is the geodesic loss [19] to represent the distance between the predicted and ground-truth rotation matrices. We aim to jointly optimize both tasks so the total loss of the framework is defined as:

$$\mathcal{L}_{\text{total}}=\mathcal{L}_{\mathcal{{\hat{\mathbf{J}}}}}+\mathcal{L}_{\text{SMPL}}+\mathcal{L}_{\text{pred}}$$

(9)

We employ two publicly available datasets to evaluate our proposed method:

mmBody dataset [20] is the first public dataset that offers Motion Capture (MoCap) ground-truth data in the SMPL-X format [23] for mmWave-based human body reconstruction task. It contains more than 100 motions performed by 20 subjects and covers different environments like dark, rain, and fog. The radar device has a range resolution of 0.4 cm, with each frame comprising thousands of points, the majority of which are noise points.

MRI dataset [21] mainly focuses on rehabilitation exercises, which contains over 160k frames from 20 subjects in a house monitoring condition. The radar device used in this dataset has a lower range resolution compared to [20], each frame has only 64 points. A method was presented in this paper to augment the number of points by fusing data from three consecutive frames. Since ground-truth data only contains 17 joint locations, mesh vertices are not considered for this dataset.

In this part, we will introduce our experimental setup in detail from the following aspects:

1) Training details: The deep learning architecture is implemented with PyTorch and trained on an NVIDIA GeForce RTX 3090 GPU. We train our method for 100 epochs with an Adam optimizer. The dropout ratio is set to 0.2 for GRU layers and linear layers. The initial learning rate is set to $1\times 10^{-3}$ with a decay rate $1\times 10^{-4}$. The batch size is set to 16. We apply the scale factors to normalize each loss term to the same scale. The number of points sampled from each frame is 1024 and 196 for mmBody dataset and MRI dataset, respectively. The training set and test set are divided according to the original dataset, and $10\%$ of the training data is used for validation. An initial bounding box is provided for our method for every test sequence.

2) Compared methods: A thorough evaluation of our proposed method is conducted on two existing methods: P4Transformer (P4Trans) [22] and mmMesh [10]. Specifically, P4Transformer is the benchmark employed to assess the validity of the mmBody dataset [20]. To further evaluate the performance of our body estimator module, besides the original implementations of these two methods, without given ground-truth bounding boxes (w/o bbx), we also compare our method with the modified version of these two methods provided with ground-truth bounding boxes (w/ bbx). This configuration is adopted to exclude the influence of noisy point clouds.

3) Evaluation metrics:

Mean Per Joint Rotation Error (MPJRE) [10] calculates the average differences between predicted joint rotations and the ground-truth rotations, which reports the pose accuracy of reconstruction results.

Mean Per Joint Position Error (MPJPE) [5, 6] calculates Euclidean distance between the estimated and ground-truth joint positions.

Mean Per Vertice Position Error (MPVPE) [6, 11] calculates Euclidean distance between the estimated and ground-truth error positions. This metric reports the overall mesh reconstruction results, which are affected by pose error, location error, and shape error.

Mean Translation Error (MTE) [10] measures Euclidean distance between predicted and ground-truth locations of the skeleton root joint. This metric represents the precision of body localization.

Mean Prediction Translation Error (MPTE) measures Euclidean distance between predicted and ground-truth locations of the skeleton root joint in our body translation predictor module.

In this section, we present the results of our comprehensive experiments, which evaluate the performance of our proposed mmBaT against existing methods. The quantitative results of our experiments are displayed in Table LABEL:tab:1.

1) Quantitative analysis: Overall, mmBaT consistently outperforms other methods across a range of metrics, with notable strengths observed in MPJRE, MPJPE, and MPVPE. These metrics are indicative of the superior ability of our method to accurately reconstruct human body poses and shapes. In particular, the significant improvement in the pose error (MPJRE) shows an average enhancement of 4.5^∘ and 9.25^∘ compared to P4Transformer and mmMesh in mmBody dataset, with ground-truth bounding box implementation. This significant improvement can be attributed to our coarse-to-fine approach, in which pose parameters are regressed based on the coarse skeleton as a priori, revealing the efficacy of our body estimator module in addressing those challenges. In terms of MTE, our approach closely matches the optimal results, deviating by only 0.7cm on average. This proximity to the optimal results is largely facilitated by the ground-truth bounding box, which aids in accurate subject localization.

By comparing the two methods with and without bounding box settings, we see great improvements across all metrics on the mmBody dataset, demonstrating the critical role of body translation prediction in accurate body reconstruction. Moreover, our method consistently achieves an MPTE under 15cm, showcasing its effectiveness in accurately predicting body translations and adeptly mitigating the influence of noisy point clouds.

In MRI dataset, all methods have relatively good performance, mainly because MRI dataset involves simple rehabilitation motions captured on a small scale and the point clouds contain fewer noise points. While mmBaT still remains the best results across all metrics, which shows the robustness of our approach in different environments.

2) Qualitative analysis: In Fig.2, we provide a visual demonstration of body reconstruction results, using mmBody dataset as a reference. The figure presents ground-truth data alongside the reconstruction results produced by two comparative methods with the ground-truth bounding box, followed by the outcomes generated by our proposed method, progressing sequentially from left to right. To facilitate a comprehensive evaluation of human poses, three representative frames have been selected. It can be visualized that our method achieves more favorable reconstruction results compared to other algorithms for pose estimation, proving the effectiveness of our novel multi-task framework and the superiority of our method.