Human Pose Estimation from Ambiguous Pressure Recordings with Spatio-temporal Masked Transformers

Pressure-based pose estimation. The development of cost-efficient and ubiquitous tactile sensing systems has allowed researchers to model human interactions and gestures using pressure recordings. For instance, the patterns of hand grasp through pressure-sensing textiles have been explored in [2], while other works have developed multi-modal solutions for hand gesture detection using a combination of cameras and e-textiles [18]. Similarly, several works have explored virtual-reality and healthcare applications via human gait recognition [19] and posture detection [7] using an array of piezoresistive pressure sensors embedded in a carpet. A recent study [9] addressed the challenging task of 3D human pose estimation using only feet pressure distributions using a deep network based on 3D convolutions. Given the limited information available in feet pressure and the ambiguity of its mapping into a 3D pose, they achieved an error of 20$cm$, which is almost ten times more than vision-based approaches [11]. In another work, to disambiguate in-bed pressure maps, generation of human-like figures from a pressure recording matrix via a pre-processing block before pose estimation was proposed [8]. In the same line of work, in [5], a synthetic in-bed pressure dataset was presented to address the limited diversity of available datasets, allowing network pre-training for better performance in real-life applications.

Problem setup. Let $X=\{x_{1},x_{2},...,x_{T}\},x_{i}\in\mathbb{R}^{WH}$ be a sequence of $T$ pressure distribution frames and $Y=\{y_{1},y_{2},...y_{T}\},y_{i}\in\mathbb{R}^{3J}$ be a sequence of $3$D ground truth joint locations. Our goal is to train a pose estimator encoder as illustrated in Figure 1. We do so via learning the mapping of $\hat{Y}=P(Y|X,\theta,\theta_{r})$, where $\theta$ and $\theta_{p}$ are the trainable parameters of the encoder and the regression head, respectively. We pre-train the encoder via self-supervision prior to pose estimation. The details of our network architecture, training steps, and implementation details are described in the following sections.

Network Architecture. Similar to video-based vision transformers, we first tokenize the input frames $X$ into space-time cubes after applying a $(2+1)D$ convolutions [20]. Next, we embed the space-time cube tokens via a patch embedding layer, $F\in\mathbb{R}^{\frac{T}{d_{t}}\frac{W}{d}\frac{H}{d}C}$, and after a linear projection layer, pass them to an Encoder, which is a ViT [12] pre-trained on ImageNet [21] followed by MS-COCO [22] datasets. Here, $d_{c}$ and $d$ are the number of temporal and spatial crops, respectively. The ViT is a sequence of transformer blocks, each consisting of multi-head self-attention (MHSA) and feed-forward network (FFN) layers. In our setting, the encoder operates on all of the input patches without using the masked tokens with an output of $F_{o}\in\mathbb{R}^{\frac{T}{d_{t}}\frac{W}{d}\frac{H}{d}C}$. As illustrated in Figure 1 (b), we adopt a simple Regression head composed of two deconvolution layers and one $1\times 1$ convolution layer on the reshaped $F_{o}$ following the common setting of the previous works [23]. Specifically, each deconvolution block up-samples the reshaped feature vectors by a factor of two, and the $1\times 1$ convolution layer predicts the joint location heatmaps for each keypoint in each frame $K\in\mathbb{R}^{T\frac{W}{4}\frac{H}{4}J}$. The last module in our network is a decoder $\hat{X}=D(X|F_{o},Mask,\theta_{d})$ used only during the self-supervised pre-training of our encoder to reconstruct the input patches from the encoder’s output latent representations. We adopt a shallow transformer network to process $F_{o}$ and connect it to a linear projection layer to match the shape of the input patches.

Self-supervised pre-training. Masked Auto-Encoders (MAE) for pre-training transformer networks has shown strong results on a variety of applications such as NLP [24], pose estimation [11], and image classification [23]. Inspired by this, we implement an MAE by masking the outputs of our encoder, $F_{o}$, and then passing the masked encoded features along with the masked tokens to our decoder. For this purpose, we adopt an asymmetric design for masked image reconstruction, where the encoder operates on fully observed data (without masked tokens), and the decoder is applied on the masked encoder output. Our encoder-decoder network reconstructs the masked patches by learning to predict the raw pixel values of each patch. We use MSE loss on the masked patches, excluding the unmasked patches, to train the network efficiently [24].

Fine-tuning for Pose Estimation. In the fine-tuning step, we first freeze the decoder and warm up the regression head and patching block of the encoder for $2$ epochs. Then, we train all of the networks via minimizing the MSE loss between the predicted and ground-truth heatmaps, as well as the limb-length loss between the ground-truth and predicted keypoint coordinates, taken via SoftMax layer on each heatmap, by:

$$L_{link}^{i}=\begin{cases}\hat{L}_{i}-L_{95\%},&\text{if $\hat{L}_{i}>L_{95\%}$}\\ L_{5\%}-\hat{L}_{i},&\text{if $\hat{L}_{i}<L_{5\%}$}\\ 0&\text{otherwise,}\end{cases}$$

(1)

where $L_{5\%}$ and $L_{95\%}$ represent the $5^{th}$ and $95^{th}$ percentile of each of the limb lengths in the given dataset.

Implementation details. We use the same architecture proposed in ViTPose [11] for our encoder and only modify the patching layer by replacing the normal convolution operator with a $(2+1)D$ convolution block. In all cases, we start the training from pre-trained weights provided for ViTPose-B [11] in eligible layers. More specifically, we only use random initialization for the fully connected layers that have different sizes than the vanilla ViTPose due to the increased number of patches. Our decoder consists of 4 blocks of transformer networks with hidden dimensions of $256$. We initialize the Temporal ViTPose weights from the pre-trained model of MAE [23], and perform masked region reconstruction with a masking rate of $75\%$. We train the encoder-decoder pipeline with a learning rate of $1e^{-3}$ using AdamW optimizer [25] for $200$ epochs with a weight decay rate of $0.1$. In the next stage, we warm up the network by training the regression head, newly initilized layers, and the patching layer of ViT in a supervised task with a learning rate of $1e^{-3}$ for $2$ epochs. Finally, we use a learning rate of $2e^{-4}$ for training the network on the supervised pose estimation task for a total of $150$ epochs.

Acknowledgement. This project was funded in part by Natural Sciences and Engineering Research Council of Canada.