Simulation-Driven Training of Vision Transformers Enabling Metal Segmentation in X-Ray Images

TL;DR

This paper introduces a simulation-based training method for vision transformers to improve metal segmentation in X-ray images, demonstrating superior performance over CNNs and highlighting the potential of CAD-generated data for clinical applications.

Contribution

The study presents a novel simulation-driven training approach using CAD models and vision transformers for metal segmentation in X-ray images, outperforming traditional CNNs.

Findings

Vision transformer with dual decoder achieves 0.90 dice score.

CAD-based data generation enhances training flexibility and generalization.

Transformer-based model outperforms CNN encoder-based networks.

Abstract

In several image acquisition and processing steps of X-ray radiography, knowledge of the existence of metal implants and their exact position is highly beneficial (e.g. dose regulation, image contrast adjustment). Another application which would benefit from an accurate metal segmentation is cone beam computed tomography (CBCT) which is based on 2D X-ray projections. Due to the high attenuation of metals, severe artifacts occur in the 3D X-ray acquisitions. The metal segmentation in CBCT projections usually serves as a prerequisite for metal artifact avoidance and reduction algorithms. Since the generation of high quality clinical training is a constant challenge, this study proposes to generate simulated X-ray images based on CT data sets combined with self-designed computer aided design (CAD) implants and make use of convolutional neural network (CNN) and vision transformer (ViT) for metal segmentation. Model test is performed on accurately labeled X-ray test datasets obtained from specimen scans. The CNN encoder-based network like U-Net has limited performance on cadaver test data with an average dice score below 0.30, while the metal segmentation transformer with dual decoder (MST-DD) shows high robustness and generalization on the segmentation task, with an average dice score of 0.90. Our study indicates that the CAD model-based data generation has high flexibility and could be a way to overcome the problem of shortage in clinical data sampling and labelling. Furthermore, the MST-DD approach generates a more reliable neural network in case of training on simulated data.