Terahertz image breast tumor detection method based on SAM and curriculum learning

Abstract

The present invention discloses a terahertz image breast tumor detection method based on SAM and curriculum learning. In order to accurately locate and identify tumors from terahertz breast images, the present invention applies the visual basic model SAM to tumor detection and constructs a dual-branch tumor detection network, in which one branch uses the target detection model DETR to perform coarse-grained tumor detection on terahertz breast images and generate mixed prompts, including bounding boxes and categories, where the categories are used as text prompts, and then input them into the SAM branch for fine-grained tumor detection. In addition, the present invention adopts the idea of ​​curriculum learning and sends terahertz breast tumor image data sets from easy to difficult into the dual-branch tumor detection network in batches for training to obtain better detection results.

Description

Technical Field

[0001] The present invention relates to computer vision processing and medical image processing technology, and in particular to a terahertz image breast tumor detection method based on SAM and curriculum learning. Background Art

[0002] The task of target detection is to find all the targets (objects) of interest in an image and determine their categories and locations. It is one of the core issues in the field of computer vision. It has a wide range of applications, including medical image analysis, face recognition, smart transportation, industrial inspection, augmented reality and virtual reality.

[0003] In the past few decades, target detection has made great progress. Before deep learning was introduced into this field, traditional target detection ideas included region selection, manual feature extraction, and classifier classification. Since manual feature extraction methods often find it difficult to meet the diverse characteristics of targets, traditional methods have never been able to solve the target detection problem well. After the rise of deep learning, neural networks can automatically learn powerful feature extraction and fitting capabilities from a large amount of data, and thus many high-performance target detection algorithms have emerged. Target detection methods based on deep learning can be roughly divided into three categories: two-stage target detection, single-stage target detection, and Tansformer-based target detection.

[0004] The two-stage object detection algorithm first extracts candidate boxes based on the image, and then performs secondary correction based on the candidate area to obtain the detection result. The detection accuracy is high, but the detection speed is slow. The pioneer of this type of algorithm is R-CNN, followed by Fast R-CNN and Faster R-CNN. Due to its outstanding performance, Faster R-CNN is still a very competitive algorithm in the field of object detection. Subsequently, algorithms such as FPN and Mask R-CNN proposed improvements to the shortcomings of Faster R-CNN, which further enriched the components of Faster R-CNN and improved its performance.

[0005] Two-stage detectors treat detection as a classification problem: a module is required to enumerate candidate boxes that are classified as foreground or background by the network. This type of algorithm is relatively slow because it requires multiple runs of the detection and classification process. The single-stage detection method only needs to be fed into the network once to predict all bounding boxes, so it is faster and very suitable for mobile devices. The most typical single-stage detection algorithms include YOLO, SSD, SqueezeDet, and DetectNet.

[0006] Transformer has been a hot topic in recent years. It was first applied in the field of NLP. According to research data, research in the fields of CV, speech and video, and multimodality has been published very rapidly in recent years. Among them, Transformer in the field of CV is generally referred to as Vision Transformer, or ViT for short. At present, the target detection algorithm based on Transformer is mainly developed in the series such as DETR and ViT. Among them, the Transformer detection model extended by DETR mainly improves and expands the object query (adding prior knowledge), attention mechanism (sparseness, focusing on meaningful areas) and label assignment (OTA, etc.) mechanism, as well as feature matching and allocation mechanism.

[0007] In my country, the incidence of tumors has remained high, and early diagnosis of tumors is very important for the prognosis of patients. Therefore, research on accurate screening of tumors has great theoretical and application value. However, according to statistics, 2% of cancer patients are caused by X-rays and CT equipment during medical treatment. Terahertz (THz) waves will not cause harmful ionization damage to biological tissues. At the same time, terahertz imaging has the characteristics of high imaging resolution, low ionization, and high penetration. It is a very potential conventional detection technology for early tumor cancer. At present, terahertz wave CT imaging technology is in its early stages, and the resulting terahertz wave CT images have many problems such as unknown spots, local blur, and artifacts. There is a lack of intelligent detection technology for tumor tissues in terahertz wave CT images at home and abroad. Therefore, the use of terahertz waves to detect breast tumors is an important research issue in the field of computer-aided diagnosis. Summary of the invention

[0008] Purpose of the invention: The purpose of the present invention is to solve the deficiencies in the prior art and to provide a terahertz image breast tumor detection method based on SAM and curriculum learning.

[0009] Technical solution: A terahertz image breast tumor detection method based on SAM and curriculum learning of the present invention comprises the following steps:

[0010] Step 1, obtain a terahertz image I with a dimension of 3×H×W, and input the terahertz image I into the target detection DETR branch and the SAM mask prediction branch at the same time;

[0011] The execution process of the target detection DETR branch is as follows:

[0012] First, the terahertz image I is subjected to feature extraction via a convolutional neural network (ResNet50) to obtain a feature map F i1 ; Then, the generated position code F posWith the feature map F i1 Add bit by bit to get the feature map F i2 , then, the feature map F i2 Reshape and input to Transformer encoder Enc, output feature map F eo ; Secondly, construct N learnable embedding vectors F l , and embed the vector F l and feature map F eo Input Transformer decoder Dec together and output feature map F do Finally, two parallel fully connected layers FC1 and FC2 that do not share weights are used to transform the feature map F do Mapped into two outputs: bounding box prediction and classification prediction;

[0013] The SAM mask prediction branch includes an image encoder Enc I , prompt encoder Enc p and mask decoder Dec m , the execution process of the SAM mask prediction branch is as follows:

[0014] The image encoder Enc I Using the MAE pre-trained visual transformer (ViT), the terahertz image I is passed through the image encoder Enc I After processing, the image is embedded in I emb ; The prompt encoder Enc p Receive the prediction results of the object detection DETR branch, i.e., bounding box prediction and classification prediction, and perform hint encoding to obtain hint embedding P emb ; Then, embed the image into I emb and prompt embedding P emb Input to the mask decoder Dec m , complete the mask prediction Mask of the tumor, convert the mask prediction Mask and the terahertz image I into the final predicted tumor bounding box Box through process, and the tumor classification result still uses the prediction Cls of the DETR branch p As the final classification prediction Cls;

[0015] Step 2: Using the idea of ​​course learning, the terahertz breast tumor image datasets from easy to difficult are trained in batches to increase the generalization and robustness of the model. The predicted bounding boxes and categories are compared with the annotated true values ​​for loss calculation, and 100 rounds of supervised learning are performed. The loss function is used to constrain the network for optimization.

[0016] Furthermore, the specific processing method of the target detection DETR branch is as follows:

[0017] Step A: The terahertz image I is first subjected to a convolutional neural network (ResNet50) to extract features. After five downsamplings, the H and W of the feature map are reduced to 1 / 2 of the original size each time, and the output feature map has a dimension of 2048×H / 32×W / 32. The 1×1 convolutional layer is then used to reduce the dimension of the feature map output by the convolutional neural network to 256, which is recorded as feature map F. i1 , whose dimension is 256×H / 32×W / 32; then generate a positional code F of size 256×H / 32×W / 32 pos , and compare it with the input feature map F i1 After bitwise addition, the dimension is still 256×H / 32×W / 32, which is recorded as feature map F. i2 ;

[0018] Feature map F i2 The calculation formula is as follows:

[0019] F i2 =ResNet50(I)+F pos ;

[0020] Step B: First, the feature map F i2 The scale size is reshaped to (H / 32×W / 32)×256, and then input to the Transformer encoder Enc, which outputs the feature map F of the same size after passing through the encoder Enc eo , feature map F eo The dimension is (H / 32×W / 32)×256; then construct N (N=100) learnable embedding vectors F of dimension 256 l , the learnable embedding vector F l and feature map F eo As the input of Transformer decoder Dec, after the decoder, the output feature map F of size 100×256 do ;

[0021] The feature F do The calculation formula is as follows:

[0022] F dpi =Dec(Enc(F i2 ),F l );

[0023] Step C: Use two parallel fully connected layers FC1 and FC2 that do not share weights to transform the output F do Mapped into two outputs: bounding box prediction and classification prediction. The classification prediction result dimension is 100×2, and the bounding box prediction is used for position regression with a dimension of 100×4.

[0024] Among them, the classification target Cls p Including benign tumors and malignant tumors, the target Box of the bounding box p Refers to the normalized center coordinates and width and height of the bounding box;

[0025] Cls p and Box p The calculation formula is as follows: Cls p =FC1(F do ); Box p =FC2(F do ).

[0026] Furthermore, the SAM mask predicts the branch

[0027] The calculation formula of the tumor mask prediction Mask is as follows:

[0028] Mask=Dec m (Enc I (I),Enc p (Cls p ,Box p ));

[0029] Box = process(Mask, I);

[0030] Cls=Cls p .

[0031] Furthermore, the loss function in step 2 is calculated as follows:

[0032] L=L cls +L reg

[0033] Among them, L cls Represents the classification loss function, that is, the loss between the target category predicted by the network and the actual target category, L reg represents the bounding box loss function, which is the loss between the target location predicted by the network and the actual target location;

[0034]

[0035] L reg = 1-IoU;

[0036] In the above formula, N is the total number of samples, y i is the category to which the i-th sample belongs, p iis the predicted value of the i-th sample; the loss function of the IoU class is based on the IoU (intersection over union) between the predicted box and the labeled box. The predicted box is P and the labeled box is G. The corresponding IoU is expressed as: IoU = P∩G / P∪G; that is, the ratio of the intersection and union of the two boxes.

[0037] Beneficial effects: Compared with the prior art, the present invention has the following advantages:

[0038] (1) The present invention combines the visual basic model SAM and the target detection model DETR to construct a new dual-branch tumor detection network, which is conducive to the accurate positioning and identification of tumors in terahertz images.

[0039] (2) The present invention uses DETR to provide specific automatic prompts for SAM, which will further enhance SAM's mask prediction capability and ultimately achieve better tumor detection results.

[0040] (3) The present invention uses the idea of ​​curriculum learning during model training, which is beneficial to enhancing the robustness and generalization of the model. BRIEF DESCRIPTION OF THE DRAWINGS

[0041] Figure 1 It is a schematic diagram of the overall process of the present invention;

[0042] Figure 2 Schematic diagram of the model framework in an embodiment of the present invention.

[0043] Figure 3 Schematic diagram of terahertz breast tumor image in an embodiment of the present invention. DETAILED DESCRIPTION

[0044] The technical solution of the present invention is described in detail below, but the protection scope of the present invention is not limited to the embodiments.

[0045] like Figure 1 As shown, a terahertz image breast tumor detection method based on SAM and curriculum learning of the present invention comprises the following steps:

[0046] Step 1. Obtain a terahertz image I with a dimension of 3×H×W, and input the terahertz image I into the target detection DETR branch and the SAM mask prediction branch at the same time; through the dual-branch parallel tumor detection network, a more fine-grained tumor detection is performed to overcome the limitations of the existing tumor detection model and accurately locate the position of the tumor in the image.

[0047] The execution process of the target detection DETR branch in this embodiment is as follows:

[0048] First, the terahertz image I is subjected to feature extraction via a convolutional neural network (ResNet50) to obtain a feature map F i1; Then, the generated position code F pos With the feature map F i1 Add bit by bit to get the feature map F i2 , then, the feature map F i2 Reshape and input to Transformer encoder Enc, output feature map F eo ; Secondly, construct N learnable embedding vectors F l , and embed the vector F l and feature map F eo Input Transformer decoder Dec together and output feature map F do Finally, two parallel fully connected layers FC1 and Fc2 that do not share weights are used to transform the feature map F do Mapped into two outputs: bounding box prediction and classification prediction.

[0049] The above SAM mask prediction branch includes the image encoder Enc I , prompt encoder Env p and mask decoder Dec m The image encoder uses a visual transformer (ViT) pre-trained on MAE and minimizes it to handle high-resolution inputs. The hint encoder can handle dense (masks) and sparse (points, boxes, text) input features. The mask decoder can effectively map image embeddings, hint embeddings, and output tags to masks. The execution process of the SAM mask prediction branch is as follows:

[0050] The image encoder Env I Using the MAE pre-trained visual transformer (ViT), the terahertz image I is passed through the image encoder Enc I After processing, the image is embedded in I emb ; The prompt encoder Enc p Receive the prediction results of the object detection DETR branch, i.e., bounding box prediction and classification prediction, and perform hint encoding to obtain hint embedding P emb ; Then, embed the image into I emb and prompt embedding P emb Input to the mask decoder Dec m , complete the mask prediction Mask of the tumor, convert the mask prediction Mask and the terahertz image I into the final predicted tumor bounding box Box through process, and the tumor classification result still uses the prediction Cls of the DETR branch p As the final classification prediction CLs.

[0051] SAM mask predicts branch

[0052] The calculation formula of the tumor mask prediction Mask is as follows:

[0053] Mask=Dec m (Enc I (I),Enc p (Cls p ,Box p ));

[0054] Box = process(Mask, I);

[0055] Cls=Cls p .

[0056] Step 2: Using the idea of ​​course learning, the terahertz breast tumor image datasets from easy to difficult are trained in batches to increase the generalization and robustness of the model. The predicted bounding boxes and categories are compared with the annotated true values ​​for loss calculation, and 100 rounds of supervised learning are performed. The loss function is used to constrain the network for optimization.

[0057] The present invention applies the visual basic model SAM to breast tumor detection in terahertz images, and uses it as one branch of the dual-branch tumor detection network to give full play to its generalization ability on new data sets and tasks, effectively solving the limitations of existing tumor detection models and the difficulty in accurately locating the position of tumors in images. At the same time, the target detection model DETR is used as the other branch of the dual-branch tumor detection network to generate specific automatic prompts for SAM, including bounding boxes and categories. Using this coarse-grained detection result as a prompt will further improve the fine-grained mask prediction effect of SAM.

[0058] The specific processing method of the target detection DETR branch in this embodiment is as follows:

[0059] Step A: The terahertz image I is first subjected to a convolutional neural network (ResNet50) to extract features. After five downsamplings, the H and W of the feature map are reduced to 1 / 2 of the original size each time, and the output feature map has a dimension of 2048×H / 32×W / 32. The 1×1 convolutional layer is then used to reduce the dimension of the feature map output by the convolutional neural network to 256, which is recorded as feature map F. i1 , whose dimension is 256×H / 32×W / 32; then generate a positional code F of size 256×H / 32×W / 32 pos , and compare it with the input feature map F i1 After bitwise addition, the dimension is still 256×H / 32×W / 32, which is recorded as feature map F. i2 ; Feature map F i2 The calculation formula is as follows:

[0060] F i2 =ResNet50(I)+F pos;

[0061] Step B: First, the feature map F i2 The scale size is reshaped to (H / 32×W / 32)×256, and then input to the Transformer encoder Enc, which outputs the feature map F of the same size after passing through the encoder Enc eo , feature map F eo The dimension is (H / 32×W / 32)×256; then construct N (N=100) learnable embedding vectors F of dimension 256 l , the learnable embedding vector F l and feature map F eo As the input of Transformer decoder Dec, after the decoder, the output feature map F of size 100×256 do ; Feature F do The calculation formula is as follows:

[0062] F do =Dec(Enc(F i2 ),F l );

[0063] Step C: Use two parallel fully connected layers FC1 and Fc2 that do not share weights to transform the output F do Mapped into two outputs: bounding box prediction and classification prediction. The classification prediction result dimension is 100×2, and the bounding box prediction is used for position regression with a dimension of 100×4.

[0064] Among them, the classification target Cls p Including benign tumors and malignant tumors, the target Box of the bounding box p Refers to the normalized center coordinates and width and height of the bounding box;

[0065] Cls p and Box p The calculation formula is as follows: Cls p =FC1(F do ); Box p =FC2(F do ).

[0066] The terahertz breast tumor image (such as Figure 3 As shown in the figure, the dataset is used in batches from easy to difficult for the network training process to improve the robustness and generalization of the model. The loss function of the training process is calculated as follows:

[0067] L=L cls +L reg

[0068] Among them, Lcls Represents the classification loss function, that is, the loss between the target category predicted by the network and the actual target category, L reg represents the bounding box loss function, which is the loss between the target location predicted by the network and the actual target location;

[0069]

[0070] N is the total number of samples, y i is the category to which the i-th sample belongs, p i is the predicted value of the i-th sample; L cls Represents the classification loss function, that is, the loss between the target category predicted by the network and the actual target category;

[0071] The loss functions of the IoU class are all based on the IoU (intersection over union) between the predicted box and the labeled box. Let the predicted box be P and the labeled box be G. The corresponding IoU can be expressed as:

[0072] IoU = P∩G / P∪G;

[0073] That is, the ratio of the intersection and union of two boxes. IoU Loss is defined as: L reg =1-IoU.

[0074] To facilitate quantitative evaluation, this embodiment adopts three widely used evaluation indicators.

[0075] (1) Average accuracy mAP:

[0076]

[0077] Among them, AP i It represents the average precision of the i-th category, N is the total number of categories, and AP (Average precision) is the evaluation indicator of the mainstream object detection model.

[0078] (2) AP50: mAP in (1) refers to the mean AP when the IOU value is increased from 50% to 95% in steps of 5%. AP50 refers to the accuracy when the IOU is 50%.

[0079] (3) AP75: AP75 refers to the accuracy when the IOU is 75%.

Claims

1. A terahertz image breast tumor detection method based on SAM and curriculum learning, characterized in that: The following steps are involved: Step 1, obtain a terahertz image I with a dimension of 3×H×W, and input the terahertz image I into the target detection DETR branch and the SAM mask prediction branch at the same time; The execution process of the target detection DETR branch is as follows: First, the terahertz image I is subjected to feature extraction by the convolutional neural network ResNet50 to obtain the feature map F i1 ; Then, the generated position code F pos With the feature map F i1 Add bit by bit to get the feature map F o2 , then, the feature map F i2 Reshape and input to Transformer encoder Enc, output feature map F eo ; Secondly, construct N learnable embedding vectors F l , and embed the vector F l and feature map F eo Input Transformer decoder Dec together and output feature map F do Finally, two parallel fully connected layers FC1 and FC2 that do not share weights are used to transform the feature map F do Mapped into two outputs: bounding box prediction and classification prediction; The SAM mask prediction branch includes an image encoder Enc I , prompt encoder Enc p and mask decoder Dec m , the execution process of the SAM mask prediction branch is as follows: The image encoder Enc I Using the MAE-based pre-trained visual converter ViT, the terahertz image I is passed through the image encoder Enc I After processing, the image is embedded in I emb ; The prompt encoder Enc p Receive the prediction results of the target detection DETR branch, i.e., bounding box prediction and classification prediction, and perform hint encoding to obtain hint embedding P emb ; Then, embed the image into I emb and prompt embedding P emb Input to the mask decoder Dec m , complete the mask prediction Mask of the tumor, convert the mask prediction Mask and the terahertz image I into the final predicted tumor bounding box Box through process, and the tumor classification result still uses the prediction Cls of the DETR branch p As the final classification prediction Cls; Step 2: Using the idea of ​​course learning, the terahertz breast tumor image datasets from easy to difficult are trained in batches to increase the generalization and robustness of the model. The predicted bounding boxes and categories are compared with the annotated true values ​​for loss calculation, and 100 rounds of supervised learning are performed. The loss function is used to constrain the network for optimization.

2. The terahertz image breast tumor detection method based on SAM and curriculum learning according to claim 1 is characterized in that: The specific processing method of the target detection DETR branch is as follows: Step A: The terahertz image I is first subjected to the convolutional neural network ResNet50 to extract features. After five downsamplings, the H and W of the feature map are reduced to 1 / 2 of the original size each time, and the output feature map has a dimension of 2048×H / 32×W / 32. Then, a 1×1 convolutional layer is used to reduce the dimension of the feature map output by the convolutional neural network to 256, which is recorded as feature map F. i1 , whose dimension is 256×H / 32×W / 32; then generate a positional code F of size 256×H / 32×W / 32 pos , and compare it with the input feature map F i1 After bitwise addition, the dimension is still 256×H / 32×W / 32, which is recorded as feature map F. i2 ; Feature map F i2 The calculation formula is as follows: F i2 =ResNet50(I)+F pos ; Step B: First, the feature map F i2 The scale size is reshaped to (H / 32×W / 32)×256, and then input to the Transformer encoder Enc, which outputs the feature map F of the same size after passing through the encoder Enc eo , feature map F eo The dimension is (H / 32×W / 32)×256; then construct N learnable embedding vectors F of dimension 256 l , the learnable embedding vector F l and feature map F eo As the input of Transformer decoder Dec, after the decoder, the output feature map F of size 100×256 do ; N = 100; The feature f do The calculation formula is as follows: F do =Dec(Enc(F i2 ),F l ); Step C: Use two parallel fully connected layers FC1 and FC2 that do not share weights to transform the output F do Mapped into two outputs: bounding box prediction and classification prediction. The classification prediction result dimension is 100×2, and the bounding box prediction is used for position regression with a dimension of 100×4. Among them, the classification target Cls p Including benign tumors and malignant tumors, the target Box of the bounding box p Refers to the normalized center coordinates and width and height of the bounding box; Cls p and Box p The calculation formula is as follows: Cls p =FC1(F do ); Box p =FC2(F do ).

3. The terahertz image breast tumor detection method based on SAM and curriculum learning according to claim 1 is characterized in that: The SAM mask predicts branches The calculation formula of the tumor mask prediction Mask is as follows: Mask=Dec m (Enc I (I),Enc p (Cls p ,Box p )); Box = process(Mask, I); Cls=Cls p 。 4. The terahertz image breast tumor detection method based on SAM and curriculum learning according to claim 1 is characterized in that: The loss function in step 2 is calculated as follows: L=L cls +L reg ; Among them, L cls Represents the classification loss function, that is, the loss between the target category predicted by the network and the actual target category, L reg represents the bounding box loss function, which is the loss between the target location predicted by the network and the actual target location; N is the total number of samples, y i is the category of the i-th sample, p i is the predicted value of the i-th sample; L reg =1-IoU; In the above formula, the loss function of the IoU class is based on the intersection-over-union ratio between the predicted box and the labeled box. Let the predicted box be P and the labeled box be G. The corresponding IoU is expressed as: IoU = P∩G / P∪G.

Images