A method for grading nuclear cataracts under non-equilibrium and small sample conditions based on self-supervised contrast learning

Abstract

The present application belongs to the field of medical image processing and artificial intelligence technology, and discloses a nuclear cataract grading method under the condition of unbalanced small sample based on self-supervised contrast learning. First, a large number of easily obtained unlabelled eye images are used to perform domain adaptive pre-training through self-supervised contrast learning, so that the model can learn discriminative general visual features from the predicament of data scarcity; then, when fine-tuning on a small-scale, unbalanced labeled data set, a label distribution perception boundary loss function is introduced, and by setting a stricter classification boundary for the sample-poor severe class, the model is forced to pay attention to and learn the subtle features of these key cases. This method effectively cooperates to solve the small sample and class imbalance problem in nuclear cataract intelligent grading, significantly improves the overall accuracy of the grading model, especially the recognition sensitivity of severe cases, and provides an efficient, reliable and automated grading auxiliary tool for nuclear cataract diagnosis.

Description

Technical Field

[0001] This invention relates to the fields of medical image processing and artificial intelligence, and in particular to an automated grading method for nuclear cataracts, which is especially suitable for achieving high-precision intelligent grading under real-world conditions where labeled data is scarce and category distribution is uneven. Background Technology

[0002] Nuclear cataracts are the main subtype of age-related cataracts, and their severity grading is crucial for clinical diagnosis and treatment decisions. The Lens Opacities Classification System III (LOCS III) standard is widely used clinically, classifying nuclear cataracts into six grades, NC1 to NC6, based on the color of the nuclear nucleus.

[0003] In recent years, deep learning-based automatic classification methods have become a research hotspot. However, these methods face two major practical problems: First, the annotation of nuclear cataract images is highly dependent on expert knowledge, which is costly and time-consuming, resulting in a very limited amount of high-quality annotated data available for training, making the model prone to the dilemma of few-shot learning.

[0004] Second, real clinical data exhibits a significant long-tail distribution, meaning that the number of severe cases is far less than that of mild cases. This severe class imbalance drives model training to be heavily biased towards the majority class, resulting in a significant decline in the ability to identify the few but clinically significant severe cases.

[0005] Existing solutions have significant shortcomings: models pre-trained directly on large-scale general image datasets are difficult to capture the pathological features unique to nuclear cataract images; and the simple sample reweighting loss function used in the fine-tuning stage only adjusts the weights at the sample level, failing to fundamentally optimize the model's decision boundary and having limited ability to improve tail category classification.

[0006] Therefore, there is a need for an intelligent grading method for nuclear cataracts that can collaboratively address the dual challenges of small sample size and class imbalance, in order to achieve reliable, high-precision, automated grading, especially sensitive to severe cases, under limited and imperfect data conditions. Summary of the Invention

[0007] The purpose of this invention is to overcome the shortcomings of existing technologies and provide a method for classifying imbalanced nuclear cataracts in small samples based on self-supervised contrastive learning and label distribution-aware boundary. This method combines domain-adaptive self-supervised pre-training with class imbalance optimization based on decision boundary control, fully utilizing unlabeled data and finely guiding model learning, thereby achieving superior overall performance and tail-class recognition capabilities on a small amount of imbalanced labeled data.

[0008] To achieve the above objectives, the technical solution adopted by the present invention is as follows: a method for grading nuclear cataracts under unbalanced small sample conditions based on self-supervised contrastive learning, specifically including the following steps; Step 1: Collect a large-scale unlabeled eye image dataset and pre-train it using a self-supervised contrastive learning framework to obtain a pre-trained visual encoder; Step 2: Obtain a small-scale and class-imbalanced dataset of nuclear cataract slit lamp image annotations, load the pre-trained visual encoder, remove the self-supervised learning prediction head and add a classification layer to build a hierarchical model; Step 3: Supervised fine-tuning of the grading model using the nuclear cataract slit-lamp image annotation dataset, with optimization performed using a loss function to mitigate class imbalance. Step 4: Input the slit-lamp image to be graded into the fine-tuned grading model, and output the severity prediction result of nuclear cataract.

[0009] The first step specifically involves: the self-supervised contrastive learning framework is the momentum contrastive learning framework MoCo-v3, whose query encoder and momentum encoder both use Vision Transformer as the backbone network; the MoCo-v3 model is pre-trained using an unlabeled eye image dataset; the image data in the unlabeled eye image dataset undergoes random cropping, flipping, rotation, color perturbation, and Gaussian blur data augmentation; the loss function used for pre-training is the InfoNCE loss, the expression of which is: in, and These are the results of the same input image after data augmentation, fed into the query encoder and momentum encoder respectively. The dot product of the two represents their similarity. Then it is a negative sample; It is a temperature coefficient used to control the degree to which the loss function pays attention to difficult samples; The parameters of the momentum encoder are updated by querying the momentum average of the encoder parameters: in, These are the parameters of the momentum encoder. To query the encoder parameters, m is the momentum coefficient.

[0010] The second step is as follows: The dataset with nuclear cataract grading labels is divided into six levels from NC1 to NC6 according to the LOCS III standard, and is further divided into training set, validation set and test set according to a preset ratio; the pre-trained visual encoder model obtained in the first step is loaded, and a new classification head is added to form a grading model; in the fine-tuning stage, the input image is preprocessed by scaling to a fixed size, center cropping and normalization, and data augmentation is applied.

[0011] The loss function in the third step is the Label Distribution Aware Boundary Loss (LDAM Loss), and its expression is as follows: ;in, Let J be the classifier weight vector for the j-th cataract grade. This represents the feature vector extracted from the input nuclear cataract slit-lamp image via the backbone network. The true classification label for this slit-lamp image of nuclear cataract is C=6, corresponding to the six classification categories of nuclear cataract, where s is the feature scaling factor. In the boundary term calculation, for the true category... Its corresponding boundary value Calculated by the following formula: , For the training set The number of samples for class classification, max_m is a preset maximum boundary hyperparameter, that is, the number of samples with true labels. Substitute into the above formula From this, the required values ​​in the loss function can be obtained. ; During the fine-tuning phase, a hierarchical weighted sampling strategy is used to load training data, and the sampling weight of samples in each batch is inversely proportional to the frequency of their respective categories.

[0012] The unlabeled eye image dataset used in the first step is composed of multiple public datasets, including Cataract-train, KaggleCataract dataset, and Cataract Classification Dataset.

[0013] In the nuclear cataract slit-lamp image annotation dataset used in the second step, the preset ratio of the six levels from NC1 to NC6 is 5:11:8:3:3:1.

[0014] The fourth step is as follows: After performing the same preprocessing as in the fine-tuning stage on the slit-lamp cataract image to be graded, input it into the grading model; the grading model outputs a 6-dimensional vector, and the Softmax function is applied to the vector to obtain the probability distribution of each level from NC1 to NC6. The level with the highest probability is taken as the final prediction result of the severity level of nuclear cataract.

[0015] The beneficial effects of this invention are as follows: This invention, based on self-supervised contrastive learning and label distribution awareness, is a nuclear cataract grading method. First, it utilizes a large number of readily available unlabeled eye images for domain-adaptive pre-training through self-supervised contrastive learning, enabling the model to learn discriminative general visual features despite data scarcity. Subsequently, during fine-tuning on small-scale, imbalanced labeled datasets, a label distribution awareness boundary loss function is introduced. By setting stricter classification boundaries for the scarce severe case category, the model is forced to focus on and learn the subtle features of these key cases. This method effectively and collaboratively solves the problems of small sample size and class imbalance in intelligent nuclear cataract grading, significantly improving the overall accuracy of the grading model, especially its sensitivity to identifying severe cases. It provides an efficient and reliable automated grading aid for nuclear cataract diagnosis. Attached Figure Description

[0016] Figure 1 This is a schematic diagram of the overall process of the unbalanced small-sample nuclear cataract grading method based on self-supervised contrastive learning provided in an embodiment of the present invention.

[0017] Figure 2 This invention relates to a smart grading framework for nuclear cataracts based on MoCo-v3. Detailed Implementation

[0018] The technical methods mentioned in the above steps are described in detail below: like Figure 1 As shown, the overall method flow of the present invention mainly consists of the following steps: Step 1: Collect a large-scale unlabeled eye image dataset and pre-train it using a momentum contrastive learning framework with Vision Transformer as the backbone network to learn a general feature representation of eye images and obtain a pre-trained visual encoder. Step 2: Obtain a small-scale nuclear cataract image annotation dataset, load the weights of the pre-trained visual encoder, remove the self-supervised learning head and add a fully connected classification layer to construct a nuclear cataract grading model. Step 3: Supervised fine-tuning of the hierarchical model using the labeled dataset. During training, the label distribution-aware boundary loss function is used as the optimization objective, and adaptive classification boundary constraints are applied to categories with different frequencies to alleviate class imbalance. Step 4: Input the slit-lamp image to be graded into the grading model that has been fine-tuned, and output the corresponding nuclear cataract severity level prediction result.

[0019] The first step includes: First, unlabeled eye images from multiple publicly available sources, including Cataract-train, Kaggle Cataract dataset, and CataractClassification Dataset, were integrated to construct a pre-training dataset containing 18,341 eye images. This dataset includes normal eyes and various types of cataract eyes, and is used only to learn general visual features of the eye without any manual annotation.

[0020] The preprocessing performed on this portion of the unlabeled eye images is as follows: All eye images were standardized, with normalization parameters based on the statistical properties of the dataset: mean: [0.372, 0.285, 0.271], standard deviation: [0.188, 0.191, 0.187]. For each image x in the dataset, a series of strongly random data augmentation operations were applied to generate two distinct augmented views x1 and x2, which served as positive sample pairs for contrastive learning. The augmentation operations included: 1) Random size cropping and scaling: Randomly crop 80% to 100% of the original image and scale it to a fixed resolution of 224×224 pixels.

[0021] 2) Random color perturbation: Randomly adjust the brightness, contrast, saturation and hue of the image, with a perturbation factor of 0.4.

[0022] 3) Random Gaussian blur: Apply Gaussian blur with a probability of 0.1, and the kernel size is randomly selected from {3, 5, 7}.

[0023] 4) Random horizontal flip: Perform a horizontal flip with a probability of 0.5.

[0024] Subsequently, self-supervised pre-training was performed using the MoCo-v3 framework. The backbone network consists of a query encoder and a momentum encoder. Both encoders use Vision Transformer-Base / 16 as feature extractors. The model segments the input image into 16×16 image patches, performs linear projection and positional encoding, and then inputs them into a stack of 12 Transformer encoder layers. The query encoder connects a three-layer multilayer perceptron as a projection head after the ViT backbone, with input and output dimensions of 768 and 2048 respectively, and a two-layer MLP as a prediction head, with input and output dimensions of 2048 and 2048 respectively. The momentum encoder contains only the same ViT backbone and projection head, but does not include a prediction head.

[0025] Meanwhile, the update mechanism of the momentum encoder is: the parameters of the momentum encoder By querying the encoder parameters Perform momentum update. The formula is: The momentum coefficient m is set to 0.99.

[0026] For the loss function, InfoNCE contrastive loss is used. For query features... and its corresponding positive sample key features and the set of negative sample key features sampled from the dynamic queue. The loss calculation formula is as follows: in and These are the results of the same input image after data augmentation, fed into two encoders respectively, and their dot product represents their similarity. Then it is a negative sample. It is a temperature coefficient used to control the degree to which the loss function pays attention to difficult samples.

[0027] Regarding the selection of training details such as the optimizer: The AdamW optimizer was used with an initial learning rate of 1e-4, followed by a 40-epoch linear warm-up and then decay using a cosine annealing strategy. The batch size was 128, and training was performed for 200 epochs on a single NVIDIA RTX 3090 GPU. After training, the ViT backbone and projection head weights of the query encoder were saved as a pre-trained visual encoder.

[0028] The second step includes: The publicly available Nuclear Cataract Dataset was used as the supervised fine-tuning dataset. This dataset contains 959 slit-lamp images, manually labeled by professional ophthalmologists into 6 levels (NC1 to NC6) according to the LOCS III standard. The specific class distribution exhibits significant imbalance, as shown in Table 1.

[0029] Table 1. Category Distribution of Nuclear Cataract Dataset

[0030] The 959 images were randomly divided into training, validation, and test sets in a 60%:20%:20% ratio, using stratified sampling to ensure consistent proportions across categories. The final set consisted of 575 images for training, 192 images for validation, and 192 images for test.

[0031] Next, the parameters of the pre-trained visual encoder from the first step, namely the ViT-B / 16 backbone network and the three-layer MLP projector head, are loaded. Simultaneously, the two-layer MLP prediction head unique to the MoCo-v3 framework is removed. After the output of the projector head, a completely new fully connected layer is added as a classification head. This classification head has an input dimension of 2048 and an output dimension of 6, aiming to correspond to the six categories NC1-NC6. This completes a nuclear cataract grading model that uses the pre-trained ViT as a feature extractor and the new classification head as a classifier.

[0032] The specific implementation plan for the third step includes: For data loading, a weighted random sampler was used to address the class imbalance problem in the training data. The principle of this weighted random sampler is to assign sampling weights to each class j. =1.0 / ,in This is the number of samples of that category in the training set, which makes it more likely that categories with few samples will be sampled in each round of training.

[0033] To further address the issue of imbalanced training data, a label distribution-aware boundary loss was chosen as the optimization objective for the fine-tuning stage. This loss function introduces a factor related to the number of training samples for each class j. Inversely proportional boundary terms This forces the classifier to learn a larger decision boundary for the class with fewer samples. The formula for the LDAM loss function is as follows: in, Let J be the classifier weight vector for the j-th cataract grade. This represents the feature vector extracted from the input slit lamp image through the backbone network. The image is assigned a true classification label, C=6 corresponds to the 6 classification categories of nuclear cataracts, and s is the feature scaling factor, set to 30. In the boundary term calculation, for the true category... Its corresponding boundary value Calculated by the following formula: , For the training set The number of samples for class grading will determine the true labels. Substitute into the above formula From this, the required values ​​in the loss function can be obtained. The constant max_m was determined to be optimal at 0.6 through a grid search on the validation set. Boundary values ​​for each category were then calculated. ≈0.54, ≈0.40, ≈0.45, ≈0.63, ≈0.62, ≈0.96.

[0034] For the optimizer selection, the SGD optimizer was used, employing a hierarchical learning rate strategy. The learning rate of the ViT backbone network was set to a relatively low 1e-5 to avoid destroying pre-trained features; the learning rate of the projection head was set to 1e-4; and the learning rate of the newly added classification head was set to a relatively high 3e-4 to allow it to quickly adapt to new tasks. Weight decay was set to 1e-4. Supervised training was performed on the training set for 50 epochs. After each epoch, the macro-average F1 score was calculated on the validation set. After training, the model weights with the highest F1 score on the validation set were selected as the final model.

[0035] The fourth step specifically includes: For a slit-lamp image I to be graded, it is first scaled to 256×256 pixels, then centered and cropped to obtain an input region of 224×224 pixels. Next, the image is normalized at the channel level using the same mean and standard deviation as in the training phase (mean=[0.372, 0.285, 0.271], std=[0.188, 0.191, 0.187]). The preprocessed image tensor is then input into the nuclear cataract grading model trained in the third step. The model sequentially passes through the ViT backbone network, projection head, and classification head, ultimately outputting a 6-dimensional logits vector z = [z1, z2, ..., z6], where z... j Corresponding to NC j The grade score. Applying the Softmax function to the logits vector z transforms it into a probability distribution: probability Indicates that image I belongs to NC j The prediction confidence level of the grade is determined. Finally, the category index with the highest probability is taken as the final prediction result for the severity grade of nuclear cataract.

[0036] By implementing the above four steps in detail, the entire process from pre-training on unlabeled data, fine-tuning with small-sample imbalanced data, to final clinical image grading prediction can be completed. Using only 959 labeled images, this method achieved a macro-average accuracy of 85.94% on an independent test set, effectively validating its superiority in solving small-sample and imbalanced problems.

Claims

1. A method for grading nuclear cataracts under unbalanced small sample conditions based on self-supervised contrastive learning, characterized in that, The specific steps are as follows; Step 1: Collect a large-scale unlabeled eye image dataset and pre-train it using a self-supervised contrastive learning framework to obtain a pre-trained visual encoder; Step 2: Obtain a small-scale and class-imbalanced dataset of nuclear cataract slit lamp image annotations; load the pre-trained visual encoder, remove the self-supervised learning prediction head and add a classification layer to build a hierarchical model; Step 3: Supervised fine-tuning of the grading model using the nuclear cataract slit-lamp image annotation dataset, with optimization performed using a loss function to mitigate class imbalance. Step 4: Input the slit-lamp image to be graded into the fine-tuned grading model, and output the severity prediction result of nuclear cataract.

2. The non-equilibrium small-sample nuclear cataract grading method based on self-supervised contrastive learning according to claim 1, characterized in that, The first step specifically involves: the self-supervised contrastive learning framework is the momentum contrastive learning framework MoCo-v3, whose query encoder and momentum encoder both use Vision Transformer as the backbone network. The MoCo-v3 model is pre-trained using an unlabeled eye image dataset; the image data in the unlabeled eye image dataset undergoes random cropping, flipping, rotation, color perturbation, and Gaussian blur data augmentation; the loss function used for pre-training is the InfoNCE loss, whose expression is: in, and These are the results of the same input image after data augmentation, fed into the query encoder and momentum encoder respectively. The dot product of the two represents their similarity. Then it is a negative sample; It is a temperature coefficient used to control the degree to which the loss function pays attention to difficult samples; The parameters of the momentum encoder are updated by querying the momentum average of the encoder parameters: in, These are the parameters of the momentum encoder. To query the encoder parameters, m is the momentum coefficient.

3. The non-equilibrium small-sample nuclear cataract grading method based on self-supervised contrastive learning according to claim 2, characterized in that, The second step is as follows: The dataset with nuclear cataract grading labels is divided into six levels from NC1 to NC6 according to the LOCS III standard, and is further divided into training set, validation set and test set according to a preset ratio; the pre-trained visual encoder model obtained in the first step is loaded, the self-supervised learning projection head is removed and a new classification head is added to form a grading model; in the fine-tuning stage, the input image is preprocessed by scaling to a fixed size, center cropping and normalization, and data augmentation is applied.

4. The non-equilibrium small-sample nuclear cataract grading method based on self-supervised contrastive learning according to claim 3, characterized in that, The loss function in the third step is the Label Distribution Aware Boundary Loss (LDAM Loss), and its expression is as follows: ;in, Let J be the classifier weight vector for the j-th cataract grade. This represents the feature vector extracted from the input nuclear cataract slit-lamp image via the backbone network. The true classification label for this slit-lamp image of nuclear cataract is C=6, corresponding to the six classification categories of nuclear cataract, where s is the feature scaling factor. In the boundary term calculation, for the true category... Its corresponding boundary value Calculated by the following formula: , For the training set The number of samples for class classification, max_m is a preset maximum boundary hyperparameter, that is, the number of samples with true labels. Substitute into the above formula From this, the required values ​​in the loss function can be obtained. ; During the fine-tuning phase, a hierarchical weighted sampling strategy is used to load training data, and the sampling weight of samples in each batch is inversely proportional to the frequency of their respective categories.

5. The non-equilibrium small-sample nuclear cataract grading method based on self-supervised contrastive learning according to claim 4, characterized in that, The unlabeled eye image dataset used in the first step is composed of multiple public datasets, including Cataract-train, KaggleCataract dataset, and Cataract Classification Dataset.

6. The non-equilibrium small-sample nuclear cataract grading method based on self-supervised contrastive learning according to claim 5, characterized in that, In the nuclear cataract slit-lamp image annotation dataset used in the second step, the preset ratio of the six levels from NC1 to NC6 is 5:11:8:3:3:

1.

7. The non-equilibrium small-sample nuclear cataract grading method based on self-supervised contrastive learning according to claim 6, characterized in that, The fourth step is as follows: after performing the same preprocessing as in the fine-tuning stage on the slit-lamp image to be graded, it is input into the grading model; the grading model outputs a 6-dimensional vector, and the Softmax function is applied to the vector to obtain the probability distribution of each level from NC1 to NC6. The level with the highest probability is taken as the final prediction result of the severity level of nuclear cataract.

Images