A tooth CBCT reconstruction method based on double-branch contrastive learning and transformer diffusion model

Abstract

The application discloses a tooth CBCT reconstruction method based on a double-branch contrast learning and a Transformer diffusion model, and relates to the technical field of medical image processing. In view of the problems that features are not aligned and global structure maintaining capability is poor when existing diffusion models process cross-modal tooth data, a double-branch variational autoencoder is first constructed, contrast learning is used to map features of in-vivo CBCT and Micro-CT to a shared latent space, and feature alignment is realized; secondly, a denoising network based on a Transformer architecture is constructed, and the global attention mechanism of the denoising network is used to learn high-frequency detail recovery from Micro-CT data guided by the aligned features. The application effectively solves the domain drift problem of clinical images and significantly improves the connectivity and clarity of micro-root canal structures.

Description

Technical Field

[0001] This invention relates to the fields of image processing and computer vision technology, specifically to a tooth CBCT reconstruction method based on bi-branch contrastive learning and a Transformer diffusion model. Background Technology

[0002] Cone-beam computed tomography (CBCT) is widely used in the diagnosis of dental pulp diseases. However, clinical in vivo CBCT images are limited by radiation dose and imaging equipment, often accompanied by artifacts and noise, making it difficult to clearly display fine root canal structures. In contrast, micro-CT images have extremely high resolution and are considered the "gold standard," but can only be used for extracted tooth samples. Existing technologies attempt to achieve super-resolution reconstruction through deep learning, but face two major challenges:

[0003] 1. Data misalignment and modal differences: In clinical practice, it is difficult to obtain paired "low-resolution in vivo" and "high-resolution in vivo" data for the same tooth under completely identical conditions. Usually, only low-resolution in vivo images can be obtained. Due to interference from skeletal backgrounds and artifacts in in vivo images, models trained directly using ex vivo data do not perform well in clinical applications.

[0004] 2. Insufficient ability to generate details: Traditional CNN-based generative models (such as U-Net) have limitations in capturing global anatomical structures and restoring extremely fine textures. Summary of the Invention

[0005] This invention proposes a tooth CBCT reconstruction method based on bi-branch contrastive learning and a Transformer diffusion model. It solves the modal difference problem by combining bi-branch VAE with contrastive learning and improves the generation quality by using a Transformer architecture.

[0006] A tooth CBCT reconstruction method based on bi-branch contrastive learning and a Transformer diffusion model includes the following steps:

[0007] S1, acquire training data, including low-resolution in vivo CBCT images, ex vivo CBCT images, and high-resolution Micro-CT images; and construct corresponding synthetic low-resolution Micro-CT images based on the high-resolution Micro-CT images;

[0008] S2. Construct a two-branch variational autoencoder for cross-modal feature decoupling and alignment: it includes an in vivo CBCT image encoder branch and an out-of-vivo CBCT image encoder branch. The in vivo CBCT image and the out-of-vivo CBCT image are respectively input into the corresponding encoder branches. Using a contrastive learning strategy, in vivo CBCT features and out-of-vivo CBCT features from the same tooth sample are defined as positive sample pairs, and features from different tooth samples are defined as negative sample pairs. By minimizing the contrastive loss function, the similarity of positive sample pairs in the latent space is maximized, while the similarity of negative sample pairs is minimized, thereby achieving cross-modal anatomical feature alignment.

[0009] S3. Construct a super-resolution network based on a latent diffusion model and adopt a denoising network based on a Transformer architecture. During the training phase, supervised training is performed using only Micro-CT data. The latent features encoded from the high-resolution Micro-CT image are added to the noise as the target. The latent features extracted from the synthetic low-resolution Micro-CT image by the encoder are used as conditional guidance and input into the denoising network to train it to predict and remove noise in order to recover the high-resolution features.

[0010] S4. The in vivo CBCT image to be processed is input into the trained in vivo CBCT image encoder branch for encoding to extract latent features. The out-of-body CBCT image to be processed is input into the trained out-of-body CBCT image encoder branch for encoding to extract latent features. The extracted latent features are used as conditional guidance information to align with the latent feature distribution of the synthesized low-resolution Micro-CT image. The latent features output from the in vivo image encoder branch and the out-of-body image encoder branch are input into the conditional generation network, and high-resolution CBCT images are generated through inverse denoising and decoding.

[0011] Optionally, the construction of the super-resolution network based on the latent diffusion model in step S3 includes the following steps:

[0012] S3.1, Construct a segmentation module: Used to segment the input noisy latent feature map into multiple patch sequences;

[0013] S3.2, Construct the embedding layer: Add position encoding to the tile sequence;

[0014] S3.3, Construct the Transformer module group: It contains multiple stacked DiT modules, each DiT module contains layer normalization, adaptive scaling and translation layers, multi-head self-attention layers and feedforward neural networks;

[0015] S3.4, Construct a conditional injection mechanism: Use the latent spatial features extracted in step S2 as conditions to inject them into the Transformer module group through an adaptive scaling and translation layer to adjust the denoising process of the features.

[0016] Optionally, to enhance the use of Micro-CT images for training, a synthetic degradation process for constructing Micro-CT data is also included:

[0017] S1.1, Perform a degradation processing procedure on the high-resolution Micro-CT image, the degradation processing procedure including at least one of downsampling, intensity enhancement, Gaussian noise injection and non-rigid deformation, to simulate the imaging degradation characteristics of clinical CBCT;

[0018] S1.2, the synthesized low-resolution Micro-CT image and its corresponding high-resolution Micro-CT image are used as paired data for supervised training of the super-resolution network in step S3.

[0019] Optionally, before step S2, a step of spatial registration of in vivo CBCT images and out vivo CBCT images is included: the in vivo CBCT data and out vivo CBCT data of the same tooth sample are spatially geometrically aligned using the iterative nearest point algorithm to maximize the mutual information between the two modalities and provide spatial pairing data for contrastive learning.

[0020] Optionally, step S2 further includes network training, wherein the network training employs a composite loss function. Its formula is:

[0021] ,

[0022] in, The reconstruction loss is used to measure the pixel difference between the generated image and the real image. KL divergence loss is used to constrain the distribution of latent variables; Perceptual loss is used to maintain semantic feature consistency; The contrast loss is used to optimize feature alignment in the latent space; These are the corresponding weighting coefficients.

[0023] Optionally, the adaptive scaling and translation layer in the DiT module receives two inputs: one is a feature sequence normalized by the layer, and the other is a modulation parameter, which is obtained by processing the multilayer perceptron after the time-step embedding and the latent features mapped to the shared latent space in step S2 are jointly input.

[0024] Beneficial effects of this invention:

[0025] 1. This invention solves the problem of high-resolution image reconstruction directly from in vivo CBCT due to missing information. Addressing the severe loss of anatomical information caused by low resolution and high noise in in vivo CBCT images, this invention introduces ex vivo CBCT as crucial guiding information. By constructing a bi-branch variational autoencoder and combining it with contrastive learning, the network is forced to learn the feature correspondence between in vivo and ex vivo images during training. It can utilize the morphological features of ex vivo images to fill in and infer missing details in in vivo images, thereby enabling the reconstruction of complete anatomical structures based on incomplete in vivo input during inference.

[0026] 2. Improved connectivity and recognizability of micro-root canal structures. Combined with the global attention mechanism of the Transformer architecture, this invention can ensure the continuity of tooth root canals in three-dimensional space while using in vitro information to guide reconstruction. This effectively solves the problems of root canal fracture and blurred margins in traditional methods, providing high-precision imaging support for the discovery of difficult root canals such as MB2 in clinical practice. Attached Figure Description

[0027] Figure 1 This is an overall flowchart of the method of the present invention.

[0028] Figure 2 This is a two-dimensional slice comparison image of the reconstruction effect of the method of the present invention.

[0029] Figure 3 These are three-dimensional root canal morphology comparison images showing the reconstruction effect of the method of this invention. Detailed Implementation

[0030] like Figure 1 As shown, a tooth CBCT reconstruction method based on bi-branch contrastive learning and a Transformer diffusion model includes the following steps:

[0031] S1, acquire training data, including low-resolution in vivo CBCT images, ex vivo CBCT images, and high-resolution Micro-CT images; and construct corresponding synthetic low-resolution Micro-CT images based on the high-resolution Micro-CT images;

[0032] S2. Construct a two-branch variational autoencoder for cross-modal feature decoupling and alignment: it includes an in vivo CBCT image encoder branch and an out-of-vivo CBCT image encoder branch. The in vivo CBCT image and the out-of-vivo CBCT image are respectively input into the corresponding encoder branches. Using a contrastive learning strategy, in vivo CBCT features and out-of-vivo CBCT features from the same tooth sample are defined as positive sample pairs, and features from different tooth samples are defined as negative sample pairs. By minimizing the contrastive loss function, the similarity of positive sample pairs in the latent space is maximized, while the similarity of negative sample pairs is minimized, thereby achieving cross-modal anatomical feature alignment.

[0033] S3. Construct a super-resolution network based on a latent diffusion model and adopt a denoising network based on a Transformer architecture. During the training phase, supervised training is performed using only Micro-CT data. The latent features encoded from the high-resolution Micro-CT image are added to the noise as the target. The latent features extracted from the synthetic low-resolution Micro-CT image by the encoder are used as conditional guidance and input into the denoising network to train it to predict and remove noise in order to recover the high-resolution features.

[0034] S4. The in vivo CBCT image to be processed is input into the trained in vivo CBCT image encoder branch for encoding to extract latent features. The out-of-body CBCT image to be processed is input into the trained out-of-body CBCT image encoder branch for encoding to extract latent features. The extracted latent features are used as conditional guidance information to align with the latent feature distribution of the synthesized low-resolution Micro-CT image. The latent features output from the in vivo image encoder branch and the out-of-body image encoder branch are input into the conditional generation network, and high-resolution CBCT images are generated through inverse denoising and decoding.

[0035] Optionally, the construction of the super-resolution network based on the latent diffusion model in step S3 includes the following steps:

[0036] S3.1, Construct a segmentation module: Used to segment the input noisy latent feature map into multiple patch sequences;

[0037] S3.2, Construct the embedding layer: Add position encoding to the tile sequence;

[0038] S3.3, Construct the Transformer module group: It contains multiple stacked DiT modules, each DiT module contains layer normalization, adaptive scaling and translation layers, multi-head self-attention layers and feedforward neural networks;

[0039] S3.4, Construct a conditional injection mechanism: Use the latent spatial features extracted in step S2 as conditions to inject them into the Transformer module group through an adaptive scaling and translation layer to adjust the denoising process of the features.

[0040] Optionally, to enhance the use of Micro-CT images for training, a synthetic degradation process for constructing Micro-CT data is also included:

[0041] S1.1, Perform a degradation processing procedure on the high-resolution Micro-CT image, the degradation processing procedure including at least one of downsampling, intensity enhancement, Gaussian noise injection and non-rigid deformation, to simulate the imaging degradation characteristics of clinical CBCT;

[0042] S1.2, the synthesized low-resolution Micro-CT image and its corresponding high-resolution Micro-CT image are used as paired data for supervised training of the super-resolution network in step S3.

[0043] Optionally, before step S2, a step of spatial registration of in vivo CBCT images and out vivo CBCT images is included: the in vivo CBCT data and out vivo CBCT data of the same tooth sample are spatially geometrically aligned using the iterative nearest point algorithm to maximize the mutual information between the two modalities and provide spatial pairing data for contrastive learning.

[0044] Optionally, step S2 further includes network training, wherein the network training employs a composite loss function. Its formula is:

[0045] ,

[0046] in, The reconstruction loss is used to measure the pixel difference between the generated image and the real image. KL divergence loss is used to constrain the distribution of latent variables; Perceptual loss is used to maintain semantic feature consistency; The contrast loss is used to optimize feature alignment in the latent space; These are the corresponding weighting coefficients.

[0047] Optionally, the adaptive scaling and translation layer in the DiT module receives two inputs: one is a feature sequence normalized by the layer, and the other is a modulation parameter, which is obtained by processing the multilayer perceptron after the time-step embedding and the latent features mapped to the shared latent space in step S2 are jointly input.

[0048] Two-dimensional slice comparison images of the reconstruction effect of the method of the present invention are shown below. Figure 2As shown, the left image E1 is the original low-resolution CBCT image, with the arrow indicating blurred root canal boundaries; the middle image E2 is the super-resolution image processed by the method of this invention, with the arrow indicating clear MB2 root canal boundaries, which is highly consistent with the Micro-CT true value image in the right image E3.

[0049] Figure 3 These are comparison images of three-dimensional root canal morphology reconstructed using the method of this invention. Left image D1 is the original CBCT three-dimensional reconstruction model, where insufficient resolution results in the fracture and loss of the portion below the root apex (the MB root canal corresponding to the yellow dashed line and the MB2 root canal corresponding to the blue dashed line). Middle image D2 is the three-dimensional model reconstructed using the method of this invention, showing that the integrity of the root apex is significantly restored, and the morphological details are close to the Micro-CT ground truth model in right image D3.

[0050] Example 1

[0051] This embodiment constructs a training set based on 40 tooth samples.

[0052] 1. Data preparation:

[0053] Modal acquisition: Acquiring in vivo CBCT, ex vivo CBCT, and high-resolution Micro-CT;

[0054] Contrast learning registration: To enable the contrast learning module to capture shared anatomical features, the ICP algorithm is used to perform spatial rigid registration of in vivo CBCT and ex vivo CBCT.

[0055] Synthetic data augmentation: A degradation pipeline is applied to Micro-CT images, including downsampling, intensity enhancement, injection of Gaussian noise, and non-rigid deformation. The generated "synthetic low-resolution image" forms a supervised pair with the original Micro-CT image, training a super-resolution network to recover details from the low-quality input.

[0056] 2. Construct a two-branch variational autoencoder for cross-modal feature decoupling and alignment:

[0057] Construct a dual-branch network architecture. One branch processes in-vivo data, and the other processes out-of-vivo data.

[0058] Features are mapped between the encoder space and the contrast space using a projection layer;

[0059] By applying contrastive loss functions such as InfoNCE, the feature distances of the same tooth in different modalities are brought closer together, while the feature distances of different teeth are pushed apart. The latent features extracted in this way contain robust anatomical structural information and remove artifact interference specific to clinical images.

[0060] 3. Construct a super-resolution network based on a latent diffusion model:

[0061] Diffusion process: T-step noise is added to the image in the latent space;

[0062] The denoising core, DiT, employs a Vision Transformer architecture. The input consists of a noisy latent map patch and a time step.

[0063] Conditional injection: The alignment features extracted in step 2 are used as conditions and injected into each module of DiT through a multilayer perceptron and an adaptive scaling and translation layer;

[0064] Attention mechanism: The multi-head self-attention layer inside DiT allows the model to integrate information globally, which better ensures the continuity of the tooth root canal structure compared to local convolution.

[0065] 4. Training settings:

[0066] Using a composite loss function: Using the AdamW optimizer, the initial learning rate... Cosine annealing strategy. Training input is cropped to... It is 0.001. It is 0.5. It is 0.5.

[0067] This method effectively utilizes the high-quality information from Micro-CT and overcomes the domain differences of clinical data through comparative learning, ultimately achieving near-Micro-CT super-resolution reconstruction results on clinical CBCT.

Claims

1. A tooth CBCT reconstruction method based on bi-branch contrastive learning and a Transformer diffusion model, characterized in that, Includes the following steps: S1, acquire training data, including low-resolution in vivo CBCT images, ex vivo CBCT images, and high-resolution Micro-CT images; and construct corresponding synthetic low-resolution Micro-CT images based on the high-resolution Micro-CT images; S2, construct a two-branch variational autoencoder for cross-modal feature decoupling and alignment: it includes an in vivo CBCT image encoder branch and an out vivo CBCT image encoder branch; it maps features from in vivo CBCT images, out vivo CBCT images, and Micro-CT images from the same tooth sample to a shared latent space to achieve feature alignment; The corresponding encoder branches are input into in vivo CBCT images and out vivo CBCT images respectively. Using a contrastive learning strategy, in vivo CBCT features and out vivo CBCT features from the same tooth sample are defined as positive sample pairs, and features from different tooth samples are defined as negative sample pairs. By minimizing the contrastive loss function, the similarity of positive sample pairs in the latent space is maximized, while the similarity of negative sample pairs is minimized, thereby achieving cross-modal anatomical feature alignment. S3. Construct a super-resolution network based on a latent diffusion model and adopt a denoising network based on a Transformer architecture. During the training phase, supervised training is performed using only Micro-CT data. The latent features encoded from the high-resolution Micro-CT image are added to the noise as the target. The latent features extracted from the synthetic low-resolution Micro-CT image by the encoder are used as conditional guidance and input into the denoising network to train it to predict and remove noise in order to recover the high-resolution features. S4, the in vivo CBCT image to be processed is input into the trained in vivo CBCT image encoder branch for encoding to extract latent features, the ex vivo CBCT image to be processed is input into the trained ex vivo CBCT image encoder branch for encoding to extract latent features, and the extracted latent features are used as conditional guidance information to align with the latent feature distribution of the synthesized low-resolution Micro-CT image; The latent features output from the in vivo image encoder branch and the out vivo image encoder branch are input into the conditional generation network, and high-resolution CBCT images are generated through inverse denoising and decoding.

2. The method according to claim 1, characterized in that, Step S3 describes the construction of a super-resolution network based on a latent diffusion model, which includes the following steps: S3.1, Construct a segmentation module: Used to segment the input noisy latent feature map into multiple patch sequences; S3.2, Construct the embedding layer: Add position encoding to the tile sequence; S3.3, Construct the Transformer module group: It contains multiple stacked DiT modules, each DiT module contains layer normalization, adaptive scaling and translation layers, multi-head self-attention layers and feedforward neural networks; S3.4, Construct a conditional injection mechanism: Use the latent spatial features extracted in step S2 as conditions to inject them into the Transformer module group through an adaptive scaling and translation layer to adjust the denoising process of the features.

3. The method according to claim 1, characterized in that, In step S1, in order to enhance the use of Micro-CT images for training, a synthetic degradation process for constructing Micro-CT data is also included: S1.1, Perform a degradation processing procedure on the high-resolution Micro-CT image, the degradation processing procedure including at least one of downsampling, intensity enhancement, Gaussian noise injection and non-rigid deformation, to simulate the imaging degradation characteristics of clinical CBCT; S1.2, the synthesized low-resolution Micro-CT image and its corresponding high-resolution Micro-CT image are used as paired data for supervised training of the super-resolution network in step S3.

4. The method according to claim 1, characterized in that, Before step S2, the method further includes a step of spatial registration of in vivo CBCT images and out vivo CBCT images: the in vivo CBCT data and out vivo CBCT data of the same tooth sample are spatially geometrically aligned using the iterative nearest point algorithm to maximize the mutual information between the two modalities and provide spatial pairing data for contrastive learning.

5. The method according to claim 1, characterized in that, Step S2 also includes network training, which employs a composite loss function. Its formula is: ， in, The reconstruction loss is used to measure the pixel difference between the generated image and the real image. KL divergence loss is used to constrain the distribution of latent variables; Perceptual loss is used to maintain semantic feature consistency; The contrast loss is used to optimize feature alignment in the latent space; These are the corresponding weighting coefficients.

6. The method according to claim 2, characterized in that, The adaptive scaling and translation layer in the DiT module receives two inputs: one is a feature sequence normalized by the layer, and the other is a modulation parameter. The modulation parameter is obtained by processing the multilayer perceptron after the time-step embedding and the latent features mapped to the shared latent space in step S2 are input together.

Images