Structure-aware mask reconstruction learning system and method for OCT image segmentation
By employing a structure-aware mask reconstruction learning system in retinal OCT image segmentation, utilizing a two-stage training approach with a ViT-Base architecture and a multi-expert adaptive reconstruction module, and combining pixel-level mean square error and gradient loss optimization, the problems of low data efficiency and weak fine structure perception in retinal OCT image segmentation are solved, achieving high-precision segmentation of retinal slice boundaries and lesion areas.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- NINGBO INST OF MATERIALS TECH & ENG CHINESE ACAD OF SCI
- Filing Date
- 2026-02-05
- Publication Date
- 2026-04-24
AI Technical Summary
Existing technologies rely on high-quality labeled data for retinal OCT image segmentation, which is costly and scarce. Self-supervised learning methods are inefficient and have difficulty effectively learning the fine structure and boundary information of the retina. Furthermore, masking strategies fail to effectively focus on clinically critical areas.
A structure-aware mask reconstruction learning system is adopted. Through unlabeled OCT image preprocessing, combined with the ViT-Base architecture and multi-expert adaptive reconstruction module, a two-stage training is carried out using random mask and attention-guided strategy. The encoder-decoder model is optimized by combining pixel-level mean square error and gradient loss to enhance the learning of the fine structure of the retina.
It significantly improves the accuracy and generalization ability of retinal OCT image segmentation, enabling high-precision segmentation of retinal interlayer boundaries and lesion areas with a small amount of labeled data.
Smart Images

Figure CN121641358B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the field of image processing technology, and specifically to a structure-aware mask reconstruction learning system and method for OCT image segmentation. Background Technology
[0002] Optical coherence tomography (OCT) is an important medical imaging technique that provides non-invasive, high-resolution cross-sectional images of the retina, crucial for the diagnosis and management of various retinal diseases. Accurate segmentation of the retinal structures using OCT images is a key step in clinical evaluation. Currently, supervised learning methods based on deep learning have achieved significant results in OCT image segmentation tasks.
[0003] For example, the invention patent with publication number CN113724262A discloses a method for segmenting choroidal neovascularization (CNV) in retinal OCT images, including: inputting the retinal OCT image into a multi-scale information fusion network, and segmenting the retinal OCT image through the multi-scale information fusion network; the multi-scale information fusion network includes an encoder-decoder network, a multi-scale adaptive perceptual deformation module (MAD), and a semantic detail aggregation module (SDA). The encoder-decoder network includes an encoder and a decoder. The encoder is used to extract semantic information and global features from the retinal OCT image. The MAD is placed on top of the encoder to aggregate contextual information. The SDA connects the encoder and decoder to fuse multi-level semantic information. The decoder is used to restore spatial resolution. A semi-supervised network structure based on a pseudo-labeled data augmentation strategy, SemiMF-Net, further improves CNV segmentation accuracy by utilizing unlabeled data.
[0004] For example, the invention patent with publication number CN118229713A discloses a method for retinal vessel image segmentation by fusing OCT and OCTA images, including: acquiring OCT images, optical coherence tomography (OCTA) images, and projection maps; performing multi-stage projection and fusion on the OCT and OCTA images to generate multimodal projection maps and intermediate-stage fused projection maps; fusing the multimodal projection maps to generate a first feature map; fusing the intermediate-stage fused projection maps to generate a second feature map; adding the projection map, the first feature map, and the second feature map to generate a third feature map; and performing projection segmentation on the first feature map, the second feature map, and the third feature map to obtain image segmentation results for retinal vessels (RV) and the foveal avascular zone (FAZ).
[0005] However, these methods heavily rely on large-scale, high-quality manually labeled datasets. But these pathological regions in fundus images are usually small in size and scattered across the entire retina, which greatly increases the difficulty of diagnosis; and the labeling process of retinal images requires the participation of experienced ophthalmologists, which is not only costly and time-consuming, but also results in a scarcity of publicly available high-quality labeled datasets, thus limiting the generalization ability and robustness of supervised learning models.
[0006] To reduce reliance on labeled data, self-supervised learning has emerged as a promising approach. Among these, mask reconstruction learning methods learn the intrinsic structural representation of an image by having the model reconstruct the masked portion of the image, demonstrating significant potential in various visual tasks. However, directly applying existing mask reconstruction learning methods to retinal OCT image analysis faces the following challenges:
[0007] Masking strategies are inefficient: Traditional masked image modeling (MIM) methods typically employ a high proportion of random masking strategies. However, lesions in OCT images (such as edema and drusen) are usually small in size and spatially concentrated, while most of the area is background. Random masking is likely to excessively obscure background areas while insufficiently masking lesion areas containing key clinical information, thereby weakening the model's ability to learn key structural features.
[0008] Reconstruction detail loss: Commonly used reconstruction loss functions, such as mean squared error, tend to recover low-frequency information (smooth regions) of the image, but perform poorly in reconstructing high-frequency details such as boundaries and textures. This leads to blurred boundaries in the reconstructed image, which is not conducive to the model learning accurate inter-layer boundary information of the retina and affects the accuracy of downstream segmentation tasks.
[0009] High data requirements: Most self-supervised learning methods require massive amounts of unlabeled data for pre-training to achieve ideal results. However, in the medical field, obtaining large-scale, high-quality OCT datasets remains extremely difficult due to factors such as patient privacy and device heterogeneity.
[0010] Therefore, there is an urgent need to develop a data-efficient, self-supervised learning method specifically designed for retinal OCT images to address the problems of insufficient ability of existing technologies to learn structural representations and poor restoration of boundary details on limited datasets. Summary of the Invention
[0011] The purpose of this invention is to provide a structure-aware mask reconstruction learning system and method for OCT image segmentation. By using unlabeled OCT images in the pre-training stage, a neural network model with an encoder-decoder structure is trained through self-supervised learning to perform retinal OCT image segmentation. This addresses the problems of low data efficiency, weak perception of fine retinal structures, and failure of masking strategies to effectively focus on clinically critical areas in existing self-supervised learning methods applied to retinal OCT image segmentation.
[0012] To achieve the above-mentioned objectives, an embodiment provides a structure-aware mask reconstruction learning system for OCT image segmentation, comprising:
[0013] A retinal OCT image preprocessing unit is used to preprocess unlabeled retinal OCT images as input and divide them into multiple image blocks;
[0014] The first-stage training unit is used to mask image patches using a random masking strategy. The unmasked image patches are input into the encoder. The encoder is based on the ViT-Base architecture and introduces a multi-expert adaptive reconstruction module to extract local and global features from the unmasked image patches and fuse them. The fused features are input into the decoder. The decoder is based on a lightweight ViT architecture and is used to reconstruct the masked image patches to obtain the first reconstructed image.
[0015] The second-stage training unit is used to mask image patches using an attention-guided masking strategy. The unmasked image patches are then input into the encoder and decoder trained in the first stage to reconstruct the second reconstructed image.
[0016] The retinal OCT image segmentation unit is used to extract features from the encoder trained in the second stage by inputting the retinal OCT image to be segmented. The extracted features are then input into the segmentation decoder to obtain the final segmentation prediction map.
[0017] In one embodiment, the preprocessing in the retinal OCT image preprocessing unit includes: random region cropping, scaling, flipping, tensor transformation, normalization, and standardization.
[0018] In one embodiment, when a random masking strategy is used to mask image patches in the first-stage training unit, a masking rate of 75% is adopted.
[0019] In one embodiment, in the first-stage training unit, inputting the unmasked image patch into the encoder includes: obtaining initial features from the unmasked image patch through linear embedding, adding positional encoding to the initial features to preserve spatial positional information, and obtaining basic features through a multi-layer Transformer in the ViT-Base architecture.
[0020] In one embodiment, the multi-expert adaptive reconstruction module in the first-stage training unit includes a local expert submodule, a global expert submodule, and a feature transformer.
[0021] The local expert submodule employs a depthwise separable convolutional network to extract local features, which include edges, textures, and corners.
[0022] The global expert submodule uses the Transformer module and employs a self-attention mechanism to capture long-distance dependencies and extract global features.
[0023] The feature transformer employs a multilayer perceptron and is used to perform nonlinear transformations and inter-channel interactions on the features extracted by the encoder.
[0024] In one embodiment, in the first-stage training unit, the multi-expert adaptive reconstruction module generates three weight coefficients for the local expert, the global expert, and the general feature transformer through the routing mechanism of the attention map; after normalizing the weight coefficients, they are used to adaptively fuse the outputs of the local expert submodule, the global expert submodule, and the feature transformer.
[0025] In one embodiment, the first-stage training unit and the second-stage training unit train the encoder and decoder by minimizing the structure-aware loss function. The structure-aware loss function is a weighted sum of pixel-level mean square error loss and gradient loss, and its calculation formula is as follows:
[0026] ,
[0027] in, The structure-aware loss function is... The pixel-level mean square error loss function is... For the original image, To reconstruct the image, Let be the gradient loss function. The gradient map of the original image. To reconstruct the gradient map of the image, These are the loss weights, used to balance the pixel-level mean squared error loss function and the gradient loss function.
[0028] In one embodiment, the pixel-level mean squared error loss is calculated by applying a mean squared error loss function to the reconstructed image and the original image to calculate the difference between pixels, which is used to evaluate the error of each pixel, as follows:
[0029] ,
[0030] in, and They represent the first Reconstructed pixel values and original pixel values of each image patch The dimension representing the vector of each image patch. This represents the total number of image patches in an image. It is a binary mask indicator, when When, it indicates the first Each image patch is masked.
[0031] In one embodiment, the gradient loss is calculated using the Scharr operator from both the horizontal and vertical directions. The calculation considers both the grayscale differences between directly adjacent pixels and the weighted summation of the pixel's neighborhood using a weighted convolution kernel. The calculated gradients are then used to calculate the loss using a gradient function, as shown in the following formula:
[0032] ,
[0033] in, and These represent the gradient magnitudes at each pixel in the original and reconstructed images, respectively. and These represent the height and width of the image, respectively. This represents the total number of pixels in the image. Representing coordinates At this point, the gradient map of the original image. gradient magnitude, Representing coordinates At this point, reconstruct the image gradient map. The gradient magnitude.
[0034] In one embodiment, the second-stage training unit performs masking of image blocks using an attention-guided masking strategy, including: using the attention weight matrix output from the last multi-head self-attention layer of the encoder after the first-stage training, aggregating and mapping the attention weight matrix back to the spatial layout of the original image blocks to generate an attention map, and using the attention map as a guiding signal to perform attention-guided masking on image blocks with high attention weights.
[0035] In one embodiment, the segmentation decoder in the retinal OCT image segmentation unit includes multiple alternately stacked upsampling layers and residual blocks, used to progressively upsample the extracted features and restore them to the same spatial resolution as the retinal OCT image to be segmented, outputting the final segmentation prediction map; the upsampling layers use transposed convolution to expand the spatial size of the extracted features and reduce the number of channels; the residual blocks include two convolutional layers and a batch normalization layer, and establish skip connections between the input and output to enhance the nonlinear expressive power of the extracted features and prevent gradient vanishing.
[0036] In one embodiment, the segmentation prediction map and the ground truth label map are compared using a segmentation loss function to calculate the error, which is used to optimize the segmentation decoder. The segmentation loss is a weighted hybrid loss function of Dice loss and cross-entropy loss, wherein Dice loss is used to measure the consistency between the segmentation prediction map and the ground truth label map in terms of spatial overlap, and cross-entropy loss is used to evaluate pixel-level classification accuracy.
[0037] This invention also provides a structure-aware mask reconstruction learning method for OCT image segmentation, which includes the following steps using the aforementioned structure-aware mask reconstruction learning system for OCT image segmentation:
[0038] Step 1: Preprocess the unlabeled retinal OCT image as input and divide it into multiple image blocks;
[0039] Step 2: The image blocks are masked using a random masking strategy. The unmasked image blocks are input into the encoder. The encoder is based on the ViT-Base architecture and introduces a multi-expert adaptive reconstruction module to extract local and global features from the unmasked image blocks and perform adaptive fusion. The fused features are input into the decoder. The decoder is based on a lightweight ViT architecture and is used to reconstruct the masked image blocks to obtain the first reconstructed image.
[0040] Step 3: Generate an attention map based on the self-attention weights of the encoder trained in Step 1. Based on the attention map, perform attention-guided masking on image blocks with high attention weights using an attention-guided masking strategy. Input the unmasked image blocks into the encoder and decoder for reconstruction. Further train the encoder to obtain the second reconstructed image.
[0041] Step 4: Input the retinal OCT image to be segmented into the encoder trained in Step 3 for feature extraction, and input the extracted features into the segmentation decoder to obtain the final segmentation prediction map.
[0042] Compared with the prior art, the beneficial effects of the present invention include at least the following:
[0043] To address the shortcomings of existing self-supervised learning methods in retinal OCT image analysis, such as low data efficiency, weak perception of fine structures like interlayer boundaries, and the failure of masking strategies to effectively focus on key retinal structural regions, a structure-aware mask reconstruction learning system and method for OCT image segmentation is proposed. In the pre-training phase, an encoder integrating ViT-Base and a multi-expert adaptive reconstruction module is used to perform mask reconstruction on unlabeled OCT images. The pre-training phase employs a two-stage training process using a random masking strategy and an attention-guided dynamic masking strategy. This allows the encoder-decoder to transition from random learning to focusing on key structural regions. Optimization is achieved using a structure-aware loss function composed of pixel loss and gradient loss, forcing the model to learn the fine boundaries and texture information of the retina. After training, the encoder, having learned rich structural knowledge, is used as the backbone network for transfer learning and combined with the segmentation decoder for downstream tasks. Fine-tuning is performed on a small amount of labeled data, significantly improving the accuracy of downstream segmentation. Attached Figure Description
[0044] To more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the accompanying drawings used in the description of the embodiments or the prior art will be briefly introduced below.
[0045] Figure 1 This is a schematic diagram of the structure-aware mask reconstruction learning system for OCT image segmentation provided by the present invention.
[0046] Figure 2 This is a schematic diagram illustrating the principle of the structure-aware mask reconstruction learning system for OCT image segmentation provided by the present invention; wherein, Figure 2 Figure (a) in the figure is a schematic diagram of the preprocessing process for an unlabeled retinal OCT image dataset. Figure 2 Figure (b) shows a schematic diagram of the structure of an asymmetric encoder-decoder neural network model. Figure 2 Figure (c) shows the training process of the first and second stage units. Figure 2 Figure (d) shows the process of validating the generalization of the pre-trained model using fine-tuned data.
[0047] Figure 3 This is a flowchart illustrating the structure-aware mask reconstruction learning method for OCT image segmentation provided by the present invention.
[0048] Figure 4 This diagram illustrates the segmentation results of the self-supervised learning method based on OCT images provided in this embodiment of the invention on three OCT datasets. Detailed Implementation
[0049] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are merely illustrative of the invention and do not limit the scope of protection of this invention.
[0050] To address the problems of low data efficiency, weak perception of fine retinal structures (especially inter-slice boundaries), and failure of masking strategies to effectively focus on clinically critical areas in existing self-supervised learning methods applied to retinal OCT image segmentation, this embodiment provides a structure-aware mask reconstruction learning system for OCT image segmentation, such as... Figure 1 As shown, the structure-aware mask reconstruction learning system includes: a retinal OCT image preprocessing unit, a first-stage training unit, a second-stage training unit, and a retinal OCT image segmentation unit.
[0051] The retinal OCT image preprocessing unit is used to preprocess unlabeled retinal OCT images as input and segment them into multiple image patches. Specifically:
[0052] In the embodiments, such as Figure 2 As shown in Figure (a), the unlabeled retinal OCT image dataset was preprocessed. The specific unlabeled dataset includes four categories: normal, diabetic macular edema (DME), drusen, and choroidal neovascularization (CNV).
[0053] A random region cropping operation is performed on the retinal OCT image, with the cropped area's proportion to the original image area randomly ranging from 0.2 to 1.0. The cropped image is then scaled to a predetermined input size (224×224) using an interpolation algorithm. This step simulates inputs at different field of view and scales during model training, improving the model's adaptability to local features and spatial deformations.
[0054] Since OCT images have a certain degree of left-right symmetry, the images are flipped horizontally with a 50% probability to increase data diversity and reduce the model's dependence on the image structure orientation.
[0055] The processed image data is converted from pixel format to tensor format in the PyTorch framework, and the pixel values are normalized from the range of 0–255 to the range of 0–1 to facilitate input processing for neural network models.
[0056] The transformed image tensor was standardized, with the mean of each channel set to 0.191 and the standard deviation to 0.202. The same parameters were applied to all three channels for normalization. This process stabilized the distribution of the input data, improving the model's convergence speed and performance.
[0057] Next, the preprocessed image is divided into a series of fixed-size, non-overlapping image blocks. In this embodiment, the image is divided into 16×16 image blocks.
[0058] The first-stage training unit masks image patches using a random masking strategy. The unmasked image patches are then input into the encoder, which is based on a ViT-Base architecture and incorporates a multi-expert adaptive reconstruction module. This module extracts and fuses local and global features from the unmasked image patches. The fused features are then input into the decoder, which uses a lightweight ViT architecture to reconstruct the masked image patches, resulting in the first reconstructed image. Details are as follows:
[0059] In the embodiments, such as Figure 2 As shown in Figure (b) I, an asymmetric encoder-decoder neural network model was constructed for learning retinal OCT image mask reconstruction. The segmented image patches were linearly embedded into vectors and positional encoding was added to preserve their spatial location information before being fed into a multi-layer Transformer encoder for processing.
[0060] In the first stage of pre-training, random masks are used to mask the input image patches, with the same 75% masking rate as the classic MAE, which prompts the model to learn the global semantic information of the image. Then, the remaining visible image patches are fed into the encoder.
[0061] The encoder employs the classic ViT-Base architecture. For retinal OCT images, a multi-expert adaptive reconstruction module is introduced into the deep network of the encoder to improve feature extraction. This module uses a multi-head attention mechanism (MH-SA) to overcome the limitations of a single network structure in processing complex medical images. This module integrates three complementary expert networks in parallel to perform multi-dimensional and multi-scale analysis of the input features.
[0062] The Local Expert submodule employs a depthwise separable convolutional network to extract local features. Its advantage lies in its ability to efficiently extract local structural information of images, such as edges, textures, and corners, at a relatively low computational cost. For OCT image segmentation, this expert is particularly suitable for accurately capturing the fine boundaries between different layers of the retina and the morphological features of tiny lesions.
[0063] The global expert submodule employs the standard Transformer module and leverages its self-attention mechanism to capture long-distance dependencies between any two locations in an image, establishing a global receptive field. This is crucial for understanding the extensive structural deformations and contextual relationships in OCT images caused by diseases such as macular edema.
[0064] The general feature transformer employs a multilayer perceptron (MLP). This expert is responsible for performing nonlinear transformations on features and deep information interaction between channels, serving as a supplement and refinement to local and global feature extraction.
[0065] Furthermore, the routing mechanism of the multi-expert adaptive reconstruction module employs an attention-map-based intelligent router. This mechanism does not simply average or concatenate the outputs of all experts, but rather generates three weight coefficients for the local expert, the global expert, and the general feature transformer based on the characteristics of the current input features (respectively...). w2, w3, w1 The contribution of each expert is determined dynamically and adaptively.
[0066] The features extracted by the encoder are fed into the decoder. In order to improve the training speed and enable the encoder to extract features better, the decoder adopts a lightweight ViT, which consists of eight Transformer blocks. The decoder reconstructs the masked region and outputs the reconstructed masked image block.
[0067] The second-stage training unit masks image patches using an attention-guided masking strategy. The unmasked image patches are then input into the encoder and decoder trained in the first stage for reconstruction, resulting in the second reconstructed image. Details are as follows:
[0068] The second stage, such as Figure 2 As shown in II of Figure (b), after 120 rounds of training with random masking, the masking method is switched from a random strategy to an attention-guided strategy. In this attention-guided strategy, the output of the last multi-head self-attention layer of the encoder is used to generate an attention map, which characterizes the importance of the model's learning of different regions in the image. Subsequently, this attention map is used as a guiding signal to prioritize masking image patches corresponding to regions with higher attention weights, thereby guiding the model to focus on learning information-rich retinal structural features. To adapt to the evolution of the model's state during training, the attention map is updated every 40 rounds to achieve adaptive adjustment of the feature-focused regions. Based on this, the unmasked image patches are input into an asymmetric encoder-decoder structure for image reconstruction, and the specific steps are consistent with the first stage.
[0069] Furthermore, such as Figure 2 As shown in Figure (c), both the first and second training phases train the encoder and decoder by minimizing the structure-aware loss function, which consists of a pixel-level mean square error loss function and a gradient loss function, calculated as follows:
[0070] ,
[0071] in, The structure-aware loss function is... The pixel-level mean square error loss function is... For the original image, To reconstruct the image, Let be the gradient loss function. The gradient map of the original image. To reconstruct the gradient map of the image, These are the loss weights, used to balance the pixel-level mean squared error loss function and the gradient loss function.
[0072] The pixel-level mean squared error loss is calculated by applying the mean squared error loss function to the reconstructed image and the original image to determine the difference between pixels, and is used to evaluate the error of each pixel. The calculation is as follows:
[0073] ,
[0074] in, and They represent the first Reconstructed pixel values and original pixel values of each image patch The dimension representing the vector of each image patch. This represents the total number of image patches in an image. It is a binary mask indicator, when When, it indicates the first Each image patch is masked.
[0075] Pixel-level mean squared error loss functions independently evaluate the error of each pixel, but this often leads to overly smooth reconstructed images, especially in high-frequency detail areas such as edges and textures, because these loss functions do not adequately consider the structural information of the image. Therefore, gradient loss is introduced to compensate for this.
[0076] In the specific implementation of gradient calculation, this invention fully considers the complexity of image quality in real-world application scenarios, especially addressing the noise artifact problem commonly found in retinal OCT images. The original gradient calculation method based on simple differences is extremely sensitive to noise, easily misidentifying noise as edge information, thus interfering with the correct optimization direction of the model.
[0077] To address this issue, the gradient loss employs the Scharr operator to calculate the pixel gradient from the horizontal direction Kx and the vertical direction Ky, respectively. The Scharr operator is an improved discrete differential operator that combines Gaussian smoothing and gradient differentiation. Its core advantage lies in that it considers both the grayscale difference between a pixel and its direct neighboring pixels during calculation, and also uses a 3×3 convolution kernel to perform a weighted summation of the pixel's neighborhood. The calculated gradient is then used to calculate the loss using a gradient function. The loss calculation formula is as follows:
[0078] ,
[0079] in, and These represent the gradient magnitudes at each pixel in the original and reconstructed images, respectively. and These represent the height and width of the image, respectively. This represents the total number of pixels in the image. Representing coordinates At this point, the gradient map of the original image. gradient magnitude, Representing coordinates At this point, reconstruct the image gradient map. The gradient magnitude.
[0080] In the above method, through attention masking and structure-aware optimization mechanism during the pre-training stage, the model of the present invention can generate more discriminative deep feature representations while maintaining global fidelity.
[0081] The retinal OCT image segmentation unit is used to extract features from the encoder trained in the second stage by inputting the retinal OCT image to be segmented. The extracted features are then input into the segmentation decoder to obtain the final segmentation prediction map.
[0082] In the embodiments, such as Figure 2 As shown in Figure (d), the fine-tuning data used included private and public datasets, with the public datasets including HCMS, RETOUCH, and AROI. The generalization ability of the pre-trained model was verified through experiments on multiple retinal layer and lesion region segmentation datasets.
[0083] Specifically, after pre-training, the original decoder is discarded, while the encoder, which has solidified knowledge of retinal structure, is retained and reused as a high-performance backbone network. For segmentation tasks, the encoder is initialized using the weights from the pre-training phase, and a dedicated segmentation decoder is then connected after the encoder to build a complete encoder-decoder segmentation architecture.
[0084] The segmentation decoder progressively upsamples the deep feature maps output by the encoder and restores them to the same spatial resolution as the input image, thereby generating corresponding pixel-level segmentation results. This decoder consists of multiple upsampling layers and residual blocks stacked alternately, maintaining the continuity and stability of feature representation while progressively restoring spatial information. The upsampling layers use transposed convolutions to expand the spatial size of the encoded features and reduce the number of channels. The residual blocks consist of two convolutional layers and a batch normalization layer, establishing skip connections between the input and output to enhance the non-linear expressiveness of the features and prevent gradient vanishing. After four layers of such upsampling and residual blocks, the final segmentation prediction map is output, with the number of output channels equal to the number of target segmentation categories.
[0085] The segmentation prediction map and its corresponding label map are compared using a segmentation loss function to calculate the error. This segmentation loss is a weighted hybrid loss function of Dice loss and cross-entropy loss. Dice loss measures the consistency between the predicted result and the ground truth region in terms of spatial overlap, effectively handling class imbalance; cross-entropy loss evaluates pixel-level classification accuracy. The combination of these two methods takes into account both global contours and local details, significantly improving the model's segmentation performance and stability.
[0086] By fine-tuning on the labeled dataset for the corresponding task, only the number of target segmentation categories needs to be modified. This allows the general features learned by the encoder to be quickly adapted to specific tasks, thereby efficiently achieving high-precision segmentation of the retinal layer or lesion areas.
[0087] like Figure 3 As shown, based on the above-mentioned structure-aware mask reconstruction learning system for OCT image segmentation, this invention also provides a structure-aware mask reconstruction learning method for OCT image segmentation, comprising the following steps:
[0088] Step 1: Preprocess the unlabeled retinal OCT image as input and segment it into multiple image blocks;
[0089] Step 2: The image blocks are masked using a random masking strategy. The unmasked image blocks are input into the encoder. The encoder is based on the ViT-Base architecture and introduces a multi-expert adaptive reconstruction module to extract local and global features from the unmasked image blocks and perform adaptive fusion. The fused features are input into the decoder. The decoder is based on a lightweight ViT architecture and is used to reconstruct the masked image blocks to obtain the first reconstructed image.
[0090] Step 3: Generate an attention map based on the self-attention weights of the encoder trained in Step 1. Based on the attention map, perform attention-guided masking on image blocks with high attention weights using an attention-guided masking strategy. Input the unmasked image blocks into the encoder and decoder for reconstruction. Further train the encoder to obtain the second reconstructed image.
[0091] Step 4: Input the retinal OCT image to be segmented into the encoder trained in Step 3 for feature extraction, and input the extracted features into the segmentation decoder to obtain the final segmentation prediction map.
[0092] To better illustrate the effectiveness of the system and method provided by this invention, qualitative and quantitative evaluations were performed on three OCT datasets, and the segmentation results on the three OCT datasets are presented. Figure 4 As shown, the first row represents the example image that has been segmented, the second row represents the segmentation label of the corresponding image, the third row represents the segmentation result without self-supervised learning, and the fourth row represents the segmentation result of this method. Compared with the method of directly training the segmentation without self-supervised learning, this method can achieve more accurate segmentation of the retinal OCT layer and the lesion area and has better generalization performance.
[0093] In summary, this method addresses the shortcomings of existing self-supervised learning methods in retinal OCT image segmentation, including low data efficiency, weak perception of fine structures such as interlayer boundaries, and the failure of masking strategies to effectively focus on retinal structural regions. Firstly, during the pre-training stage, this invention utilizes a powerful encoder integrating ViT-base and a multi-expert adaptive reconstruction module to perform masked reconstruction on unlabeled OCT images. This stage innovatively employs an attention-guided dynamic masking strategy, enabling the model to transition from random learning to focusing on key structural regions. This is further optimized using a structure-aware loss function composed of pixel loss and gradient loss, thereby forcing the model to learn the fine boundaries and texture information of the retina. After pre-training, the encoder, having learned rich structural knowledge, is transferred to a high-performance backbone network and combined with the segmentation decoder for downstream segmentation tasks, undergoing fine-tuning on a small amount of labeled data. By combining an attention-guided masking method with structure-aware loss optimization during the pre-training stage, this invention deeply enhances the encoder's feature representation capabilities during self-supervision, ultimately significantly improving the accuracy of downstream retinal layer segmentation.
[0094] The specific embodiments described above illustrate the technical solution and beneficial effects of the present invention in detail. It should be understood that the above description is only the most preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, additions, and equivalent substitutions made within the scope of the principles of the present invention should be included within the protection scope of the present invention.
Claims
1. A structure-aware mask reconstruction learning system for OCT image segmentation, characterized in that, include: A retinal OCT image preprocessing unit is used to preprocess unlabeled retinal OCT images as input and divide them into multiple image blocks; The first-stage training unit is used to mask image patches using a random masking strategy. The unmasked image patches are input into the encoder. The encoder is based on the ViT-Base architecture and introduces a multi-expert adaptive reconstruction module to extract local and global features from the unmasked image patches and fuse them. The fused features are input into the decoder. The decoder is based on a lightweight ViT architecture and is used to reconstruct the masked image patches to obtain the first reconstructed image. The second-stage training unit is used to mask image patches using an attention-guided masking strategy. The unmasked image patches are then input into the encoder and decoder trained in the first stage to reconstruct the second reconstructed image. The retinal OCT image segmentation unit is used to extract features from the encoder trained in the second stage by inputting the retinal OCT image to be segmented. The extracted features are then input into the segmentation decoder to obtain the final segmentation prediction map.
2. The structure-aware mask reconstruction learning system for OCT image segmentation according to claim 1, characterized in that, In the first stage training unit, the input of the unmasked image patch into the encoder includes: obtaining initial features from the unmasked image patch through linear embedding, adding position encoding to the initial features to preserve spatial position information, and obtaining basic features through a multi-layer Transformer in the ViT-Base architecture.
3. The structure-aware mask reconstruction learning system for OCT image segmentation according to claim 1, characterized in that, The multi-expert adaptive reconstruction module in the first-stage training unit includes a local expert sub-module, a global expert sub-module, and a feature transformer; The local expert submodule employs a depthwise separable convolutional network to extract local features, which include edges, textures, and corners. The global expert submodule uses the Transformer module and employs a self-attention mechanism to capture long-distance dependencies and extract global features. The feature transformer employs a multilayer perceptron and is used to perform nonlinear transformations and inter-channel interactions on the features extracted by the encoder.
4. The structure-aware mask reconstruction learning system for OCT image segmentation according to claim 3, characterized in that, In the first training phase, the multi-expert adaptive reconstruction module generates three weight coefficients for the local expert, the global expert, and the general feature transformer through the routing mechanism of the attention map. After normalizing the weight coefficients, they are used to adaptively fuse the outputs of the local expert submodule, the global expert submodule, and the feature transformer.
5. The structure-aware mask reconstruction learning system for OCT image segmentation according to claim 1, characterized in that, Both the first and second training phases train the encoder and decoder by minimizing the structure-aware loss function. This structure-aware loss function is a weighted average of pixel-level mean square error loss and gradient loss, calculated as follows: , in, The structure-aware loss function is... The pixel-level mean square error loss function is... For the original image, To reconstruct the image, Let be the gradient loss function. The gradient map of the original image. To reconstruct the gradient map of the image, These are the loss weights, used to balance the pixel-level mean squared error loss function and the gradient loss function.
6. The structure-aware mask reconstruction learning system for OCT image segmentation according to claim 5, characterized in that, The pixel-level mean squared error loss is calculated by applying the mean squared error loss function to the reconstructed image and the original image to determine the difference between pixels, and is used to evaluate the error of each pixel. The calculation is as follows: , in, and They represent the first Reconstructed pixel values and original pixel values of each image patch The dimension representing the vector of each image patch. This represents the total number of image patches in an image. It is a binary mask indicator, when When, it indicates the first Each image patch is masked; The gradient loss is calculated using the Scharr operator from both the horizontal and vertical directions. The calculation considers both the grayscale differences between directly adjacent pixels and the weighted summation of the pixel's neighborhood using a weighted convolution kernel. The calculated gradients are then used to calculate the loss using a gradient function. The loss calculation formula is as follows: , in, and These represent the gradient magnitudes at each pixel in the original and reconstructed images, respectively. and These represent the height and width of the image, respectively. This represents the total number of pixels in the image. Representing coordinates At this point, the gradient map of the original image. gradient magnitude, Representing coordinates At this point, reconstruct the image gradient map. The gradient magnitude.
7. The structure-aware mask reconstruction learning system for OCT image segmentation according to claim 1, characterized in that, In the second-stage training unit, image patches are masked using an attention-guided masking strategy. This includes: using the attention weight matrix output from the last multi-head self-attention layer of the encoder after the second-stage training, aggregating and mapping the attention weight matrix back to the spatial layout of the original image patches to generate an attention map, and using the attention map as a guiding signal to perform attention-guided masking on image patches with high attention weights.
8. The structure-aware mask reconstruction learning system for OCT image segmentation according to claim 1, characterized in that, The segmentation decoder in the retinal OCT image segmentation unit includes multiple alternately stacked upsampling layers and residual blocks, which are used to progressively upsample the extracted features and restore them to the same spatial resolution as the retinal OCT image to be segmented, outputting the final segmentation prediction map. The upsampling layers use transposed convolution to expand the spatial size of the extracted features and reduce the number of channels. The residual blocks include two convolutional layers and a batch normalization layer, and establish skip connections between the input and output to enhance the nonlinear expressive power of the extracted features and prevent gradient vanishing.
9. The structure-aware mask reconstruction learning system for OCT image segmentation according to claim 1, characterized in that, In the retinal OCT image segmentation unit, the output segmentation prediction map and the ground truth label map are compared using a segmentation loss function to calculate the error, which is used to optimize the segmentation decoder. The segmentation loss is a weighted hybrid loss function of Dice loss and cross-entropy loss, where Dice loss is used to measure the consistency between the segmentation prediction map and the ground truth label map in terms of spatial overlap, and cross-entropy loss is used to evaluate pixel-level classification accuracy.
10. A structure-aware mask reconstruction learning method for OCT image segmentation, characterized in that, The structure-aware mask reconstruction learning system for OCT image segmentation according to any one of claims 1 to 9 includes the following steps: Step 1: Preprocess the unlabeled retinal OCT image as input and divide it into multiple image blocks; Step 2: The image blocks are masked using a random masking strategy. The unmasked image blocks are input into the encoder. The encoder is based on the ViT-Base architecture and introduces a multi-expert adaptive reconstruction module to extract local and global features from the unmasked image blocks and perform adaptive fusion. The fused features are input into the decoder. The decoder is based on a lightweight ViT architecture and is used to reconstruct the masked image blocks to obtain the first reconstructed image. Step 3: Generate an attention map based on the self-attention weights of the encoder trained in Step 1. Based on the attention map, perform attention-guided masking on image blocks with high attention weights using an attention-guided masking strategy. Input the unmasked image blocks into the encoder and decoder for reconstruction. Further train the encoder to obtain the second reconstructed image. Step 4: Input the retinal OCT image to be segmented into the encoder trained in Step 3 for feature extraction, and input the extracted features into the segmentation decoder to obtain the final segmentation prediction map.
Citation Information
Patent Citations
CNV segmentation method in retina OCT image
CN113724262A
OCT and OCTA image fused retinal blood vessel image segmentation method
CN118229713A
Visual localization and anaphora segmentation method, system and device based on mask anaphora modeling and storage medium
CN118734091A
Medical image segmentation method and system based on guiding information and multi-dimensional attention mechanism
CN120807920A