Cataract multi-granularity classification method based on slit lamp image classification multi-task model training
Patent Information
- Application Number
- CN202410133750.X
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2024-01-31
- Publication Date
- 2026-09-29
- Estimated Expiration
- 2044-01-31
AI Technical Summary
[0009]为了克服已有为了克服裂隙灯图数据集中存在的噪声问题，为了获得更多粒度的白内障分类结果，本发明提供了一种基于裂隙灯图像分类多任务模型训练的白内障多粒度分类方法，结合图像分类任务和图像生成任务，该框架的第一个任务是图像分类，旨在将裂隙灯图像准确地分类为不同类型和程度的白内障；与此同时，引入了图像生成任务作为第二个任务，旨在通过消除噪声的影响学习到无噪声的特征，最终提高白内障多粒度分类的准确率和可靠性
[0059]本发明的有益效果主要表现在：利用多任务学习框架结合图像分类和图像生成任务，最大程度地利用数据，并提高模型的泛化能力；利用扩散模型进行无噪特征的提取，有效地提高了噪声的鲁棒性从而提高白内障分类的准确性；利用交叉融合模块动态调整特征权重，实现两种任务特征的有效融合，增强了模型表征能力和鲁棒性，当数据被大量噪声干扰时，本发明依旧能够得到较好的输出结果，在实际应用中简单且有效。
Smart Images
Figure CN117876788B_ABST
Abstract
Description
Technical Field
[0001] This invention pertains to artificial intelligence-based medical image processing and analysis methods and systems, and specifically relates to a method for training a multi-task model for noisy slit lamp image classification based on diffusion models and multi-task learning, thereby achieving multi-granularity cataract classification. Background Technology
[0002] Cataracts are a common eye disease prevalent worldwide, especially among the elderly. Cataracts refer to blurred vision and astigmatism caused by a decrease in the transparency of the eye's lens.
[0003] Slit-lamp imaging is widely used in clinical practice and medical research for the accurate diagnosis and classification of cataracts. A slit lamp is a common ophthalmic examination device that helps doctors assess eye diseases by shining a strong light on the eye and observing the details of its structures.
[0004] However, noise in slit-lamp images has a certain impact on the multi-granularity classification of cataracts. Noise is a phenomenon that degrades image quality due to factors such as image sensor, ambient light, and photographic conditions. This noise can interfere with the details of eye structures, making it more difficult for doctors to accurately identify and analyze the type and severity of cataracts in the images.
[0005] Addressing noise issues in slit-lamp images is crucial for accurate cataract classification. Recent advancements in image processing and computer vision have offered new opportunities to solve this problem. Image processing techniques can be used to denoise and enhance slit-lamp images, thereby improving image quality and the visibility of eye structures. Simultaneously, machine learning and deep learning algorithms can be used to automatically analyze and classify the type and severity of cataracts by training and learning from a large number of labeled slit-lamp images.
[0006] However, it is important to note that noise remains a challenge, as it can obscure crucial feature information, making it difficult for algorithms to accurately distinguish between different types of cataracts. Therefore, continued research and development of methods for reducing and eliminating noise in slit-lamp images are essential for improving the accuracy and reliability of multi-granularity cataract classification.
[0007] In summary, cataracts are a common eye disease, and slit-lamp images play a crucial role in clinical diagnosis and classification. However, noise in slit-lamp images can negatively impact multi-granularity cataract classification, necessitating the integration of image processing and machine learning algorithms to address this issue. Continuous research and technological innovation will contribute to improving the accuracy and reliability of cataract diagnosis, thereby better assisting patients in making treatment decisions and providing individualized care.
[0008] Currently, the main artificial intelligence-based cataract diagnosis methods include: Naive Bayes (…). Methods such as Bayes (NB), Random Forest (RF), Support Vector Machine (SVM), and Convolutional Neural Network (CNN) are used. However, these methods do not explore how noise generated by environmental factors, hardware devices, and other factors in the real world affects the performance of classification models or provide solutions. Summary of the Invention
[0009] To overcome the noise problem in existing slit-lamp image datasets and obtain more granular cataract classification results, this invention provides a multi-granular cataract classification method based on a multi-task model trained on slit-lamp image classification. This method combines image classification and image generation tasks. The first task is image classification, aiming to accurately classify slit-lamp images into different types and degrees of cataracts. Simultaneously, an image generation task is introduced as the second task, aiming to learn noise-free features by eliminating the influence of noise, ultimately improving the accuracy and reliability of multi-granular cataract classification. This invention utilizes a diffusion model for noise-free feature extraction. By applying partial differential equations, noise is gradually propagated and smoothed while maintaining the integrity of image edges and structures. This makes it easier to distinguish and analyze cataract features in the diffused image, thereby reducing the influence of noise in slit-lamp images. This invention uses cross-attention to fuse features from both tasks. The cross-attention mechanism allows the network to consider information from the other task while learning the feature representation of each task. Specifically, we use an attention mechanism to dynamically adjust the feature weights between different tasks to better utilize their correlations.
[0010] The technical solution adopted by this invention to solve its technical problem is:
[0011] A cataract multi-granularity classification method based on a slit-lamp image classification multi-task model training includes the following steps:
[0012] Step 1: Build a multi-granularity cataract classification network framework; construct a multi-granularity cataract classification network model based on a diffusion model to improve the robustness of the network to slit-lamp image noise when learning multi-granularity features of cataracts.
[0013] Step 2: Train the cataract multi-granularity classification network model; train the original multi-granularity cataract classification branch, the diffusion model-based image generation branch, and the classifier through a three-stage training method to improve the robustness of the cataract multi-granularity classification network model to noise in slit-lamp images.
[0014] Step 3: Dataset preprocessing; cropping the initial slit lamp images and dividing the cropped slit lamp images into nine categories based on existing labels;
[0015] Step 4: Cataract classification test; Input the slit-lamp image into the trained cataract multi-granularity classification network model to test the model's accuracy.
[0016] Furthermore, step one includes the following sub-steps:
[0017] Step 1.1: The cataract multi-granularity classification network model includes an original multi-granularity cataract classification branch, an image generation branch based on a diffusion model, a cross-fusion module, and a classifier;
[0018] Step 1.2: The original multi-granularity cataract classification branch consists of a Transformer encoder. This encoder is a sequence of multiple Transformer blocks, each containing a multi-head self-attention mechanism and a feedforward neural network for feature encoding of the input image. Before being input into this branch, the original image is first normalized to N×N pixels. Then, a convolutional layer with the same kernel size and padding step size is used to calculate the slit lamp image. After calculation, the image is uniformly divided into fixed-size blocks, i.e., image patching, for subsequent multi-head sub-attention calculation. Simultaneously, to learn multi-granularity cataract features, this convolutional layer enlarges the original 3-channel image to a feature space with a dimension greater than or equal to 128. Each small patch is then passed as a feature vector to the Transformer encoder for feature encoding and classification.
[0019] Step 1.3: The image generation branch based on the diffusion model consists of a diffusion process and a generator. The diffusion process constructs images by performing contour diffusion on noise, generating different images based on different noise patterns. The generator is used to generate higher-quality images based on the initial noise obtained from the diffusion process. The generator employs a residual convolutional encoder to encode the generated image, and a decoder converts the encoding back into an image. The diffusion model also applies a normalized flow method to align the generated image with the distribution of the real image, thereby improving the generator's performance.
[0020] Step 1.4: The cross-fusion module includes two fully connected layers; the two fully connected layers have the same structure, the purpose of which is to map the two features output by the original multi-granularity cataract classification branch and the diffusion model-based image generation branch to a high-dimensional space, and then add the two features calculated by the Softmax function to learn the similarity between the two features;
[0021] Step 1.5: The classifier of the multi-granularity cataract classification network model consists of a fully connected layer; this fully connected layer maps the number of output neurons of the cross-fusion module to the number of multi-granularity cataract categories. This fully connected layer has weights and biases, both defined as floating-point variables.
[0022] Furthermore, step two includes the following sub-steps:
[0023] Step 2.1: Train the original multi-granularity cataract classification branch; the training process is as follows:
[0024] Step 2.1.1: Prepare a slit lamp image and adjust its size to N×N pixels, then input it into the Transformer encoder;
[0025] Step 2.1.2: Use the pre-trained Transformer encoder as the initial model and start multi-round iterative training; during the training process, the learning rate of the initial few rounds of training needs to be set to linear growth, and the learning rate of subsequent rounds needs to be set to non-linear decrease. The multi-head self-attention mechanism and feedforward neural network in Transformer extract and interact with the features of the input data.
[0026] Step 2.1.3: Use a temporary classifier to classify the extracted features. This classifier has the same structure as the final classifier.
[0027] Step 2.1.4: Calculate the loss between the model output and the true label in each iteration;
[0028] Step 2.1.5: Minimize the loss function using the AdamW optimizer and save the optimal model;
[0029] Step 2.2: Train the image generation branch based on the diffusion model; the training process is as follows:
[0030] Step 2.2.1: Prepare a slit lamp image and adjust its size to N×N pixels, then input it into the generator;
[0031] Step 2.2.2: Construct the diffusion process; at each step of the diffusion, Gaussian noise is gradually added to the original image to disrupt the image distribution and shift the samples to a higher distribution direction;
[0032] Step 2.2.3: Train the generator; use the samples generated during the diffusion process to train the generator, gradually restore the noisy images to the original images, and estimate the noise level;
[0033] Step 2.2.4: Iteratively fit the model parameters in the generator using stochastic gradient descent to make the sampled data distribution fit the true data distribution as closely as possible. Similarly, the parameters for estimating the noise level during the diffusion process are also iteratively updated using the stochastic gradient descent method.
[0034] Step 2.2.5: Minimize the loss function using the Adam optimizer and save the optimal model;
[0035] Step 2.3: Train the classifier; the training process is as follows:
[0036] Step 2.3.1: Freeze the parameters of the Transformer encoder trained in Step 2.1 and the diffusion model encoder trained in Step 2.2;
[0037] Step 2.3.2: Prepare the slit lamp image and adjust the image size to N×N pixels, then input it into the Transformer encoder and the generator encoder respectively;
[0038] Step 2.3.3: Input the features output from the Transformer encoder and the generator encoder into the cross-fusion module for feature fusion;
[0039] Step 2.3.4: Input the fused features into the classifier and output the probability value of each type of cataract;
[0040] Step 2.3.5: Calculate the cross-entropy loss based on the calculated probability value and corresponding label for each type of cataract;
[0041] Step 2.3.6: Minimize the loss function using the AdamW algorithm and save the optimal generator model as the training output.
[0042] Step three includes the following sub-steps:
[0043] Step 3.1: Cropping the slit-lamp image; the cropping requirements include: the lens structure is fully presented in the image, and removing tissues that are not related to the multi-granularity classification features of cataracts, such as the eyelids, iris, sclera, etc.
[0044] Step 3.2: Fine slit lamp dataset; The final cataract multi-granularity categories are: nuclear type II, nuclear type III, ≥ nuclear type IV, mild cortical type, intermediate cortical type, severe cortical type, fully mature, posterior cystic, and healthy. The cataract multi-granularity categories correspond to the final output of the cataract multi-granularity classification network.
[0045] Step 3.3: Divide the dataset; Given the existing labels, divide the data into training, validation and test sets in a certain proportion, and ensure that the distribution of each cataract multi-granularity class is consistent across these datasets.
[0046] Step four includes the following sub-steps:
[0047] Step 4.1: Test the saved optimal model on the test set. The test metrics include precision, recall, F1-score, and average accuracy.
[0048] Step 4.2: For the nine granular categories of cataracts, calculate the accuracy of predicting each category and compare it with other more effective methods for handling imbalanced data, and demonstrate the effectiveness of the method through experiments;
[0049] Step 4.3: Add random noise to the source dataset, perform classification tests on the images containing the noise, and verify the robustness of the model to the noise.
[0050] Preferably, in sub-step 1.3, the generator described has an encoder that is a deep residual network, whose overall structure consists of multiple repeating residual blocks, each of which is composed of multiple convolutional layers, so that features are extracted from the noisy image through conventional convolution operations; its decoder consists of multiple modules composed of bilinear interpolation and convolutional layers.
[0051] In sub-step 1.2, the height N and width N of the initial image of the multiple Transformer blocks are the same as the output image shape of each Transformer block.
[0052] In sub-step 1.3, each Transformer block performs the calculation of the multi-head self-attention mechanism. The required operations are: LayerNorm normalization, calculation of the multi-head self-attention mechanism MSA, random deactivation of neurons (Dropout), and calculation of the multilayer perceptron (MLP). The weights and biases in both the MSA and MLP layers are defined as floating-point variables. The multi-head self-attention mechanism formula is as follows:
[0053] MSA(Q,K,V)=Concat(head1,…,head h )·W O ,
[0054]
[0055] MSA stands for Multi-Head Self-Attention Mechanism (MSA), W O W Q W K W V All of these are learnable parameters, defined as floating-point variables;
[0056] In addition, residual connections are used to make the training process more stable after each random dropout.
[0057] In sub-step 1.4, within the Transformer block, the input size of the corresponding LayerNorm, Multi-Head Attention (MSA) mechanism, and Multilayer Perceptron (MLP) must correspond to the image size; at the same time, the relationship between the number of attention heads h and the feature dimension dim in multi-head self-attention must satisfy dim%h=0.
[0058] In step 2.1.2, during iterative training, when setting the training learning rate, the number of training rounds in the warm-up phase should be less than or equal to 10% of the total number of rounds, and the maximum learning rate should be less than or equal to 1 × 10⁻⁶. -4 The minimum learning rate is greater than or equal to 1×10. -8 In the process of nonlinear learning rate decay, a sine function is chosen.
[0059] The beneficial effects of this invention are mainly reflected in the following aspects: by using a multi-task learning framework to combine image classification and image generation tasks, the data is utilized to the maximum extent and the generalization ability of the model is improved; by using a diffusion model to extract noise-free features, the robustness to noise is effectively improved, thereby improving the accuracy of cataract classification; by using a cross-fusion module to dynamically adjust feature weights, the effective fusion of features from the two tasks is achieved, enhancing the model's representation ability and robustness. Even when the data is interfered with by a large amount of noise, this invention can still obtain good output results. It is simple and effective in practical applications. Attached Figure Description
[0060] Figure 1 This is a schematic diagram of the cataract classification network framework used in this invention.
[0061] Figure 2 This is a schematic diagram of the cross-fusion module of the present invention.
[0062] Figure 3 This is the original multi-granularity cataract classification branch training flowchart of the present invention.
[0063] Figure 4 This is a flowchart of the training process for the image generation branch based on the diffusion model of the present invention.
[0064] Figure 5 This is a flowchart of the cross-fusion module and classifier training process of the present invention.
[0065] Figure 6 This is a flowchart of the model reasoning process of the present invention.
[0066] Figure 7 This is an overall flowchart of the method of the present invention. Detailed Implementation
[0067] The present invention will now be further described with reference to the accompanying drawings.
[0068] Reference Figures 1 to 7 A cataract multi-granularity classification method based on slit lamp image classification multi-task model training is proposed. The method includes four processes: construction of cataract multi-granularity classification network framework, training of cataract multi-granularity classification network, dataset preprocessing, and cataract classification testing.
[0069] The method includes the following steps:
[0070] Step 1: Construct a multi-granularity cataract classification network framework; build a multi-granularity cataract classification network model based on a diffusion model to improve the network's robustness to slit-lamp image noise when learning multi-granularity features of cataracts, including the following steps:
[0071] Step 1.1: The cataract multi-granularity classification network model includes an original multi-granularity cataract classification branch, a diffusion model-based image generation branch, a cross-fusion module, and a classifier. The original multi-granularity cataract classification branch is used to extract key features for multi-granularity cataract classification based on the image classification task; the diffusion model-based image generation branch is used to extract noise-free features from the initial noisy slit-lamp image based on the image generation task; the cross-fusion module is used to learn the correlation between the features extracted by the first two tasks; the classifier consists of a fully connected layer, which can map the features processed by the first part of the cross-fusion module, especially the cataract-related features, to nine multi-granularity cataract categories;
[0072] Step 1.2: The original multi-granularity cataract classification branch consists of a Transformer encoder; this encoder is a sequence of 16 Transformer blocks. Specifically, each Transformer block contains a multi-head self-attention mechanism and a feedforward neural network, sequentially performing LayerNorm_1, MSA, Dropout_1, LayerNorm_2, MLP, and Dropout_2. The LayerNorm normalization layer matches the dimension of the features, with no weights or biases, as its purpose is to normalize features and prevent gradient vanishing or exploding problems. The random deactivation Dropout has a probability of 0.1, with no weights or biases, to prevent overfitting. The Multilayer Perceptron (MLP) consists of two linear mappings... The fully connected layers are Linear(feature, 4*feature) and Linear(4*feature, feature), where feature represents the feature dimension of the output feature of the previous operation. The weights are W_2, b_2, W_3, and b_3, respectively, all defined as floating-point variables. A GELU activation function is then applied after each fully connected operation. The multi-head self-attention MSA implementation is a fully connected layer Linear(dim, 3*dim) operation with weights W_1 and bias b_1, both defined as floating-point variables. The multi-head self-attention MSA calculation formula is as follows:
[0073] MSA(Q,K,V)=Concat(head1,…,head h )·W O ,
[0074]
[0075] Where Q, K, and V are the three different layers of features obtained after the features pass through the learnable weights of the fully connected layer, and d k The formula, representing the feature dimension, is used to measure the correlation between different patches within a window. Q1, Q2, ..., Q h ,K1,K2,…,K h ,V1,V2,…,V h Based on Q, K, and V, the features are divided into multiple different features, where h represents the number of attention heads, which is set to 16.
[0076] Before being input into the Transformer encoder, the image needs to be patched, which involves dividing the image into fixed-size blocks. This is achieved using a convolutional layer with 3 in-channels, 768 kernels, a kernel size of 16×16, and a stride of 16. Simultaneously, the image of the lens portion cropped from the slit-lamp image is scaled to 224×224. After this operation, the image is divided into... In addition to the 196 patches, for classification tasks, a ClassToken based on Position Embedding is added and overlaid on the 196 patches to extract global class features and position information. Therefore, before being input into the Transformer encoder, there are a total of 197 features with 768 dimensions to be trained.
[0077] Additionally, after each random deactivation Dropout, the residual connection is added to the previous features;
[0078] Step 1.3: The image generation branch based on the diffusion model consists of a diffusion process and a generator; the diffusion process is constructed by contour diffusion on noise, and different images can be generated according to different noise morphologies; specifically, forward diffusion is the noise addition process, if x t If the image is at time t, then the image is only related to x at the previous time. t-1 This process can be viewed as a Markov process, satisfying:
[0079]
[0080] Where t∈[1,T],β t It is predefined, gradually increasing with time t, satisfying β1 < β2 < ... < β T Specifically, the value of T is 1000, and β... t The value ranges from [0.0001, 0.02] and increases linearly with time t. I represents a standard normal distribution.
[0081] The generator is used to generate the target image based on the initial noise obtained from the diffusion process, i.e., backdiffusion. Specifically, the generator contains an encoder and a decoder.
[0082] The encoder contains five convolutional modules. The first convolutional module contains one convolutional layer and one max-pooling layer. The convolutional kernel size is 7×7, the stride is 2, the padding is 3, and the number of kernels is 64, denoted as Conv(7,2,3,64). The max-pooling layer has a kernel size of 3×3, a stride of 2, and a padding of 3, denoted as MaxPool(3,2,3). The second convolutional module performs the convolution operation Conv(1,1,0,64) → Conv(3,1,1,64) → Conv(1,1,0,256) three times. The third convolutional module performs... The first convolutional module performs the following operations: Conv(1,1,0,128) → Conv(3,1,1,128) → Conv(1,1,0,512), repeated 4 times. The second convolutional module performs the following operations: Conv(1,1,0,256) → Conv(3,1,1,256) → Conv(1,1,0,1024), repeated 23 times. The third convolutional module performs the following operations: Conv(1,1,0,512) → Conv(3,1,1,512) → Conv(1,1,0,2048), repeated 3 times. Each module uses residual connections to prevent overfitting.
[0083] The decoder consists of five deconvolution modules, each containing a bilinear interpolation operation and a Conv(3,1,1,O) convolution operation; where O has values of 1024, 512, 256, 128, and 64 in each deconvolution module. Features are passed between each encoder and decoder module via skip connections, ultimately converting the encoded features into an image.
[0084] Similarly, before inputting the image into the encoder, we need to scale the image of the lens portion after cropping in the slit lamp image to a size of 224×224;
[0085] Step 1.4: The cross-fusion module includes two fully connected layers. These two fully connected layers maintain a consistent structure to map the two features output from the original multi-granularity cataract classification branch and the diffusion-based image generation branch to a high-dimensional space. The two features, calculated using the Softmax function, are then summed to learn the similarity between the two feature types. Specifically, both fully connected layers have a structure of Linear(feature, 3*feature), with weights W_4, b_4, W_5, and b_5, all defined as floating-point variables. The features extracted from the original multi-granularity cataract classification branch are represented as μ, and the features extracted from the diffusion-based image generation branch are represented as... The features after the two fully connected layers are represented as Q. μ,K μ V μ as well as That is, the Q, K, V information in MSA. Based on this, the cross-fusion module f DF The expression is as follows:
[0086]
[0087] Then, LayerNorm is used to scale the summed feature values to prevent gradient explosion.
[0088] Step 1.5: The classifier of the multi-granularity cataract classification network model consists of a fully connected layer; this fully connected layer maps the number of output neurons of the cross-fusion module to the number of multi-granularity cataract categories. Specifically, the classifier consists of a fully connected layer with 768 input neurons and 9 output neurons, corresponding to 9 multi-granularity cataract categories, with weights of W_6 and biases of b_6, both defined as floating-point variables.
[0089] Step 2: Train the cataract multi-granularity classification network model; a three-stage training method is used to train the original multi-granularity cataract classification branch, the diffusion model-based image generation branch, and the classifier respectively, improving the robustness of the cataract multi-granularity classification network model to noise in slit-lamp images; including the following steps:
[0090] Step 2.1: Train the original multi-granularity cataract classification branch; the training process is as follows:
[0091] Step 2.1.1: Prepare the slit lamp image and adjust the image size to 224×224 pixels. After patching, input it into the Transformer encoder.
[0092] Step 2.1.2: Next, we will use the pre-trained model on the ImageNet-1K dataset as the initial model and begin training for a total of 50 epochs. The first 1-5 epochs are a warm-up phase with a linearly increasing learning rate until the learning rate increases to 0.0001. Then, the next 6-50 epochs are a training phase with a sinusoidally decaying learning rate until the learning rate decreases to 0, accelerating the convergence of the model.
[0093] Step 2.1.3: Use a temporary classifier to classify the extracted features. This classifier has the same structure as the final classifier. Specifically, it consists of a fully connected layer that maps the number of output neurons of the cross-fusion module to the number of cataract multi-granularity categories. Specifically, the classifier consists of a fully connected layer with 768 input neurons and 9 output neurons, corresponding to 9 cataract multi-granularity categories. The weights are W_7 and the bias is b_8, both defined as floating-point variables.
[0094] Step 2.1.4: Calculate the loss between the model output and the true label in each iteration using the soft cross-entropy loss function;
[0095] Step 2.1.5: Minimize the loss function using the AdamW optimizer and save the optimal cataract classification model;
[0096] Step 2.2: Train the image generation branch based on the diffusion model; the training process is as follows:
[0097] Step 2.2.1: Prepare the slit lamp image and adjust the image size to 224×224 pixels, then prepare to add noise to the original image;
[0098] Step 2.2.2: Constructing the diffusion process; at each step of the diffusion, Gaussian noise is gradually added to the original image to disrupt the image distribution and shift the samples towards a higher distribution direction. Specifically, for each sample, we sample at random time t∈[0,T] and add noise to the image through the forward diffusion process;
[0099] Step 2.2.3: Train the generator; use the samples generated during the diffusion process to train the generator, gradually restoring the noisy images to the original images and estimating the noise level. Specifically, begin training for a total of 1000 epochs. Epochs 1-40 are a warm-up phase with a linearly increasing learning rate until the learning rate reaches 0.001. Then, epochs 41-1000 are a training phase with a linearly decreasing learning rate until the learning rate drops to 0.
[0100] Step 2.2.4: Iteratively fit the model parameters in the generator using stochastic gradient descent to make the sampled data distribution fit the true data distribution as closely as possible. Similarly, the parameters for estimating the noise level during the diffusion process are also iteratively updated using the stochastic gradient descent method. Specifically, we calculate the mean squared error loss function (MSE Loss) for both the noise predicted by the model and the noise we added.
[0101] Step 2.2.5: Minimize the loss function using the Adam optimizer and save the optimal model;
[0102] Step 2.3: Train the classifier; the training process is as follows:
[0103] Step 2.3.1: Freeze the parameters of the Transformer encoder trained in Step 2.1 and the generator encoder trained in Step 2.2;
[0104] Step 2.3.2: Prepare slit lamp images and adjust the image size to 224×224 pixels, then input them into the Transformer encoder and the encoder of the diffusion model generator respectively;
[0105] Step 2.3.3: Input the features output from the Transformer encoder and the diffusion model generator into the cross-fusion module for feature fusion;
[0106] Step 2.3.4: Input the fused features into the classifier and output the probability value of each type of cataract. Specifically, the training of the classifier for the original multi-granularity cataract classification branch is similar, starting with a total of 50 epochs of training. The first 1-5 epochs are a warm-up phase with linearly increasing learning rate until the learning rate increases to 0.0001. Then, the next 6-50 epochs are a training phase with sinusoidal cosine decay of learning rate until the learning rate decreases to 0, accelerating the convergence of the model.
[0107] Step 2.3.5: Calculate the Soft Cross Entropy Loss function based on the calculated probability value and corresponding label for each type of cataract;
[0108] Step 2.3.6: Minimize the loss function using the AdamW algorithm. Save the optimal cataract multi-granularity classification model as the training output.
[0109] Step 3: Dataset preprocessing; cropping the initial slit lamp images and dividing the cropped slit lamp images into nine categories based on existing labels;
[0110] Step 3.1: Cropping the slit-lamp image; the cropping requirements include: the lens structure is fully presented in the image, and tissues that are not related to the multi-granularity classification features of cataracts, such as the eyelids, iris, sclera, etc. are removed. In particular, the iris contains personal identification information, and removing it can prevent this part from interfering with model training while maintaining privacy.
[0111] Step 3.2: Fine slit lamp dataset; The final cataract multi-granularity categories are: nuclear type II, nuclear type III, ≥ nuclear type IV, mild cortical type, intermediate cortical type, severe cortical type, fully mature, posterior cystic, and healthy. The cataract multi-granularity categories correspond to the final output of the cataract multi-granularity classification network.
[0112] Step 3.3: Divide the dataset; With the existing labels, divide the data into training set, validation set and test set in a certain ratio, and ensure that the distribution of each cataract multi-granularity class is consistent across these datasets, at a ratio of 6:2:2.
[0113] Step four, for the aforementioned cataract multi-granularity classification test, the steps are as follows:
[0114] Step 4.1: Input the test set into the trained cataract multi-granularity classification network model. We calculate the precision, recall, F1-Score and mean accuracy of the final results to evaluate the overall performance of the model and its ability to classify cataracts in a multi-granularity manner under noise.
[0115] Step 4.2: In particular, for the nine granular categories of cataracts, we calculate the prediction accuracy for each category and compare it with other methods. This comparison method can better reflect that when the original image is interfered with by noise, the method that introduces the diffusion model can reduce the impact of noise on the model's inference and judgment, thereby improving the robustness of the model.
[0116] Step 4.3: Finally, to further verify the performance of this method in multi-granular cataract classification under increasingly severe noise in slit-lamp images, we added noise to the test data to different degrees and then let the model perform inference verification. Under this comparison method, we proved the robustness of the multi-granular cataract classification method. As the noise increases, the model's accuracy loss is smaller, and it can better classify the various categories of cataracts. Therefore, our multi-granular cataract classification method based on slit-lamp images is quite effective.
[0117] In this embodiment, slit lamps are commonly used in the actual clinical diagnosis of cataracts, providing detailed characteristics of the eye, particularly the opacity of the lens. There are many methods for capturing slit lamp images, one of which is using a digital slit lamp microscope. The slit lamp is connected to a digital camera, resulting in images with higher contrast and resolution, and enriching the functionality of the lens area, enabling ophthalmologists to determine the patient's condition from these images. This application uses slit lamp images for multi-granular classification of cataracts.
[0118] This embodiment also includes a system for a multi-granularity cataract classification method based on slit-lamp image model training, comprising:
[0119] The cataract multi-granularity classification network framework construction module is used to build a cataract multi-granularity classification network framework.
[0120] The cataract multi-granularity classification network training module is used to train the cataract multi-granularity classification network.
[0121] The dataset preprocessing module is used to preprocess datasets;
[0122] The cataract classification test module is used to classify and test cataracts.
[0123] The above modules correspond to steps one through four of embodiment one.
[0124] This embodiment applies to a cataract multi-granularity classification device based on slit-lamp images, including a memory and one or more processors. The memory stores executable code, and when the one or more processors execute the executable code, they are used to implement a cataract multi-granularity classification method based on a multi-task model trained on slit-lamp images.
[0125] This embodiment applies to a computer-readable storage medium on which a program is stored. When the program is executed by a processor, it implements the aforementioned method for multi-granular classification of cataracts based on a slit-lamp image multi-task model training.
[0126] The embodiments described in this specification are merely examples of implementations of the inventive concept and are for illustrative purposes only. The scope of protection of this invention should not be considered limited to the specific forms described in these embodiments; rather, it extends to equivalent technical means conceived by those skilled in the art based on the inventive concept.
Claims
1. A cataract multi-granularity classification method based on a slit-lamp image classification multi-task model training, characterized in that, The method includes the following steps: Step 1: Build a multi-granularity cataract classification network framework; construct a multi-granularity cataract classification network model based on a diffusion model to improve the robustness of the network to slit-lamp image noise when learning multi-granularity features of cataracts. Step 2: Train the cataract multi-granularity classification network model; train the original multi-granularity cataract classification branch, the diffusion model-based image generation branch, and the classifier through a three-stage training method to improve the robustness of the cataract multi-granularity classification network model to noise in slit-lamp images. Step 3: Dataset preprocessing; cropping the initial slit lamp images and dividing the cropped slit lamp images into nine categories based on existing labels; Step 4: Cataract classification test; Input the slit-lamp image into the trained cataract multi-granularity classification network model to test the model's accuracy; Step one of the method includes the following sub-steps: Step 1.1: The cataract multi-granularity classification network model includes an original multi-granularity cataract classification branch, an image generation branch based on a diffusion model, a cross-fusion module, and a classifier; Step 1.2: The original multi-granularity cataract classification branch consists of a Transformer encoder; this encoder is a sequence of multiple Transformer blocks, each containing a multi-head self-attention mechanism and a feedforward neural network, used to encode features of the input image. Before being input into this branch, the original image first needs to be normalized and adjusted to... The image is processed using pixels, and then a convolutional layer with the same kernel size and padding step size is used to calculate the slit lamp image. After the calculation, the image is evenly divided into fixed-size blocks, i.e., image patching, in order to perform subsequent multi-head sub-attention calculations. At the same time, in order to learn the multi-granularity features of cataracts, the original 3-channel image is enlarged to a feature space with a dimension greater than or equal to 128 using this convolutional layer. Subsequently, each small patch is passed as a feature vector to the Transformer encoder for feature encoding and classification. Step 1.3: The image generation branch based on the diffusion model consists of a diffusion process and a generator. The diffusion process is a process of constructing a model by spreading contours over noise. Different images can be generated depending on the different shapes of the noise. The generator is used to generate higher quality images based on the initial noise obtained from the diffusion process. In the generator, a residual convolutional encoder is used to encode the generated image, and the decoder converts the encoding into an image. The diffusion model also applies the normalized flow method to align the distribution of the generated image with that of the real image, thereby improving the performance of the generator. Step 1.4: The cross-fusion module includes two fully connected layers; the two fully connected layers have the same structure, the purpose of which is to map the two features output by the original multi-granularity cataract classification branch and the diffusion model-based image generation branch to a high-dimensional space, and then add the two features calculated by the Softmax function to learn the similarity between the two features; Step 1.5: The classifier of the multi-granularity cataract classification network model consists of a fully connected layer; this fully connected layer maps the number of output neurons of the cross-fusion module to the number of multi-granularity cataract categories. This fully connected layer has weights and biases, both defined as floating-point variables.
2. The cataract multi-granularity classification method based on a slit-lamp image classification multi-task model training as described in claim 1, characterized in that, Step two includes the following sub-steps: Step 2.1: Train the original multi-granularity cataract classification branch; The training process is as follows: Step 2.1.1: Prepare slit lamp images and adjust the image size. The pixels are then input into the Transformer encoder; Step 2.1.2: Use the pre-trained Transformer encoder as the initial model and begin multiple rounds of iterative training; During training, the learning rate needs to be set to linearly increasing for the first few rounds of training, and then non-linearly decreasing for subsequent rounds. The multi-head self-attention mechanism and feedforward neural network in Transformer extract and interact with the features of the input data. Step 2.1.3: Use a temporary classifier to classify the extracted features. This classifier has the same structure as the final classifier. Step 2.1.4: Calculate the loss between the model output and the true label in each iteration; Step 2.1.5: Minimize the loss function using the AdamW optimizer and save the optimal model; Step 2.2: Train the image generation branch based on the diffusion model; The training process is as follows: Step 2.2.1: Prepare slit lamp images and adjust the image size. The pixels are then input into the generator; Step 2.2.2: Constructing the diffusion process; At each step of the diffusion process, Gaussian noise is gradually added to the original image to disrupt the image distribution and shift the samples towards a higher distribution direction. Step 2.2.3: Train the generator; The generator is trained using samples generated during the diffusion process, gradually restoring the noisy images to the original images and estimating the noise level. Step 2.2.4: By iteratively fitting the model parameters in the generator using stochastic gradient descent, the sampled data distribution is made to fit the real data distribution as closely as possible. Similarly, the parameters for estimating the noise level during the diffusion process are also iteratively updated using the stochastic gradient descent method. Step 2.2.5: Minimize the loss function using the Adam optimizer and save the optimal model; Step 2.3: Train the classifier; The training process is as follows: Step 2.3.1: Freeze the parameters of the Transformer encoder trained in Step 2.1 and the diffusion model encoder trained in Step 2.2; Step 2.3.2: Prepare slit lamp images and adjust the image size. The pixels are then fed into the Transformer encoder and the generator encoder, respectively. Step 2.3.3: Input the features output from the Transformer encoder and the generator encoder into the cross-fusion module for feature fusion; Step 2.3.4: Input the fused features into the classifier and output the probability value of each type of cataract; Step 2.3.5: Calculate the cross-entropy loss based on the calculated probability value and corresponding label for each type of cataract; Step 2.3.6: Minimize the loss function using the AdamW algorithm and save the optimal generator model as the training output.
3. The cataract multi-granularity classification method based on a slit-lamp image classification multi-task model training as described in claim 1, characterized in that, Step three includes the following sub-steps: Step 3.1: Cropping the slit-lamp image; the cropping requirements include: the lens structure is fully presented in the image, and tissues unrelated to the multi-granularity classification features of cataracts, such as the eyelids, iris, and sclera, are removed; Step 3.2: Fine slit lamp dataset; The final cataract multi-granularity categories are: nuclear type II, nuclear type III, ≥ nuclear type IV, mild cortical type, intermediate cortical type, severe cortical type, fully mature, posterior cystic, and healthy. The cataract multi-granularity categories correspond to the final output of the cataract multi-granularity classification network. Step 3.3: Divide the dataset; Given the existing labels, divide the data into training, validation and test sets in a certain proportion, and ensure that the distribution of each cataract multi-granularity class is consistent across these datasets.
4. The cataract multi-granularity classification method based on a slit-lamp image classification multi-task model training as described in claim 1, characterized in that, Step four includes the following sub-steps: Step 4.1: Test the saved optimal model on the test set. The test metrics include precision, recall, F1-score, and average accuracy. Step 4.2: For the nine granular categories of cataracts, calculate the accuracy of predicting each category and compare it with other more effective methods for handling imbalanced data, and demonstrate the effectiveness of the method through experiments; Step 4.3: Add random noise to the source dataset, perform classification tests on the images containing the noise, and verify the robustness of the model to the noise.
5. The cataract multi-granularity classification method based on a slit-lamp image classification multi-task model training as described in claim 1, characterized in that, In sub-step 1.3, the generator described has an encoder that is a deep residual network. Its overall structure consists of multiple repeating residual blocks, and each residual block is composed of multiple convolutional layers. The generator extracts features from the noisy image through conventional convolution operations. Its decoder consists of multiple modules composed of bilinear interpolation and convolutional layers.
6. The cataract multi-granularity classification method based on a slit-lamp image classification multi-task model training as described in claim 1, characterized in that, In sub-step 1.2, the height N and width N of the initial image of the multiple Transformer blocks are the same as the output image shape of each Transformer block. .
7. The cataract multi-granularity classification method based on a slit-lamp image classification multi-task model training as described in claim 1, characterized in that, In sub-step 1.3, each Transformer block performs the calculation of the multi-head self-attention mechanism. The required operations are: LayerNorm normalization, calculation of the multi-head self-attention mechanism MSA, random deactivation of neurons (Dropout), and calculation of the multilayer perceptron (MLP). The weights and biases in both the MSA and MLP layers are defined as floating-point variables. The multi-head self-attention mechanism formula is as follows: ， ， （1）； in This indicates the multi-head self-attention mechanism (MSA). , , , All of these are learnable parameters, defined as floating-point variables; In addition, residual connections are used to make the training process more stable after each random dropout.
8. The cataract multi-granularity classification method based on a slit-lamp image classification multi-task model training as described in claim 1, characterized in that, In sub-step 1.4, within the Transformer block, the input sizes of the corresponding LayerNorm, Multi-Head Attention (MSA) mechanism, and Multilayer Perceptron (MLP) must correspond to the image size; simultaneously, the number of attention heads in the multi-head self-attention mechanism must also correspond to the image size. The relationship between the feature dimension dim and the feature dimension dim must satisfy the following conditions. .
9. The cataract multi-granularity classification method based on a slit-lamp image classification multi-task model training as described in claim 2, characterized in that, In step 2.1.2, during iterative training, when setting the training learning rate, the number of training rounds in the warm-up phase should be less than or equal to 10% of the total number of rounds, and the maximum learning rate should be less than or equal to... The minimum learning rate is greater than or equal to In the process of nonlinear learning rate decay, a sine function is chosen.