Needle tracking in ultrasound images has lots of medical applications, such as image guidance for robotic needle intervention, brachytherapy, and biopsy. However, it is hard because needles are not always visible and that ultrasound images are usually noisy. Most previous research uses classical line detectors, such as Hough Transform [4] or RANSAC [6] to detect the needle But these methods assume that the needle is clearly visible in the ultrasound images. In fact, the needle can be bent in the images, and the visibility of the needle often changes from time to time. All these factors can reduce the accuracy of the classical line detection algorithms. Recently, deep learning-based segmentation is becoming popular because of its potential to learn the segmentation in difficult situations.

Transformer [16] was first utilized in NLP, because of its ability to handle long-range dependencies in sequence to sequence task. Unlike RNNs, each word in the sequential input has direct access to other words through self-attention in Transformers, leaving no room for any information loss in back-propagation through time. Ever since Transformer was proposed, it has been the most popular and powerful framework in the NLP community [3, 1]. Recently, Transformers have been applied in the field of computer vision as well. [5] proposed Vision Transformer that beats the current state-of-the-art model in image classification. [2] came up with an end-to-end Detection Transformer that performs better than other algorithms in object detection. Besides, Transformer based networks can also achieve good results on image segmentation tasks [20]. In terms of object tracking, TrackFormer [8] unified object detection and tracking with Transformer, where the detection result from the last frame is used to guide the detection in the current frame. Inspired by the performance of Transformer, we propose to apply Transformer to ultrasound needle tracking for its capabilities in attending to certain parts of the image as well as its exceptional performance in computer vision.

Deep Learning for Needle Tracking: Most research uses CNN-based architectures in needle segmentation or localization. Mwikirize et al. [10] propose a network model with a fully convolutional network and fast R-CNN to localize the bounding box of the needle in 2D ultrasound. Zhang et al. [18] propose attention U-Net to segment multiple needles in 2D ultrasound images. Zhang et al. [19] also apply Mask R-CNN in a similar task. Mukhopadhyay et al. [9] use an encoder and FCN network to find the bounding box of the needle. But most previous research does not consider the information in time series. To the best of our knowledge, only Mwikirize et al. [11] utilizes the time-series information by using a CNN-based architecture to segment the needle tip on the difference of consecutive frames. But the information in time sequence is not fully utilized since it is only introduced in the subtraction of different frames. Another drawback is that they are only detecting the needle tip, which is an easier task than segmenting the entire needle shaft.

Transformer: In medical image analysis, there have been quite a number of works focusing on applying Transformer-based frameworks to this field. [15] outperformed other medical image segmentation networks, such as U-Net [13], U-Net++[21], and U-Net with axial attention [17] in ultrasound image segmentation. It uses gated axial attention as the basic building block while having a local branch and a global branch of the network. The local branch takes in the image patches rather than the whole image whereas the global branch takes in the entire image. The local branch is more complex with the attention modules and has more layers compared with the global branch. [12] also proposed a U-Transformer that uses self-attention and cross-attention to improve the performance of medical image segmentation. The attention module in this work leverages the feature maps from the decoder branch that have rich semantic information and the feature maps from the skip connection which have information from high-resolution layers, allowing more accurate segmentation at the boundaries. In terms of Transformer in object tracking, [8] proposed to use a Transformer framework that takes in the CNN-encoded image features and the tracking result from the previous frame to produce the result for the current frame. [14] used the objects of interest as the queries in the attention module.

Learning to track objects in computer vision is hard since the algorithm might completely fall apart if the errors accumulate in every frame. Especially in ultrasound needle tracking, needles are sometimes not visible, leading to bad tracking results on certain frames. Therefore, it is important to quantify the certainty of the tracking results in each frame and incorporate that information into the framework to avoid the accumulation of errors. We model such certainty with a separate network besides our tracking network.

In this paper, we model the tracking of needles as a temporal segmentation of needles, i.e. we consider the segmentation of the previous frame when segmenting the current frame. We perform tracking with a Transformer-based framework. Our proposed method consists of two major innovations: (1) a confidence-driven tracking framework that models the certainty of the tracking results of the previous frame, (2) to the best of our knowledge, we are the first group to use a Transformer-based method in the ultrasound needle-tracking task.

To study which module can increase the accuracy, we use an ablation study to train the model in three different modes as shown in Table 1. For each mode, we test the tracking model with U-Net, U-Transformer with full attention (U-Trans), and U-Transformer with axial attention (Axial-U-Trans). The model to calculate the confidence map is always a U-Net. In total, we have nine different models (3 different tracking models times 3 modes).

U-Net-S, U-Net-T and U-Net-C are trained using a Quadro P5000, with CUDA 10.2 and PyTorch. Other networks (Transformer-based networks) are trained on a Titan RTX, with CUDA 11 and PyTorch. For the model parameters, we have 5 downsampling layers in the encoder side of the networks, with the number of filters as 64, 128, 256, 512, and 1024 respectively. The decoder side follows the same setting in reverse order. Adam [7] optimizer is used with a learning rate of 0.0001. The dataset includes 15 trials for training, 5 trial for validation, and 7 trials for testing. In total there are 1004 images for training, 302 images for validation, and 427 images for testing.

We compare three quantitative results:

• Shaft error: we first use Otsu-threshold to threshold the output of the model, then calculate the weight average distance between the positive pixels to the ground truth pixels. The distance is defined as the smallest distance of the pixel to all the ground truth pixels.

• Tip error: we first use Otsu-threshold to threshold the output of the model. Since in the dataset, the needle is inserted from the right edge of the image, we define the needle tip as the positive pixel with the smallest x- coordinate, and the max y-coordinates on that column (that is, the bottom-left pixel in the label).

• Dice coefficient: it is twice the IoU between the ground truth label and the segmentation.

The quantitative results in Table 2 show that the U-Net tracking (U-Net-T) model works the best in shaft fitting and localize the needle tip. The U-Net tracking with confidence (U-Net-C) model works slightly better in the dice coefficient. Among the transformer families, the U-Transformer with axial attention with confidence (Axial U-Trans-C) performs the best.

From these results, we can see that the segmentation map from the previous frame improves the U-Net performance, but the confidence slightly decreases the shaft fitting and tip localization accuracy but giving a higher dice coefficient. It is possible that the Otsu thresholding does not find the best threshold to convert the prediction to a segmentation map. For all the U-Transformer models, the confidence can improve the segmentation performance since Axial U-Transformer-C and U-Transformer-C has lower shaft fitting and tip localization error. The axial attention performs better than the full attention, and the network with attention modules performance worse than U-Net. The results show that a simpler model works better in this case. It is possible that our dataset is too small that it is harder to train a good large model.

In this project, we test different kinds of models for needle tracking in ultrasound images. Our results show that in this task, the smaller models perform better than the large models. Our results also prove that the information from the previous segmentation can improve the segmentation accuracy of the model, but the effect of the generated confidence is not clear since it has different effects on the U-Net and the Transformers. Our future work will include exploring more ideas on U-Net-based tracking models, such as introducing LSTM into the models, to preserve the time-series information better.