Enhancement of Needle Localization Using Semi-Supervised Deep Learning

Ultrasound (US) imaging is widely used to monitor needle insertions. However, needle visibility is still not guaranteed especially with deep and steep punctures. Learning-based approaches are promising but are hindered by the lack of labeled data and their generalization capabilities. In previous work, a small dataset was acquired and used to train a U-Net like architecture for the needle localization task. Despite performing well on this dataset, the network failed to generalize when tested on different out of distribution datasets [1]. This work proposes using semi-supervised learning (SSL) techniques to improve model generalization across datasets despite limited annotated training data. Two SSL techniques were implemented: 1) Pseudo Labeling (PL) uses a U-Net like architecture with Monte Carlo (MC) dropout layers for uncertainty measurement so that only trusted pseudo-labels are used for retraining. To avoid noisy training, labels were further modified by an algorithm that uses the Hough transform. 2) Cross pseudo supervision (CPS) uses two identical U-Nets trained simultaneously on labeled and unlabeled data. For unlabeled data, Dice loss is calculated between the predictions of one network and the corresponding predictions of the other as ground truths. This loss enforces output consistency thereby enhancing robustness and generalizability. Both techniques were trained using two datasets (D1 and D2) with different distributions and their performance was benchmarked by a fully supervised model trained on only data from D1 called Lower Bound (LB) and on data from D1 and D2 called Upper Bound (UB). While the LB model lacked generalization ability with a tip location error of 0.62 mm, an angle error of 3.35°, and a 5.52% shaft failure rate, both SSL techniques enhanced generalizability and performance, achieving errors on par with the UB model with an average of 0.36 mm tip location error (mean absolute difference between true and predicted tip position), 2.04° angle error, and 0% shaft failure rate.