A New Paradigm for Spinal Keypoint Detection: Pixel-Level Contrastive Learning + Differentiable Shape Modeling

Spinal X-ray keypoint detection is a critical step for assessing the severity of spinal deformities, building biomechanical models, and preoperative planning. Traditional methods typically use heatmap regression, but they often struggle to distinguish the corner points of adjacent vertebrae. A recent paper proposes redefining anatomical keypoint detection as a pixel-level contrastive representation learning problem and incorporating a learnable shape prior, achieving performance competitive with state-of-the-art methods on multiple datasets.

From Heatmap Regression to Pixel-Level Contrastive Learning

Traditional keypoint detection methods typically model the problem as heatmap regression: the network predicts the probability that each pixel belongs to a keypoint, and then the peak is taken as the predicted location. However, this approach faces challenges in spinal X-rays—the corner points of adjacent vertebrae are close together, and heatmaps tend to be blurry, resulting in insufficient discriminability.

This study takes a different approach, treating keypoint detection as a pixel-level contrastive representation learning problem. Specifically, the network learns to map each pixel into an embedding space such that the embeddings of pixels belonging to the same keypoint are close to the keypoint prototype, while the embeddings of pixels belonging to different keypoints are far apart. Surprisingly, even with a generic U-Net backbone, this formulation can produce heatmaps with greater discriminability than traditional heatmap regression, thereby better distinguishing the corner points of different vertebrae.

The significance of this finding lies in the fact that it shows that through a contrastive learning objective, the model can automatically learn finer-grained features without requiring complex network designs. According to the paper, this method achieves keypoint accuracy and clinical parameter estimation comparable to state-of-the-art methods on the public AASCE dataset and two private spinal X-ray datasets.

Heatmap regression is a classic paradigm in human pose estimation and keypoint detection. In this approach, a Gaussian heatmap is generated for each keypoint as a supervision signal; the network outputs the probability that each pixel is that keypoint, and at inference time the coordinates of the heatmap peak are taken as the predicted location. The advantage of heatmap regression is that it preserves spatial structure information, achieves high localization accuracy for keypoints, and its output format is naturally compatible with convolutional networks. But its limitations are also obvious: when multiple keypoints are close together, the Gaussian heatmaps overlap and blur, making peaks difficult to distinguish; in addition, the heatmap resolution is affected by downsampling, requiring additional offset regression or upsampling to recover accuracy. In spinal X-rays, the corner points of adjacent vertebrae are closely spaced and similar in appearance, which is exactly the typical scenario where heatmap regression tends to fail.

Contrastive representation learning aims to pull positive sample pairs closer in the embedding space and push negative sample pairs farther apart, thereby learning highly discriminative feature representations. A common approach such as the InfoNCE loss approximates a lower bound on mutual information through a large number of negative samples. Pixel-level contrastive learning pushes this idea down to each pixel: pixels belonging to the same keypoint are treated as positive samples, pixels belonging to different keypoints are treated as negative samples, and a set of learnable keypoint prototypes is introduced as anchors. In this way, the network must not only determine whether a pixel "is a keypoint" but also "which keypoint it is," thereby naturally separating the corner points of different vertebrae in the embedding space. This objective provides a finer-grained supervision signal than per-pixel classification, and it is also the fundamental reason why this method's heatmaps have stronger discriminability.

Differentiable Shape Modeling: Disentangling Positional Ambiguity

Although pixel-level contrastive learning improves the discriminability of heatmaps, relying solely on local appearance information may still produce ambiguity in some cases, for example when a keypoint is occluded or when image contrast is low. To address this, the research team further introduced a learnable shape prior to model the global configuration of spinal keypoints.

Specifically, they combine the heatmap with the shape prior and optimize keypoint locations in a differentiable manner. The shape prior can capture the spatial relationships among spinal keypoints (such as the relative positions and angles of vertebrae), thereby providing global constraints when local information is unreliable and resolving positional ambiguity. This combination forms a simple yet practical pipeline for spinal keypoint detection.

According to the paper, this method shows competitive performance on the AASCE dataset and two private datasets, not only matching existing methods in keypoint localization accuracy but also achieving good results in estimating clinical parameters such as the Cobb angle.

Differentiable optimization refers to implementing an optimization process (such as minimizing the energy function of a shape prior) as a differentiable module, allowing gradients to propagate from the optimization result back to the front-end network for end-to-end training. In keypoint detection, a common approach is to jointly formulate heatmap responses and shape constraints as energy terms and obtain the final keypoint coordinates through several steps of differentiable iterative solving (such as gradient descent or a differentiable fitting layer). Compared with non-differentiable pipelines that first take peaks and then perform post-processing, differentiable shape modeling allows the shape prior to influence feature learning during training, making the heatmaps output by the network mutually adapt to global anatomical constraints. The cost is the need to design an appropriate energy function and number of iteration steps, and to trade off accuracy against computational overhead.

Open Source and Reproducibility

To facilitate follow-up research, the authors have open-sourced the code and pretrained weights at https://github.com/BioMedTP/Pix2Point. This initiative helps other researchers quickly reproduce and validate the method and promotes the development of the field of spinal imaging analysis.

Conclusion and Outlook

By redefining keypoint detection as a pixel-level contrastive learning problem and combining it with differentiable shape modeling, this study provides a new approach for keypoint detection in spinal X-rays. Experimental results show that the method is competitive on multiple datasets, and the open-source implementation lowers the barrier to use. In the future, this framework combining representation learning and shape priors is expected to be extended to keypoint detection tasks for other anatomical structures.

It should be noted that this method has currently been validated mainly on spinal X-rays, and its generalization ability to other modalities (such as CT and MRI) still needs further exploration. In addition, the design of the shape prior depends on the global configuration of keypoints, and its adaptability to severe deformities or abnormal anatomical structures also remains to be studied.

One-line Summary

Reframes spinal keypoint detection as pixel-level contrastive learning and integrates a differentiable shape prior to improve the discrimination of adjacent vertebral corner points.

Paragraph Summary

This article introduces a new method for keypoint detection in spinal X-rays. To address the difficulty traditional heatmap regression has in distinguishing the corner points of adjacent vertebrae, the authors redefine keypoint detection as a pixel-level contrastive representation learning problem, enhancing discriminability by matching pixel embeddings with keypoint prototypes; even with a generic U-Net backbone, it can generate clearer heatmaps than heatmap regression. To further eliminate positional ambiguity under occlusion or low contrast, the method introduces a learnable shape prior that models the global configuration of spinal keypoints in a differentiable manner and provides constraints. On the public AASCE dataset and two private datasets, the method's keypoint accuracy and estimation of clinical parameters such as the Cobb angle are comparable to state-of-the-art methods, and the code and pretrained weights have been open-sourced.