Dental CBCT three-dimensional tooth rapid labeling method based on SAM large model

Abstract

The invention discloses a dental CBCT three-dimensional tooth rapid labeling method based on an SAM large model. According to the method, an improved SAM network is utilized to distinguish teeth from periodontal membrane and alveolar bone in a CBCT image, and the method comprises the following specific steps: stage 1, conversion and processing from tooth data to trainable data; stage 2, constructing an improved SAM network model; stage 3, training a network model; and a fourth stage of tooth labeling and three-dimensional visualization based on CBCT data. According to the invention, the Lora fine tuning technology is used to help the SAM network adapt to the image characteristics of the medical image data set, so that the defect that the SAM network is only trained on a natural image originally is overcome. In addition, a CNN bypass network is adopted to make full use of spatial information in a three-dimensional slice and help capture edge information in the slice. And finally, a previous prediction result is ingeniously utilized to help to generate a target box prompt to give full play to a prompt mechanism in the SAM network. According to the method, a better tooth labeling effect can be obtained when the same small amount of labels are used for training.

Description

Technical Field

[0001] The present invention belongs to the intersection of dental medicine and image processing, mainly involves 3D image segmentation and classification, and proposes a dental CBCT 3D tooth fast annotation method based on the SAM large model. The present invention uses a deep learning semantic segmentation model to distinguish teeth from the periodontal membrane and alveolar bone in the cone beam computed tomography (CBCT) images of the dental oral cavity, so that the morphology, position and relationship of the teeth with the surrounding structures can be analyzed in detail. Background Art

[0002] Dental CBCT (cone beam computed tomography) in medical imaging technology can obtain three-dimensional CT images of teeth, which has many important meanings in the oral field. CBCT provides three-dimensional images, which can show comprehensive information of oral structures such as teeth, pulp, root, alveolar bone, etc., thus providing more accurate diagnosis basis for clinicians. Using three-dimensional tooth segmentation technology to segment the tooth area can provide key information for the diagnosis of dental diseases, treatment plan formulation, oral surgery simulation, etc.

[0003] Traditional tooth segmentation methods are usually based on manually designed features and rules. These methods often require a lot of manual intervention and are not adaptable to complex and variable tooth shapes. With the development of deep learning, deep learning methods based on convolutional neural networks have achieved remarkable success in image segmentation tasks and have also been applied to 3D tooth segmentation tasks. However, deep learning methods usually require a large amount of labeled data for training, and obtaining a large amount of accurately labeled tooth segmentation data is a time-consuming task that requires expertise.

[0004] SAM (Segment Anything) is the first large-scale model network structure for visual images released by the Meta Research team in 2023. It is designed for prompt-based segmentation and can adapt to new image distributions and tasks with zero samples. After its release, SAM, as a pioneering visual foundation model, has promoted rapid application in different fields with its zero-sample segmentation capability and prompt-based approach. However, in the field of medical imaging, although SAM has a fairly high accuracy in segmenting objects in natural images, it has difficulties in dealing with objects outside its training field, such as medical images.

[0005] The PEFT (Parameter-Efficient Fine-Tuning) method first appeared in the field of natural language processing. Since the cost of adjusting all parameters of large model networks in the field of natural language processing is too high, and if the training data is sparse, overfitting may occur, which will make the effect worse. Therefore, the PEFT method aims to improve the performance of pre-trained models on new tasks by minimizing the number of fine-tuning parameters and computational complexity, thereby alleviating the training cost of large pre-trained models, and the required training data is relatively small.

[0006] Since the labeling of three-dimensional dental CBCT images requires high professional knowledge and the labeling process is extremely time-consuming, there is less 3D dental CBCT labeling data. Therefore, fine-tuning and modification based on the SAM network is adopted to minimize the number of three-dimensional CTs required for training while achieving better segmentation prediction results. Summary of the invention

[0007] The purpose of the present invention is to provide a dental CBCT three-dimensional tooth rapid annotation method based on a SAM large model in view of the deficiencies in the prior art.

[0008] The technical solution provided by the present invention uses the improved SAM network to segment CBCT tooth data. The specific steps are as follows:

[0009] Phase 1: Conversion and processing of dental data into trainable data;

[0010] Phase 2: Building an improved SAM network model;

[0011] Phase 3: Network model training;

[0012] Phase 4: Tooth annotation and 3D visualization based on CBCT data;

[0013] Furthermore, the specific implementation steps of the conversion and processing of the stage 1 tooth data into trainable data are as follows:

[0014] 1-1. Compress the original CBCT dental data into nii.gz format, divide it into training set, validation set and test set, and perform binary operation on the labels.

[0015] 1-2. Read in the training set, resize the image to 400*512*512 pixels, then perform intensity normalization (linearly convert the intensity value of the image from the specific range of the CT scan (such as from -1024 to 1252) to the range of [0.0, 1.0]), random cropping (randomly crop the image and label according to the given hyperparameters (256, 256, 256). This step balances the position of positive and negative samples in the label to generate diverse training samples, which helps to improve the generalization ability of the model), random affine transformation, random intensity offset and a series of preprocessing operations.

[0016] Furthermore, the second stage improves the SAM network model, and the specific implementation steps are as follows:

[0017] In order to target dental data, a CNN bypass network is added to the original SAM network model to supplement the extraction of spatial information of adjacent slices and improve the edge capture capability of dental details. At the same time, a Lora fine-tuning module is added to the Encoder part of the SAM network model to adapt to medical image information and make full use of the original SAM pre-training data. Finally, the target box information generated by the adjacent previous slice is used as the current slice input to the Prompt Encoder network. The overall network structure is as follows: Figure 1 shown.

[0018] The CNN bypass network mainly includes point-by-point convolutional layers (with 3 filters) and CNN networks (using the CNN architecture in traditional ResNet), and sends the features extracted by CNN together with the transformer layer in the corresponding SAM Encoder to the channel fusion module (hereinafter referred to as CLI). The channel fusion module contains a combination of four 1*1 convolutional layers, a ReLU layer, and a 1*1 convolutional layer. Its overall structure is as follows: Figure 2 shown.

[0019] Among them, the Lora fine-tuning strategy is mainly used in each Transformer Encoder layer in the SAM Encoder, and a LoRA module is added to the q and v projection layers before the multi-head attention layer to fine-tune the network.

[0020] Assume that the output of the network layer under the original pre-trained weights is as follows:

[0021] z=W0x

[0022] The formula for Lora's fine-tuning strategy is as follows:

[0023] z'=W0x+BAx

[0024] Among them, z is the original output, z' is the output after fine-tuning, W0 represents the weight matrix after pre-training, and B and A are low-rank matrices respectively.

[0025] In the process of looping through the number of channels of the training data, the target frame information is generated by predicting the adjacent previous slice, and the prediction result in the adjacent previous channel is saved. The prediction result is sent to the target frame generation module to obtain the target frame information (that is, the pixel points at the lower left corner and the upper right corner of the prediction result are expanded by ten pixels each to obtain the target frame of the slice, and the target frame is expressed in the form of two coordinate points (lower left corner and upper right corner)), and the target frame information (that is, the two coordinate points) is sent to the Prompt Encoder network.

[0026] Furthermore, the specific implementation steps of the stage three network model training are as follows:

[0027] 3.1 During the training process, the training data is traversed and processed. For each set of training data, the image part in the training data is traversed according to the number of channels. In each loop, the two-dimensional slice under the current loop is taken out according to the number of channels. The slice image channel is copied three times using the repeat function to form a three-channel image (CHW=3*512*512) named input_sam_image. The two-dimensional slice of the label under the corresponding channel number is taken out together, and the zoom function is used to adjust its height and width to low_res, that is, 64 to obtain low_res_label (low_res is img_embedding_size*4, and img_embedding_size is a hyperparameter set in advance in the sam network structure). Then the input_image is sent to the SAM Encoder using the Lora fine-tuning strategy.

[0028] 3.2 In the above process of looping through the number of channels for training data, in addition to obtaining input_image_sam, the three slices between (c-2,h,w) and (c+2,h,w) (c is the current channel number. If c is the rightmost or leftmost, copy the edgemost slice to replace the part beyond the edge) are taken out and saved as a new image named input_image_cnn, and then input_cnn_image is input into the CNN bypass network in the sam encoder.

[0029] 3.3 In the process of looping through the number of channels of the training data, the prediction result in the previous channel (that is, the prediction result for the image (c-1, h, w)) is saved and sent to the target box generation module to obtain the target box information, and the target box information is sent to the Prompt Encoder network.

[0030] The Prompt Encoder module is part of the original SAM model, and is mainly used to process input prompts. These prompts can be user inputs in the form of points, boxes, or masks to tell the model which area or object needs to be segmented.

[0031] The SAM Mask Decoder module is part of the original SAM model and is mainly responsible for generating the final segmentation mask. Its main function is to receive the feature map extracted by the Image Encoder and the prompt information processed by the PromptEncoder, and combine these two parts of information to predict the segmentation mask of the target object.

[0032] 3.4 After passing through 12 layers of Transformer Encoder layers with Lora fine-tuning modules (i.e., after passing through the SAMEncoder part), the extracted features and the features extracted from the Prompt Encoder network are sent to the SAM MaskDecoder module to obtain a dictionary file, which contains the keywords "masks" and "low_res_logits". The obtained "low_res_logits" is a small-sized mask image, the size of which is proportionally reduced from the image size. "Masks" is the "low_res_logits" magnified by the SAM Mask Decoder module. The "low_res_logits" obtained by the low-resolution mask and the two-dimensional slice of the label under the corresponding channel of the input_sam_image are processed by the zoom function to obtain the lows_res_label (i.e., the lows_res_label in step 3.1) to calculate the loss function. The loss function adopts a combination of the cross entropy loss function and the dice loss function. The overall flow chart of model training is as follows: Figure 4 It is worth noting that in the SAM Encoder during the network training process, only the LoRa fine-tuning module participates in the training and performs the forward propagation process to learn parameters, while the rest use the original pre-training parameters without change. As for the SAM Mask Decoder module, CNN bypass network and prompt encoder module, they all perform forward propagation training normally.

[0033] Furthermore, the specific implementation steps of the stage 4 of tooth annotation and 3D visualization based on CBCT data are as follows:

[0034] After completing the network training, the unlabeled CBCT is predicted to obtain the prediction results. The test set is predicted to obtain the following results, which are compared with multiple original images and corresponding prediction images. Figure 4The method of the present invention is compared with other methods, both of which use the same training sample size, i.e., CBCT images of 5 patients, and the same number of iterations, i.e., 250 epochs. The comparison table is shown in FIG. Figure 5 In order to better demonstrate the effect of tooth annotation of the present invention, the prediction results are visualized in 3D using 3Dslicer software. The comparison effect is shown in Figure 6 As shown, it can be clearly seen in the red frame that the tooth marking result of the present invention can accurately separate the teeth.

[0035] Beneficial effects of the present invention:

[0036] The core technical feature of the present invention is to address the problem of small amount of CBCT data annotation so as to make full use of the visual large model SAM network, and make improvements on this basis, using Lora fine-tuning technology to help the SAM network adapt to the image characteristics of the medical image data set to make up for its shortcomings due to being trained only on natural images. In addition, a CNN bypass network is used to make full use of the spatial information in the 3D slices and help capture the edge information in the slices. Finally, the previous prediction result is cleverly used to help generate the target box prompt to give full play to the prompt mechanism in the SAM network. When the present invention uses the same small amount of annotations for training, better tooth annotation effects can be obtained.

[0037] The present invention utilizes an improved SAM network to distinguish teeth from periodontal ligament and alveolar bone in CBCT images, thereby being able to perform a detailed analysis of the shape and position of teeth and their relationship with surrounding structures.

[0038] Compared with existing patents and papers, the present invention can obtain better tooth annotation effect when using the same small amount of annotations for training. BRIEF DESCRIPTION OF THE DRAWINGS

[0039] Figure 1 This is the improved SAM network structure diagram.

[0040] Figure 2 This is the overall structure of the CNN bypass network.

[0041] Figure 3 This is the overall training flow chart of the network.

[0042] Figure 4 This is a 2D result display diagram.

[0043] Figure 5 This is a comparison chart of the result indicators of multiple methods.

[0044] Figure 6 3D preview results display for multiple methods. DETAILED DESCRIPTION

[0045] In order to make the purpose, technical solutions and advantages of the present invention more clearly understood, embodiments are described below.

[0046] The present invention aims at the limitations of existing CBCT image tooth segmentation technology in terms of efficiency and accuracy, and proposes a new method based on deep learning technology. The main challenges faced when using deep learning technology to process cone beam computed tomography (CBCT) images include the small amount of annotated data, the difficulty of manual annotation, and the limited ability to capture tooth details. The purpose of the present invention is to use an improved SAM network to distinguish teeth from the periodontal ligament and alveolar bone in CBCT images, so that the morphology, position and relationship of teeth with surrounding structures can be analyzed in detail.

[0047] The application field of the present invention is mainly concentrated in the field of dental medical image processing, especially in scenarios that require high-precision tooth segmentation, such as planning before dental implant surgery, monitoring during orthodontic treatment, and diagnosis of oral diseases. The method provided by the present invention can improve the doctor's work efficiency, improve the quality of diagnosis, and provide patients with more personalized and efficient treatment plans.

[0048] Step s1: Collect oral 3D CBCT images from (Hangzhou Qiantang) Dental Hospital. A small number of oral 3D CBCT images (40 patients) were annotated by doctors from West China University of Medical Sciences, and the annotated oral 3D CBCT images of 40 patients were obtained. For these 40 patients' annotated oral 3D CBCT images, they were divided into training set and test set in a ratio of 7:3 (the test set is for the convenience of comparison with other networks). The original data collected in the hospital is a dcm file, but in order to facilitate network training, prediction and viewing, it is uniformly compressed into nii.gz format.

[0049] Step s2: Define a function named load_data_volume to assist in reading the above training set data (annotated oral 3D CBCT images of 28 patients) and pre-set the preprocessing hyperparameters, so as to resize the CBCT images with an initial size of 101*565*565 pixels to 400*512*512 pixels. Then, after the CBCT image (i.e., image) and its corresponding annotation (i.e., label) are stored as keywords "image" and "label" respectively, they are subjected to intensity normalization (linearly converting the intensity value of the image from a specific range of CT scans (e.g., from -1024 to 1252) to the range of [0.0, 1.0]), random cropping (randomly cropping images and labels according to the given hyperparameters (256, 256, 256). This step balances the positions of positive and negative samples in the labels, thereby generating diverse training samples, which helps to improve the generalization ability of the model), random affine transformation, random intensity offset and other preprocessing operations.

[0050] Step s31: During the training process, the above processed training data is traversed for training. For each set of training data, the image part in the training data is traversed according to the number of channels. In each cycle, the two-dimensional slice under the current cycle is taken out according to the number of channels, and the slice image channel is copied three times using the repeat function to form a three-channel image (CHW=3*512*512) named input_sam_image. The two-dimensional slice of the label under the corresponding channel number is taken out together, and the zoom function is used to adjust its height and width to low_res, that is, 64 to obtain low_res_label (low_res is img_embedding_size*4, and img_embedding_size is a hyperparameter set in advance in the sam network structure). Then input_image is sent to SAM Encoder, and the 2D slice image is first sent to patch embed. After all patches pass through the linear projection layer, they are added with position encoding and special characters cls and sent to the Transformer Encoder layer. Because the SAM-B pre-training model is used, there are 12 layers of Transformer Encoder layers. In each Transformer Encoder layer, the LoRA fine-tuning module is added to the q and v projection layers before the multi-head attention layer to fine-tune the network.

[0051] Step s32: In the above process of looping through the number of channels for training data, in addition to obtaining input_image_sam, three slices between (c-2,h,w) and (c+2,h,w) (c is the current channel number, if c is the rightmost or leftmost, the edgemost slice is copied to replace the part beyond the edge) are taken out and saved as a new image named input_cnn_image, and input_cnn_image is input into the CNN bypass network in the sam encoder. The CNN bypass network is as follows Figure 2 As shown. In the CNN bypass network in the sam encoder, input_cnn_image first passes through a point-by-point convolution layer (the number of filters is 3) to obtain three feature maps with a channel number of 1, and then is sent to the CNN network (using the CNN architecture in the traditional ResNet). The features after passing through the CNN network are sent to the CLI module (CNN Transformer Interaction) to combine with the corresponding Transformer Encoder layer in the sam encoder layer, and then the feature maps fused by the CLI module are sent to the next CNN layer and Transformer Encoder layer respectively (that is, the number of CNN layers and Transformer Encoder layers are consistent and one-to-one corresponding). In the CLI module, the shape of the features obtained by the CNN module is (c i ,h i , w i ), reshape the token in the transformer layer into (c, h i , w i ), and then send it to the projection layer to get (c i ,h i , w i ), and then the two features are superimposed along the 0th dimension to obtain a new feature with a shape of (2×c i ,h i , w i ). Then the feature is sent to the channel fusion module to obtain a shape of (c i ,h i , w i ) and add the new features to the original features of each layer and send them to the next CNN layer and transformer layer. The channel fusion module contains four 1*1 convolutional layers (the number of filters is c i ), a combination of ReLU layer and 1*1 convolution layer

[0052] Step s33: In the process of looping through the number of channels of the training data, the prediction result in the previous channel (i.e. the prediction result for the image (c-1, h, w)) is saved and sent to the target frame generation module to obtain the boundary prompt frame (i.e. the pixel points at the lower left corner and the upper right corner of the prediction result are expanded by ten pixels each to obtain the target frame of the slice, and the target frame is expressed in the form of two coordinate points (lower left corner and upper right corner)), and the boundary prompt frame (i.e. the two coordinate points) is sent to the prompttencoder network.

[0053] Step s4: After 12 layers of Transformer Encoder layers with fine-tuning modules, the SAMEncoder part is passed. The extracted features are sent to the SAM MaskDecoder module together with the features extracted in the Prompt Encoder module to obtain a dictionary file, which contains the keywords "masks" and "low_res_logits". The "low_res_logits" of the low-resolution mask is obtained, and the loss function is calculated with the lows_res_label in the above data dictionary file. The loss function combines the cross entropy loss function and the dice loss function. The overall network structure is as follows: Figure 1 It is worth noting that in the SAM Encoder during the network training process, only the LoRa module participates in the forward propagation process to learn parameters, and the rest use the original pre-training parameters without change. As for the SAM Mask Decoder module, forward propagation is normally performed for training. The overall training process is as follows: Figure 3 shown.

[0054] Step s4: After completing the network training, the unlabeled CBCT from the (Hangzhou Qiantang) Dental Hospital is predicted to obtain the prediction results, that is, after an unlabeled oral 3D CBCT passes through the above network structure, the "masks" in the dictionary file are output in the final SAMMask Decoder module and the sliced ​​masks are combined to obtain the final prediction results and stored in the form of nii.gz. The test set is predicted to obtain the following results, which are compared with multiple original images and corresponding prediction images. Figure 4 The method of the present invention is compared with other methods, both of which use the same training sample size, i.e., CBCT images of 5 patients, and the same number of iterations, i.e., 250 epochs. The comparison table is shown in FIG. Figure 5 In order to better demonstrate the effect of tooth annotation of the present invention, the prediction results are visualized in 3D using 3Dslicer software. The comparison effect is shown in Figure 6 As shown, it can be clearly seen in the circle that the tooth marking result of the present invention can accurately separate the teeth.

[0055] The present invention introduces large model technology into the field of tooth annotation, and the training parameter amount is much smaller than that of previous networks. At the same time, the present invention considers the spatial information of 3D images and achieves better results.

Claims

1. A dental CBCT 3D teeth rapid annotation method based on SAM large model, characterized by: The improved SAM network is used to distinguish teeth from periodontal ligament and alveolar bone in CBCT images. The specific steps are as follows: Phase 1: Conversion and processing of dental data into trainable data; Phase 2: Building an improved SAM network model; Phase 3: Network model training; Phase 4: Tooth annotation and 3D visualization based on CBCT data.

2. The dental CBCT three-dimensional tooth rapid annotation method based on the SAM large model according to claim 1 is characterized in that: The specific implementation steps of Phase 1 are as follows: 1-1. Compress the original CBCT dental data into nii.gz format, divide it into training set, validation set and test set, and perform binary operation on the labels; 1-2. Read in the training set, resize the image to 400*512*512 pixels, and then perform a series of preprocessing operations such as intensity normalization, random cropping, random affine transformation, and random intensity offset.

3. The dental CBCT three-dimensional tooth rapid annotation method based on the SAM large model according to claim 1 is characterized in that: The specific implementation steps of Phase 2 are as follows: A CNN bypass network is added to the original SAM network model to supplement the extraction of spatial information of adjacent slices and improve the ability to capture edge details of teeth. At the same time, a Lora fine-tuning module is added to the Encoder part of the SAM network model to adapt to medical image information and make full use of the original SAM pre-training data. Finally, the target box information generated by the adjacent previous slice is used as the current slice input into the Prompt Encoder module.

4. The dental CBCT three-dimensional tooth rapid annotation method based on the SAM large model according to claim 3 is characterized in that: The CNN bypass network mainly includes a point-by-point convolution layer and a CNN network in sequence, and sends the features extracted by the CNN network together with the transformer layer in the Encoder part of the corresponding SAM network model into the channel fusion module CLI.

5. The dental CBCT three-dimensional tooth rapid annotation method based on the SAM large model according to claim 4 is characterized in that: The channel fusion module contains a combination of four 1*1 convolutional layers, a ReLU layer, and a 1*1 convolutional layer.

6. The dental CBCT three-dimensional tooth rapid annotation method based on the SAM large model according to claim 4 is characterized in that: The Lora fine-tuning module is applied to each Transformer Encoder layer in the Encoder part of the SAM network model. The Lora fine-tuning module is added to the q and v projection layers before the multi-head attention layer to fine-tune the network; assuming that the output of the network layer under the original pre-trained weights is as follows: z=W0x The formula after Lora fine-tuning the module is as follows: z′=W0x+BAx Among them, z is the original output, z' is the output after fine-tuning, W0 represents the weight matrix after pre-training, and B and A are low-rank matrices respectively.

7. The dental CBCT three-dimensional tooth rapid annotation method based on the SAM large model according to claim 4 is characterized in that: In the process of looping through the number of channels of the training data, the target frame information is generated by using the prediction of the adjacent previous slice, and the prediction result in the adjacent previous channel is saved. The prediction result is sent to the target frame generation module to obtain the target frame information, and the target frame information is sent to the Prompt Encoder module.

8. The dental CBCT three-dimensional tooth rapid annotation method based on the SAM large model according to claim 4, characterized in that: The specific implementation steps of Phase 3 are as follows: 3-1. Perform traversal training on the training data; For each set of training data, loop through the image part of the training data according to the number of channels; in each loop, take out the 2D slice image under the current loop according to the number of channels, use the repeat function to copy the slice image channel three times to form a three-channel image, named input_sam_image; take out the 2D slice image of the label under the corresponding channel number of the three-channel image, use the zoom function to adjust the height and width of the 2D slice image to low_res, get low_res_label, and then send input_sam_image to the Encoder part of the SAM network model using the Lora fine-tuning module; 3-2. In the process of looping through the number of channels of training data, in addition to obtaining input_sam_image, the three slices between (c-2,h,w) and (c+2,h,w) are taken out and saved as a new image, named input_cnn_image, and then input_cnn_image is input into the CNN bypass network; 3-3. In the process of looping through the number of channels of the training data, the prediction result of the previous channel is saved, that is, the prediction result for the image (c-1, h, w); then the prediction result of the previous channel is sent to the target frame generation module to obtain the target frame information, and the target frame information is sent to the Prompt Encoder module; 3-4. The features extracted after 12 layers of Transformer Encoder layers with Lora fine-tuning modules and the features extracted in the Prompt Encoder module are sent to the SAM Mask Decoder module to obtain a dictionary file. The dictionary file contains the keywords "masks" and "low_res_logits". The "low_res_logits" of the low-resolution mask is processed by the zoom function with the two-dimensional slice of the label under the corresponding channel of the input_sam_image to obtain the lows_res_label to calculate the loss function. The loss function adopts a combination of the cross entropy loss function and the dice loss function.

9. The dental CBCT three-dimensional tooth rapid annotation method based on the SAM large model according to claim 8, characterized in that: The Prompt Encoder module is used to process input prompt information, which includes user input in the form of points, boxes or masks to confirm which area or object to segment.

10. The dental CBCT three-dimensional tooth rapid annotation method based on the SAM large model according to claim 8, characterized in that: The SAM Mask Decoder module is used to generate the final segmentation mask. The module receives the feature map extracted from the image encoder and the prompt information processed by the Prompt Encoder module, and combines these two parts of information to predict the segmentation mask of the target object.

Images