Automatic prediction method for dose distribution in brachytherapy plan

Abstract

The invention relates to an automatic prediction method for dose distribution in a brachytherapy plan. Comprising the following steps: acquiring a structure image and a real dose distribution diagram; constructing an improved DiT model and a noise predictor; preprocessing the structure image and the real dose distribution diagram, and inputting the preprocessed structure image and the preprocessed real dose distribution diagram into an improved DiT model for model training; and preprocessing a noise image and a to-be-predicted structure image, inputting the preprocessed noise image and the to-be-predicted structure image into the trained improved DiT model, carrying out reverse denoising on the improved DiT model under the guidance of the feature information of the structure image, and finally decoding to obtain a predicted dose distribution diagram. The improved DiT model is a diffusion model using a Diffusion Transform architecture, the prediction precision is maintained, the consumption of computing resources is reduced, the efficiency of the radiotherapy plan is improved, and the controllability of model generation is further improved by combining the semantic features of the structural image.

Description

Technical Field

[0001] The invention relates to the field of radioactive preoperative planning, and in particular to an automatic prediction method for dose distribution in a brachytherapy plan. Background Art

[0002] Radiotherapy is one of the mainstream treatments for cancer patients and has made remarkable progress in the past few decades. In order to achieve good treatment results, high-quality radiotherapy plans are required to distribute sufficient radiation dose to the planning target volume (PTV) while minimizing radiation hazards to organs at risk (OARs). To achieve this, radiotherapy plans need to be manually adjusted by dosimetrists in a trial-and-error manner, which is very labor-intensive and time-consuming. In addition, the quality of treatment plans may vary due to different expertise and experience of radiologists.

[0003] Diffusion models have demonstrated their significant potential in modeling complex image distributions in certain vision tasks. Unlike other Deep Learning (DL) models, diffusion models do not make any additional assumptions about the target data distribution during training, thus avoiding the averaging effect and alleviating the over-smoothing problem. The diffusion-based model predicts dose maps with the shape of the ray penetration area and clearer boundaries. Therefore, introducing diffusion models into dose prediction tasks is a worthy endeavor.

[0004] The advantages of Transformer-based networks over convolutional networks have been widely established in various studies, demonstrating properties such as robustness, effective modal fusion and scalability. Similarly, the results of multimodal fusion compared to the CNN-based generator U-Net are consistent with our observations in this study. This advantage is not only due to the high-quality alignment achieved in the second training stage, but also due to the multi-head attention-based fusion mechanism, which is good at modeling long dependencies. This mechanism effectively integrates the compositional semantic information, more effectively guides the generation of visual latent vectors, and produces images that are closely aligned with the input image control. These findings highlight the unique advantages of the Diffusion Transformer architecture in effectively fusing multimodal information.

[0005] In this paper, we explored the feasibility of applying diffusion models to dose prediction tasks and proposed a diffusion-based model to automatically predict clinically acceptable dose distributions for lung cancer. Specifically, it consists of a forward process and a reverse process. In the forward process, the model uses a Markov chain to gradually transform the dose distribution, mapping the complex distribution to a Gaussian distribution by gradually adding predefined noise. Then, in the reverse process, given a pure Gaussian noise, the model gradually removes the noise in multiple steps and finally outputs a predicted dose map. In this process, the noise predictor is trained to predict the noise added in the corresponding step of the forward process. To further ensure the accuracy of the predicted dose distribution of PTV and OAR, we designed a DL-based structural encoder to extract lesion site information from CT images and segmentation masks of PTV and OAR, indicating the structure and relative position of the organs. By combining the lesion site information, the noise predictor can know the dose constraints between PTV and OAR, thereby allocating more appropriate doses to them and generating more accurate dose distribution maps.

[0006] Therefore, the present invention provides an automatic prediction method for dose distribution in brachytherapy planning to solve the above problems. Summary of the invention

[0007] In view of the above, the present invention proposes a method for automatic prediction of dose distribution in brachytherapy plan. On the basis of using diffusion model for dose prediction in radiotherapy plan, transformer architecture is introduced to ensure the accuracy of dose prediction while reducing the consumption of computing resources and enhancing the generation capability. Conditions are added to low-dimensional feature information to improve the controllability of model generation.

[0008] To achieve the above-mentioned object of the invention, the present invention proposes a method for automatically predicting dose distribution in a brachytherapy plan, comprising the following steps:

[0009] A method for automatically predicting dose distribution in a brachytherapy plan comprises the following steps:

[0010] S1: Acquire a structural image containing information about the patient's lesion location and a real dose distribution map corresponding to the structural image, wherein the structural image containing information about the patient's lesion location includes a CT image of the patient, a planned target volume, and a segmentation mask of an organ at risk;

[0011] S2: construct a conditional model, use the pre-trained DiT model as the base model, and mount the conditional model on the base model to form an improved DiT model;

[0012] S3: Preprocessing the structural image and the true dose distribution map to obtain characteristic maps of the structural image and the dose distribution map respectively;

[0013] S4: constructing a noise predictor, wherein the noise predictor is used to predict the noise added in the noise adding process;

[0014] S5: Freeze the base model, input the feature map of the structural image into the conditional model, input the feature map of the real dose distribution map into the base model, and train the improved DiT model. During the training process, the improved DiT model performs forward noise addition on the feature map of the real dose distribution map under the guidance of the structural image, and the noise predictor predicts the noise added during the forward noise addition process. The training goal is to minimize the difference between the noise predicted by the noise predictor and the noise actually added by the improved DiT model.

[0015] S6: Obtain a noise image. The feature maps of the noise image and the structure image to be predicted after preprocessing are respectively input into the base model and the conditional model in the trained improved DiT model. Under the guidance of the structure image to be predicted, the improved DiT model reversely denoises the feature map of the noise image according to the noise predicted by the noise predictor. Finally, the denoised feature map is reconstructed and decoded to obtain the predicted dose distribution map.

[0016] Preferably, in step S2, the conditional model includes N DiT blocks and N+1 zero linear layers, wherein the first zero linear layer is used as the first layer of the conditional model and is serially connected to the N DiT blocks, and the output of each DiT block is processed by a zero linear layer and then output to the base model;

[0017] The base model is mounted with the conditional model specifically as follows: the output of the first zero linear layer in the conditional model is added and combined with the input of the base model as the input of the first DiT block in the conditional model; the output of the i-th DiT block in the conditional model is processed by the zero linear layer and then added and combined with the output of the i-th DiT block in the base model as the input of the i+1-th DiT block in the base model; the number of DiT blocks in the conditional model is less than the number of DiT blocks in the base model.

[0018] Preferably, the pretreatment method described in step S3 is consistent with the pretreatment method described in step S6, specifically:

[0019] For the image to be preprocessed, an encoder is used to compress it into a low-dimensional latent space representation. The low-dimensional latent space representation is first patched into multiple tokens of the same size and non-overlapping. Multiple tokens are rearranged into a token sequence, and then the token sequence is encoded by a position embedding method, so that the encoded token sequence carries the information of the relative position of each token in the image. The encoded token sequence is the feature map obtained by preprocessing.

[0020] Preferably, in step S4, the specific method of the noise predictor predicting the noise is: fusing the feature maps of the structural image and the true dose distribution map at time t in the forward noisy process by element-by-element addition to obtain the predicted noise at time t.

[0021] Preferably, in step S5, the improved DiT model performs forward noise addition on the feature map of the real dose distribution map under the guidance of the structural image, specifically:

[0022]

[0023] Among them, x t-1 、x t is the characteristic diagram of the real dose distribution map with noise at time t-1 and time t in the forward noise addition process, x 0 is the characteristic diagram of the real dose distribution map without noise, z 1 、z t is the noise added at the first and tth moments in the forward noise addition process. All added noises obey the normal distribution. t are the preset hyperparameters.

[0024] Preferably, the objective function of the training in step S5 is:

[0025]

[0026] Among them, θ represents the trainable parameters in the conditional model, represents the trainable parameters in the optimization conditional model to minimize the function value, ε is the actual added noise at time t in the forward noise addition process, and ε t,θ is the noise added at time t predicted by the noise predictor, ‖‖ is the Euclidean distance of the vector, E (x,y) represents the expectation under the Euclidean distance, x is the structural image input to the conditional model, and y is the true dose distribution map.

[0027] Preferably, in step S6, acquiring the noise image refers to obtaining the noise image by sampling from a Gaussian distribution.

[0028] Preferably, in step S6, the improved DiT model performs reverse denoising on the feature map of the noise image according to the noise predicted by the noise predictor under the guidance of the structure image to be predicted, specifically:

[0029]

[0030] Among them, y′ t-1 ,y′ t ~N(0,1) is the feature map of the noise image at time t-1 and time t during the reverse denoising process, α t , β tis the hyperparameter at time t, ε t,θ is the noise added to the prediction of the noise predictor at time t, θ represents the trainable parameters in the conditional model, δ t is the noise intensity removed cumulatively at time t during the reverse denoising process, and ε′ is a normal distribution.

[0031] On the other hand, the present invention also proposes an automatic prediction system for dose distribution in radiotherapy plan based on improved DiT, which is used to implement the above automatic prediction method.

[0032] Compared with the prior art, the present invention has the following beneficial technical effects:

[0033] 1. The present invention uses an improved DiT model to automatically predict the radiation dose distribution of a cancer patient's treatment plan, thereby reducing the burden on dosimetry technicians and speeding up the radiotherapy procedure.

[0034] 2. The present invention replaces the UNet architecture used in the traditional diffusion model with the DiT architecture, thereby improving the generation capability of the model, that is, improving the accuracy of the dose distribution prediction map.

[0035] 3. The present invention solves the problem of difficulty in retaining basic structural information and unstable results by introducing structural images (CT images, PTV and OARs), thereby generating a more accurate dose map. BRIEF DESCRIPTION OF THE DRAWINGS

[0036] Figure 1 A block diagram of a method for automatically predicting dose distribution in a brachytherapy plan provided in an embodiment of the present application;

[0037] Figure 2 A schematic diagram of the forward noise addition process of the improved DiT model provided in an embodiment of the present application;

[0038] Figure 3 A schematic diagram of the architecture of the feature extraction condition model provided in the embodiment of the present application;

[0039] Figure 4 A schematic diagram of a patching process provided in an embodiment of the present application;

[0040] Figure 5 A schematic diagram of noise-adding training for the model forward process provided in an embodiment of the present application;

[0041] Figure 6 A schematic diagram of the inference process of the diffusion model inverse denoising provided in an embodiment of the present application;

[0042] Figure 7 Schematic diagram of the inverse denoising process of the diffusion model provided in an embodiment of the present application. DETAILED DESCRIPTION

[0043] To make the purpose, technical solution and advantages of the present invention more clearly understood, the present invention is further described in detail below in conjunction with the accompanying drawings and embodiments. It should be understood that the specific implementation methods described herein are only used to explain the present invention and do not limit the scope of protection of the present invention.

[0044] An embodiment of the present invention provides an automatic prediction method for dose distribution in a brachytherapy plan, which completes the radiation dose prediction in lung cancer treatment through a diffusion model. Compared with the traditional dose prediction method that relies on the professional knowledge and experience of radiologists, the method reduces the excessive dependence on radiologists and improves the efficiency of radiotherapy plan execution. The embodiment combines Transformer with the diffusion model to further improve the accuracy of dose prediction, and enhances the controllability of the prediction process by using a conditional architecture to extract the semantic features of the structural image.

[0045] The automatic prediction method of dose distribution in the radiotherapy plan of the present invention is implemented based on an improved DiT model. The improved DiT model is a diffusion model using the Diffusion Transformer architecture. Different from the traditional diffusion model based on the U-Net structure, the improved DiT model of the present invention is composed of the traditional DiT (Diffusion Transformer) model as the base model and the conditional model of the present invention. The present invention utilizes the advanced characteristics of the DiT architecture to improve the accuracy and efficiency of radiation dose prediction in cancer treatment.

[0046] like Figure 1 As shown, an embodiment provides a method for automatically predicting dose distribution in a brachytherapy plan, comprising the following steps:

[0047] S1, obtaining a structural image containing information about the patient's lesion location and a true dose distribution map corresponding to the structural image as a training data set.

[0048] The data set of this embodiment is a lung cancer data set, and the structural image x includes the patient's CT image, the planning target volume (PTV), and the segmentation mask of the risk organs (OARs). The planning target volume is a regional image artificially divided in the CT image, and the segmentation mask of the risk organs is also obtained by segmenting the CT image. The real dose distribution map is the dose distribution map of the drug in the patient's body marked on the basis of the CT image.

[0049] Wherein, the size of x is (2+O)*H*W, 2 represents the CT image and PTV, O represents the number of OARs, and H and W define the length and width of the image. The present invention utilizes the above data training to improve the ability of the DiT model to generate a predicted dose distribution map, and the generated predicted dose distribution map can maximize the therapeutic effect of the drug on the target planning volume while minimizing the potential damage to the dangerous organs.

[0050] In the embodiment, each structural image and true dose distribution map is a segmentation map of size 512*512*3, where 3 represents the number of channels.

[0051] S2, constructing a conditional model, and mounting the conditional model on the base model to form an improved DiT model.

[0052] like Figure 3 As shown, the present invention proposes an innovative architecture, called conditional architecture, as the architecture of the conditional model. The architecture combines the advantages of the ControlNet structure with the flexibility of the Transformer architecture to achieve more efficient model performance and more precise task control.

[0053] The conditional model includes N DiT blocks and N+1 zero-linear layers, where the first zero-linear layer is used as the first layer of the conditional model and is serially connected to the N DiT blocks. The output of each DiT block is processed by a zero-linear layer and then output to the base model.

[0054] In order to achieve the integration of ControlNet structure and Transformer architecture, the present invention selectively applies ControlNet structure to the initial N basic blocks of Transformer model. This selective integration allows the model to introduce the control capability of specific tasks through ControlNet structure while maintaining the parallel processing capability of Transformer architecture. During training, only the conditional model is trained, and the parameters of the base model are frozen.

[0055] The base model adopts the traditional DiT (Diffusion Transformer) model. The traditional DiT model is a diffusion model based on Transformer, which uses the diffusion process to gradually improve the feature representation of the token sequence. The base model includes multiple DiT blocks. The structure of the DiT block is well known in the art and is not described in detail in the present invention. In one embodiment of the present invention, N trainable DiT blocks corresponding to the first N DiT blocks of the base model are generated. Each trainable DiT block is independently trained and can be optimized according to specific task requirements. The outputs of these trainable DiT blocks are the key to model learning and adaptation because they allow the model to switch and adjust between different tasks.

[0056] The structural image is input to the first zero linear layer in the conditional model, the output of which is added to the input of the base model as the input of the first DiT block in the conditional model, and the output of the i-th trainable DiT block in the conditional model is connected to a zero linear layer. The output processed by the zero linear layer is added to the output of the i-th DiT block of the corresponding base model to form a combined output. Subsequently, the combined output is used as the input of the i+1-th DiT block of the subsequent base model. The number of DiT blocks in the conditional model is less than the number of DiT blocks in the base model. This design allows the model to utilize the control information provided by the ControlNet structure at each processing stage while maintaining the efficient data processing capabilities of the Transformer architecture.

[0057] In one embodiment of the present invention, N=13 is used as the final model parameter.

[0058] S3, performing data preprocessing on the structural image and the true dose distribution map, and processing the image into a long sequence low-dimensional feature map.

[0059] S3.1 In the present invention, the encoder part of the variational autoencoder (VAE) (hereinafter referred to as VAE encoder) is used to achieve image compression.

[0060] In one embodiment of the present invention, specifically, the VAE encoder receives an image of size 512×512×3 obtained in step S1 and compresses it into a low-dimensional latent space representation z. The size of this latent space is 64×64×4, where the image size I is equal to 64 and the number of channels C is equal to 4.

[0061] The output z of the S3.2 VAE encoder is further processed by “patchify”. Patching is a method of dividing an image into small non-overlapping regions (or “patches”) according to a specified patch size, which can then be processed or analyzed independently.

[0062] In the present invention, the patches obtained by patching the spatial representation z are called tokens, and each token represents a patch in the image.

[0063] like Figure 4 As shown in Figure 1, in order to ensure that the input data can be processed efficiently later, the present invention defines precise input specifications. Specifically, given a spatial representation of shape I×I×C (the noisy latent image from VAE), it is first “patched” into The patches are rearranged in sequence on the same dimension to obtain a token sequence of length T, where the hidden dimension is d.

[0064] In this embodiment, the value of the patch size p is set to 4. In this way, the compressed 64×64 low-dimensional feature image is divided into non-overlapping 16×16 patches, each of which contains a 4×4 pixel area in the original image, and the dimension of each patch is 4 (corresponding to the number of color channels). These patches are rearranged and transformed to obtain a token sequence of length T=256.

[0065] S3.3 The patched token sequence is then input into the first layer of the traditional DiT architecture. This step only uses the first layer of the DiT architecture, which is responsible for linearly embedding the input token sequence, and each embedded token has a dimension of d.

[0066] To further enhance the model's understanding of spatial structure, we use a frequency-based position embedding method, the sine-cosine version, to encode all input tokens. This position embedding method exploits the periodicity of sine and cosine functions and provides the model with fine-grained information about the relative position of tokens in the image.

[0067] S4, constructing a noise predictor, wherein the noise predictor is used to predict the noise added in the noise adding process.

[0068] The noise predictor is a key component used to predict and process noise in the present invention. It predicts the noise added by the improved DiT model at each moment in the forward noise adding process through formula reasoning.

[0069] The specific method of the noise predictor for predicting noise is: the feature map of the structural image and the real dose distribution map at time t in the forward noise addition process is fused by element-by-element addition to obtain the predicted noise at time t, thereby allowing the structural information to be transmitted in the noise predictor.

[0070] The calculation of the noise predictor is expressed as follows:

[0071] ε t,θ =f(g(x),g(y t ),γ t )

[0072] Among them, ε t,θ is the noise added by the improved DiT model predicted by the noise predictor at time t in the forward noise addition process, θ represents the trainable parameters of the conditional model, g(x) is the feature map after structural image preprocessing, g(y t ) is the characteristic diagram of the real dose distribution diagram at time t in the forward noise addition process, γ t is the accumulated noise intensity at time t during the forward noise addition process.

[0073] S5: Inputting the preprocessed feature map of the structural image and the feature map of the true dose distribution map into an improved DiT model to train the improved DiT model.

[0074] The present invention performs training optimization by updating the parameters of the conditional model so that the noise predicted by the noise predictor is as equal as possible to the noise added in the forward noise adding process. The base model is frozen during the training process, and the noise predictor is also not involved in the parameter update. The specific training process is as follows Figure 5 As shown, the feature map of the structural image and the true dose distribution map are input into the improved DiT model. Specifically, the feature map of the structural image is input into the conditional model in the improved DiT model, and the feature map of the true dose distribution map is input into the base model in the improved DiT model. The improved DiT model forward-noises the feature map of the true dose distribution map under the guidance of the structural image, and the noise predictor predicts the noise added in the forward noising process.

[0075] Forward Noise: The forward noise process of the improved DiT model is the process of gradually applying noise to the real data. Assume that the real dose distribution map of lung cancer in the embodiment is x 0 ~q(x), the forward process is to continuously add noise to the image until the image becomes pure noise, such as Figure 2 As shown in the figure, the characteristic graph of the real dose distribution graph with noise at each moment obtained during the noise addition process is defined as x 1 ,x 2 ,…,x T , the specific formula is as follows:

[0076]

[0077]

[0078] Among them, q(x t |x t-1 ) is the conditional probability of generating the t-th feature map from the t-1th feature map, x t-1 、x t is the characteristic diagram of the real dose distribution diagram with noise at time t-1 and time t in the forward noise addition process, is the hyperparameter at time t, I is the unit matrix, N(·) represents Gaussian distribution, q(x 1：T |x 0 ) is the conditional probability of generating the Tth feature map from the initial unnoised true dose distribution feature map, T is the number of forward denoising iterations, It represents the multiplication from time t=1 to time T.

[0079] In the process of adding noise, it is necessary to use the reparameterization technique to calculate the x at each moment. t To take a sample:

[0080] Assume α t =1-β t ,and Expand x t You can get:

[0081]

[0082] Among them, z i is the noise, which obeys the normal distribution, z 1 , z 2 , ...~N(0,I),α t and β t is the hyperparameter preset in the sampling process, x 0 Refers to the feature map at time 0, x t It is the feature map at time t after adding noise.

[0083] If you want to randomly sample a sample of a Gaussian distribution from a certain distribution, this process cannot backpropagate the gradient. Therefore, it is necessary to make it differentiable through the above-mentioned reparameterization technique so that the gradient can be backpropagated.

[0084] Due to the additivity of independent Gaussian distributions, that is, so

[0085] So we can mix two Gaussian distributions to get a standard deviation of A mixture of Gaussian distributions.

[0086] In the process of forward noise addition, as t increases, x t As T approaches infinity, x T This is complete Gaussian noise. In this embodiment, T=1000.

[0087] Training objective: The objective function of model training is defined as follows:

[0088]

[0089] Among them, θ represents the trainable parameters in the conditional model, represents the trainable parameters in the optimization conditional model to minimize the function value, ε is the actual added noise at time t in the forward noise addition process, and ε t ,θ is the noise added at time t predicted by the noise predictor, |||| is the Euclidean distance of the vector, E (x，y) represents the expectation under the Euclidean distance, x is the structural image input to the conditional model, and y is the true dose distribution map.

[0090] S6, using the trained improved DiT model, the dose distribution is predicted according to the structure image to be predicted.

[0091] like Figure 6 and Figure 7 As shown, the model in this step performs a denoising process, gradually denoising a random noise matrix until an image is generated. In the denoising process, noise markers are first sampled from a Gaussian distribution as noise images. The structural image x to be predicted includes CT images of lung cancer patients, planned target volumes, and segmentation masks of dangerous organs. The noise image and the structural image x are preprocessed separately to obtain feature maps of the noise image and feature maps of the structural image, which are used as inputs of the base model and the conditional model, respectively. Under the guidance of the structural image, the improved DiT model reversely denoises the feature map of the noise image according to the noise predicted by the noise predictor. After completing K iterations of the denoising process, it is reconstructed and decoded according to the position information embedded in the preprocessing process to obtain the final predicted dose distribution image. In a specific embodiment of the present invention, the number of denoising iterations K is set to 30.

[0092] It should be noted that the method of reconstructing and decoding the image is a common method in the field and will not be described in detail in the present invention.

[0093] For this purpose, the inverse denoising function is defined as follows:

[0094]

[0095] Among them, y′ t-1 , y′ t ~N(0,1) is the feature map of the noise image at time t-1 and time t during the reverse denoising process, α t , β t is the hyperparameter at time t, ε t，θ is the noise added to the prediction of the noise predictor at time t, θ represents the trainable parameters in the conditional model, δ t is the noise intensity accumulated at time t in the reverse denoising process, and ε′~N(0,I) is normally distributed. In this example, the hyperparameter α t Set to 0.98, β t Set to 0.02. The above formula only defines a denoising function for one iteration from time t to time t-1. After K iterations, the feature map of the predicted dose distribution image is obtained.

[0096] After denoising for K times from time t = K, the denoised feature map at time t = 0 is obtained in the form of a token sequence with a token length of T = 256. The token sequence is rearranged to transform the size of the latent space to 64 × 64 × 4, where the image size I is equal to 64 and the number of channels C is equal to 4. The VAE decoder receives a latent space image size to restore the image. The restored image size is 512 × 512 × 3, which is the final predicted dose distribution map.

[0097] The specific implementation methods described above provide a detailed description of the technical solutions and beneficial effects of the present invention. It should be understood that the above is only the most preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, supplements and equivalent substitutions made within the scope of the principles of the present invention should be included in the protection scope of the present invention.

Claims

1. A method for automatically predicting dose distribution in brachytherapy planning, characterized in that: The steps include: S1: Acquire a structural image containing information about the patient's lesion location and a real dose distribution map corresponding to the structural image, wherein the structural image containing information about the patient's lesion location includes a CT image of the patient, a planned target volume, and a segmentation mask of an organ at risk; S2: construct a conditional model, use the pre-trained DiT model as the base model, and mount the conditional model on the base model to form an improved DiT model; S3: Preprocessing the structural image and the true dose distribution map to obtain characteristic maps of the structural image and the dose distribution map respectively; S4: constructing a noise predictor, wherein the noise predictor is used to predict the noise added in the noise adding process; S5: Freeze the base model, input the feature map of the structural image into the conditional model, input the feature map of the real dose distribution map into the base model, train the improved DiT model, during the training process, the improved DiT model performs forward noise addition on the feature map of the real dose distribution map under the guidance of the structural image, and the noise predictor predicts the noise added in the forward noise addition process; The training objective is to minimize the difference between the noise predicted by the noise predictor and the noise actually added by the improved DiT model; S6: Obtain a noise image. The feature maps of the noise image and the structure image to be predicted after preprocessing are respectively input into the base model and the conditional model in the trained improved DiT model. Under the guidance of the structure image to be predicted, the improved DiT model reversely denoises the feature map of the noise image according to the noise predicted by the noise predictor. Finally, the denoised feature map is reconstructed and decoded to obtain the predicted dose distribution map.

2. The method for automatically predicting dose distribution in brachytherapy plan according to claim 1, characterized in that: In step S2, the conditional model includes N DiT blocks and N+1 zero linear layers, wherein the first zero linear layer is used as the first layer of the conditional model and is serially connected to the N DiT blocks, and the output of each DiT block is processed by a zero linear layer and then output to the base model; The base model mounting the conditional model is specifically: The output of the first zero linear layer in the conditional model is added and combined with the input of the base model as the input of the first DiT block in the conditional model. The output of the i-th DiT block in the conditional model is processed by the zero linear layer and added and combined with the output of the i-th DiT block in the base model as the input of the i+1-th DiT block in the base model. The number of DiT blocks in the conditional model is less than the number of DiT blocks in the base model.

3. The method for automatically predicting dose distribution in brachytherapy plan according to claim 1, characterized in that: The pretreatment method described in step S3 is consistent with the pretreatment method described in step S6, specifically: For the image to be preprocessed, an encoder is used to compress it into a low-dimensional latent space representation. The low-dimensional latent space representation is first patched into multiple tokens of the same size and non-overlapping. Multiple tokens are rearranged into a token sequence, and then the token sequence is encoded by a position embedding method, so that the encoded token sequence carries the information of the relative position of each token in the image. The encoded token sequence is the feature map obtained by preprocessing.

4. The method for automatically predicting dose distribution in brachytherapy plan according to claim 3, characterized in that: In step S3, the encoder uses the encoder part of a variational autoencoder.

5. The method for automatically predicting dose distribution in brachytherapy plan according to claim 1, characterized in that: In step S4, the specific method of the noise predictor predicting the noise is: fusing the feature maps of the structural image and the true dose distribution map at time t in the forward noise addition process by element-by-element addition to obtain the predicted noise at time t.

6. The method for automatically predicting dose distribution in brachytherapy plan according to claim 1, characterized in that: In step S5, the improved DiT model performs forward noise addition on the feature map of the real dose distribution map under the guidance of the structural image, specifically: Among them, x t-1 、x t is the characteristic graph of the real dose distribution graph with noise at time t-1 and time t in the forward noise process, x0 is the characteristic graph of the real dose distribution graph without noise, z1, z t is the noise added at the first and tth moments in the forward noise addition process. All added noises obey the normal distribution. t are the preset hyperparameters.

7. The method for automatically predicting dose distribution in brachytherapy plan according to claim 1, characterized in that: The objective function of the training in step S5 is: Among them, θ represents the trainable parameters in the conditional model, represents the trainable parameters in the optimization conditional model to minimize the function value, ε is the actual added noise at time t in the forward noise addition process, and ε t,θ is the noise added at time t predicted by the noise predictor, ‖‖ is the Euclidean distance of the vector, E (x,y) represents the expectation under the Euclidean distance, x is the structural image input to the conditional model, and y is the true dose distribution map.

8. The method for automatically predicting dose distribution in brachytherapy plan according to claim 1, characterized in that: In step S6, the acquiring of the noise image refers to acquiring the noise image by sampling from a Gaussian distribution.

9. The method for automatically predicting dose distribution in brachytherapy plan according to claim 1, characterized in that: In step S6, the improved DiT model performs reverse denoising on the feature map of the noise image according to the noise predicted by the noise predictor under the guidance of the structure image to be predicted, specifically: Among them, y t ′ -1 ,y t ′ ~N(0,1) is the feature map of the noise image at time t-1 and time t during the reverse denoising process, α t , β t is the hyperparameter at time t, ε t,θ is the noise added to the prediction of the noise predictor at time t, θ represents the trainable parameters in the conditional model, δ t is the noise intensity removed cumulatively at time t during the reverse denoising process, ε ′ is a normal distribution.

10. An automatic prediction system for dose distribution in radiotherapy planning based on improved DiT, characterized in that: The system is used to implement the automatic prediction method according to claim 1, and the system comprises: A training data acquisition module, used to obtain structural images and true dose distribution maps; A model building module, used for building a conditional model and mounting the conditional model on a base model to form an improved DiT model; A noise predictor module, used for constructing a noise predictor and predicting noise using the noise predictor; Model training module, used to train and improve the DiT model; The automatic prediction module is used to obtain the predicted dose distribution map through the trained improved DiT model.

Images