Semi-supervised object detection based on single-stage detector for thighbone fracture localization

CAD has been extensively studied in the past decade[21, 22], and CAD system based on deep learning has been developed to diagnose a wide range of Pathology such as detection of covid-19[23, 24], mass and calcification features in mammography[25] and brain tumor diagnosis[26]. In the fracture detection method based on deep learning[27], FAMO[7] constructed the Feature Ambiguity Mitigate Operator model to mitigate feature ambiguity in bone fracture detection on radiographs of various body parts. Due to the requirements of medical professional knowledge, the labor cost of large-scale annotations is expensive which hinders the development of CAD solutions based on deep learning. Computer aided detection using SSL method is an emerging task in recent years, such as Yirui Wang et al. proposed the adaptive asymmetric label sharpening (AALS) algorithm using the teacher-student model paradigm, which solves the label imbalance problem unique to the medical field[28].

Object detection is one of the core tasks in computer vision. At present, the object detector based on CNN can be divided into single-stage and two-stage detectors. FasterRcnn[18] is the representative two-stage detector, which uses RPN network for proposal extraction and RCNN head for regional prediction and extraction of objects. The single-stage detector only uses the features extracted by the feature extraction network for regression and classification. For example, SSD[29] uses the feature pyramid method to complete target regression and classification on different scale features at the same time. Chen et al. developed the YOLOF that only uses C5 feature for detection as shown in Figure 1 in which the complex Multiple-in-Multiple-out encoder is replaced by the simple Single-in-Single-out encoder, YOLOF containing two key components: dilated encoder and uniform decoder.

SSL method plays a leading role in image classification[31, 32, 33, 34, 35]. Because the object detector has complex architecture design and multi task learning (classification and regression), it is not a simple work to transfer the SSL method to the object detection task. The current SSOD method mainly has two directions: Consistency Regularization[36] and Pesudo Label[8]. The former uses two deep convolution neural networks to learn the consistency between different data augmentation[37] (horizontal flip, different contrast, brightness, etc.) of the same unlabeled image, and make the image prediction to small disturbance the same. The latter uses the pre-training model learned on labeled data to infer the unlabeled data. In recent years, semi-supervised object detection method has attracted people’s attention[38, 39, 40]. STAC[19] first applies pseudo label method to SSOD, it apply weak data augmentation to unlabeled data, and uses the trained teacher model to generate pseudo labels of unlabeled images. Unbiased teacher[41] uses focal loss[16] to solve the imbalance between positive and negative samples. Instant teaching[42] trains two models at the same time to check and correct pseudo labels for each other, so as to effectively suppress the accumulation of false predictions. Almost all the above work is based on the two-stage detector, such as FasterRcnn, which is not convenient to develop in the medical field with limited resources. Inspired by the above works, we designed a fast semi-supervised detection model based on the single-stage detector.

In each training iteration, unlabeled images and labeled images are extracted according to a certain data sampling ratio. The data are preprocessed by two different preprocessing methods to obtain strong augmented labeled images, weak augmented and strong augmented unlabeled images. The student network is trained with the pseudo boxes generated by teacher model and ground truth boxes in labeled images. The total loss function can be expressed by (1) :

where $L_{\text{sup }}$ and $L_{\text{unsup}}$ represent the loss function of labeled images and the unlabeled images respectively, and $\lambda$ represent the weight of unsupervised loss in the total loss function.

At the beginning of training, the student model and teacher model adopt random initialization. For labeled data, the student network uses its ground truth to calculate the loss $L_{\text{sup }}$ as (2) and update the student model parameters by the gradient descent method. For unlabeled data, the teacher network first deduces the result of weakly augmented unlabeled data, then filtering out the low-quality results according to the confidence threshold $\sigma$. The retained high-quality predictions is considered to be the pseudo labels for unlabeled data. Then, the student calculates the unsupervised loss $L_{\text{unsup }}$ of unlabeled data as (3), and updates the parameters by the gradient descent method. Finally, the teacher network is updated with the exponential moving average (EMA) for the student model.

where $N_{\text{label }}$ indicates the number of positive samples in the labeled data, $N_{\text{unlabel }}$ indicates the number of positive samples in pseudo data with retained classification scores above the threshold $\sigma$. $v_{t}$ and $\hat{v}_{t}$ represent the student predicted category and ground truth of the t-th positive sample, $\theta_{t}$ and $\hat{\theta}_{t}$ represent the student predicted regression result and corresponding target of the t-th positive sample. $v_{t}^{*}$ and $\theta_{t}^{*}$ represent the category and box of the positive sample whose classification score from pesudo labels is higher than the threshold $\sigma$.

The quality of pseudo labels determines the performance of SSL model, in our framework, pseudo labels with low classification scores will be filtered out with confidence threshold $\sigma$ to ensure the quality of pseudo labels. Neither the high confidence threshold setting nor low confidence threshold setting is suitable for the semi-supervised framework based on single-stage detectors. For the SSOD method based on single-stage detector, directly using the same setting as the two-stage detector (confidence threshold = 0.7 in [41]) will greatly increase the number of negative samples, and only few part can be predicted in positive samples[40]. This imbalance will result in not getting enough pseudo labels, which limits the performance of the semi-supervised single-stage detector. On the contrary, low threshold threshold will easily bring in more low-quality pseudo labels with noise. When we adopt confidence $\sigma$ = 0.5, as shown in Figure 4, AP50 decreases to a great extent with the increase of training iterations. We believe with training iterations increasing, Teacher have been able to generate enough pseudo labels with confidence greater than $\sigma$. The number of pseudo labels has been much more than that in the early stage of training, but limited to low threshold settings, the classification score of pseudo labels has not been greatly increased with the increase of training iterations. We need to adopt a more strict confidence screening strategy in the middle stage of training and provide higher score pseudo labels.

We propose the adaptive difficult sample oriented (ADSO) moudle, which inherits the advantages of the flexibility of the end-to-end framework and fully utilizes the information provided by the teacher model. Specifically, we traverse the pseudo labels after confidence filtering to evaluate the reliability of each label generated by teacher. This labels generated by teacher are regarded as real pseudo labels. In this process, deal with the confidence of each pseudo box without considering background boxes. The processing function is shown in the Figure 3:

when the confidence is close to the confidence threshold setting, the confidence is reduced to a lower value through the function, at this time when the confidence is already high, we take the current confidence as the evaluation standard of prospect reliability in (4), so that the classification loss function of unsupervised part is as (5). Liu et al.[41] have proved by a large number of experiments that selecting 0.7 as the confidence threshold in the two-stage network will achieve better results. Therefore, we select $r_{i}$=0.7 to make the pseudo boxes’ classification score approach the two-stage network. When the confidence is low, a penalty item will be added to the loss calculation of the pseudo box which is similar to the hard negative sample mining method. We find that this method makes the classification score after a long iteration higher than without ADSO as shown in Figure 5.

where $L_{\mathrm{unsup}}^{cls}$ is unsupervised classification loss, $x$ indicates the classification score of the i-th box, $\omega_{i}$ indicates the classification score after function transformation as (4).

Different from the classification quality evaluation of pseudo labels, the regression quality of pseudo boxes is a difficult index to evaluate. We visualized the pseudo labeled images of the teacher model in the training process, As shown in Figure 6. Because the prediction accuracy of the teacher model is not high enough, multiple pseudo boxes will be predicted around ground truth boxes of the unlabeled images. Most of these pseudo boxes can not provide reliable and accurate positioning information for student but after merging these pseudo boxes, the new box location is closer to the ground truth box. Therefore, we introduce Fusion Box module to reduce the inaccurate location impact of pseudo boxes. Specifically, for an unlabeled image predicted by teacher model, after confidence filtering, Fusion Box module select whether to synthesize each other through the similarity $\xi$ and the pseudo algorithm of Fusion Box is as Algorithm 1.

Fusion Box collects the euclidean distance of the center point between each other for all the pseudo boxes generated in each iteration, and then calculates the similarity between each other. The two boxes whose similarity is less than the threshold $\mu$ are combined into a new pseudo box. The calculation process of the similarity $\xi$ of these pseudo boxes can be expressed by (6).

where$\text{ mean(Bbox }[\mathrm{n}])$ indicates that the center coordinates of each pseudo box. $\operatorname{scale}(w)$ and $\operatorname{scale}(h)$ represent the width and height of the whole image scaled after data augmentation. After Fusion Box module, the unsupervised regression loss function is calculated as follows:

where $L_{\mathrm{unsup}}^{reg}$ is unsupervised regression loss, $N_{pos_{\text{fusion }}}$ indicates the number of pseudo boxes after confidence filtering and Fusion Box, $\theta_{t}^{*}$ represents the regression target after Fusion Box.

As shown in Figure 8, we design the Dex encoder to replace the dilated encoder in YOLOF. Dilated encoder stacks standard convolution and dilated convolution, then merge the original feature to the feature map containing expanded receptive field. For reducing the parameter quantity, we cancel the dilated block with dilation rate = 2 in Dex encoder and choose to add 1×1 convolution after Deformable convolution v2 (DCNv2) module[44], the convolution layer reduces channels to 128 and maintains the reduced number of channels throughout the network.

Data augmentation need to be used for the input images in SSOD. The unlabeled images after different data augmentation will be collected to mark pseudo labels in the unsupervised section. Inputting these deformed images to network will lead to coordinates and angle changes in the ground truth box corresponding to Figure 7 while the deformed box will seriously affect the network training. Compared with the traditional fixed window convolution, deformable convolution [43] can effectively deal with gtbox deformation including box movement, size scaling and rotation, because its local receptive field is learnable and oriented to the whole image. Deformable convolution is to increase the spatial sampling position of additional offset and does not need additional supervision. Therefore, we choose to apply DCNv2 to YOLOF.

As shown in Figure 9, (a) is fixed window convolution, (b) is deformation convolution, and saturation color points is the actual sampling position of convolution kernel, which is offset from the standard position. (c) is dilated convolution which can be regarded as the special form of deformation convolution, has the ability to expand the receptive field. The convolution kernel used to generate the output feature and the deformation convolution kernel used to generate the offset are synchronously learned in DCNv2 while the offset is obtained by back propagation using interpolation algorithm.

We validate our method on the thighbone fracture dataset[45] which was collected from the radiology department of Linyi people’s hospital. All X-ray images are produced by the latest digital radiography (DR) technology. The dataset consists of 3842 thigh fracture images in 24 bit JPG format. The train set contains 3484 labeled images. In addition, The test set of 358 images is provided for serious performance verification. The following settings are adopted for performance verification:

Referring to the verification method of STAC[19], we use 1%, 5% and 10% of trainset as labeled training data, and unselected images in trainset as unlabeled data. In addition, our knowledge distillation result is compared with the supervised thighbone fracture detection methods proposed by Guan et al.[45] and Wang et al.[46] on this dataset to prove our advantage, following the convention to report the performance of val with mAP and AP50 as evaluation indicators.

The experiment was conducted on 4 NVIDIA geforce RTX 1080ti. We use Resnet-50 as feature extraction network to compare with previous methods. The backbone is initialized on ImageNet with pre-trained weight. Our implementation and hyperparameters are based on mmdetection[47], anchors with 5 scales and 1 aspect ratios are used. In the comparison between using partial labeled data sets and being compared with supervised algorithms, the training parameters under the two settings are slightly different because there is a large discrepancy in the amount of labeled training data.

Partial annotation data. the model is trained on 4 GPUs for 60K iterations, and each GPU has 8 images. During SGD training, the initial learning rate is 0.02, and the learning rate of backbone part is set to 0.005, which is divided by 10 in the 1k-th iteration and 1.5k-th iteration. The warmup phase is delayed to 1500iter. The weight attenuation and momentum are set to 0.0001 and 0.9, respectively. The confidence threshold is set to 0.5, the data sampling ratio SR is set to 0.25, and gradually decreases to 0 in the last 1K iterations.

Compared with supervised algorithms. the data sampling ratio SR is set to 1, other settings are the same as the experience of partial annotation data.

In order to estimate the reliability of Fusion Box, the Fusion Box threshold $\mu$ is selected as 0.05 according to the size of the input image and select the fused pseudo box for box regression. In addition, we use weak data augmentation for teacher model, strong data augmentation for student model as shown in Table 1. (Unlabeled (T) and Unlabeled (S) respectively represent the data augmentation used for the unlabeled images input to the teacher and student models while Labeled represents the data augmentation used for labeled images. Weak Aug and Strong Aug respectively represent the weak data augmentation and strong data augmentation.)

In this section, we compare our method with other state-of-the-art methods proposed in recent years. We first evaluate on Partial annotation data setting and compare our method with the results in STAC, Soft teacher and Unbiased teacher. Also the baseline of YOLOF and FasterRcnn is compared with our semi-supervised method in Table 2. We found that our method performed better than other methods in fracture detection. Specifically, the mAP of our method is 15.7%, 22.5% and 11.0% higher than the baseline on the 1%, 5% and 10% data sets, at the same time 2.4%, 0.2% and 2.3% higher than the previous methods. We evaluate the loss of model training and show the results in Figure 10. In order to compare the prediction results more intuitively, we visualized the baseline prediction using YOLOF and the prediction results of our semi-supervised framework in Figure 11.

Then we compared our performance of SSL framework on full dataset and 50% labeled dataset in Table 3 and Table 4. We only use 50% data for semi-supervised training, and the AP exceeds Cascade r-cnn, a two-stage network which parameters are much more than YOLOF. By using Dex encoder, we only need 3 divided blocks to increase the mAP of YOLOF by 1.6% and AP75 by 2.6%, which proves the effectiveness of Dex encoder. In addition, using our semi-supervised framework to distill the knowledge of the single-stage network can further improve 2.6% AP50 and 9.6% AP75. Compared with the algorithm for fracture detection and the current popular object detection network, our semi-supervised framework has higher accuracy. In addition, by introducing our semi-supervised framework, our model achieves the same performance as the fully supervised training of two-stage detectors with simpler network structure and fewer parameters. We also compare the model inference speed on test set and moel complexity with the latest two methods in Table 7.

In this section, we validate key designs. All ablation experiments are performed on YOLOF using 10% labeled dataset.

The influence of critical design in semi-supervised model. Table 5 shows the impact of different modules on the performance of the model. When all three modules are adopted, the model performance can reach the best. Compared with the original SSL model, mAP is increased by 3.3% and AP50 by 2.4%. The validity of modules is proved.

The impact of confidence filter threshold. We select different confidence thresholds to compare the performance of our semi-supervised model in Table 8. When we select a threshold of 0.5, the performance of the model reaches the best. When we select a threshold of 0.5 ± 0.1, a high threshold will lead to a faster decline in the mAP, indicating that our semi-supervised model is sensitive to a high threshold. However, whether we select a high threshold or a low threshold, the accuracy of the model does not decline much, indicating that ADSO module has a certain adjustment function to the confidence threshold.

The impact of dilated block quantity.We compared the effect of different number of expansion blocks in Table 6. Considering the balance between accuracy and parameter quantity, we finally choose to use 3 dilated blocks on the Dex encoder.

The impact of Fusion Box threshold. In Table 9, we study the box regression Fusion Box threshold. The best performance is achieved when the threshold is set to 0.05.