STEM diffraction image self-supervised learning method and system based on multi-loss function fusion

Abstract

The invention relates to the technical field of image feature extraction, in particular to an STEM diffraction image self-supervised learning method and system based on multi-loss function fusion, and the method comprises the steps: firstly carrying out the targeted preprocessing of a diffraction image based on physical characteristics; then, a universal base model based on Vision Transform is constructed, a teacher-student framework is adopted, and self-supervised pre-training is carried out by fusing DINO loss, iBOT loss and Gram loss specially used for capturing feature correlation, so that the model learns a global mode and local details at the same time; finally, the pre-training model can be directly used for various downstream tasks such as grain segmentation and orientation recognition. According to the method, a transferable general feature extraction base model is provided for the 4D-STEM field for the first time, and the feature quality and the cross-task generalization ability are remarkably improved.

Description

A Self-Supervised Learning Method and System for STEM Diffraction Images Based on Multi-Loss Function Fusion Technical Field

[0001] This invention relates to the field of image feature extraction technology, and in particular to a self-supervised learning method and system for STEM diffraction images based on multi-loss function fusion. Background Technology

[0002] In materials science and crystallography research, the analysis of scanning transmission electron microscopy (STEM) diffraction images is a key technique for determining crystal structure and orientation. 4D-STEM (four-dimensional scanning transmission electron microscopy) technology can simultaneously acquire spatial location and diffraction information, providing rich data for material characterization. However, the 4D-STEM field currently faces the following key challenges: a lack of a universal feature extraction backbone model: existing 4D-STEM image analysis methods are primarily designed for specific tasks, lacking a universal, transferable feature extraction backbone model. Different research teams need to retrain models for each specific application (such as grain segmentation, orientation recognition, defect detection, etc.), leading to resource waste and inefficiency. The 4D-STEM field urgently needs a universal backbone model, similar to ImageNet pre-trained models in natural image processing or BERT / GPT in language modeling, capable of providing high-quality feature representations for various downstream tasks.

[0003] Existing methods struggle to capture multi-level features simultaneously: Current self-supervised learning methods (such as MAE and iBOT) primarily employ a single loss function when processing STEM diffraction images, such as mean squared error loss (MSE loss) or iBOT (ImageBERT Pre-training) loss. This makes it difficult to simultaneously capture both global features at the image level and detailed features at the local patch level. STEM diffraction images possess unique physical characteristics, requiring both the preservation of the global diffraction pattern structure and attention to subtle differences in local details. A single loss function cannot fully extract these multi-level feature information.

[0004] Feature representations lack generalization ability across datasets and tasks: Existing 4D-STEM image analysis methods exhibit poor generalization ability across different datasets and tasks. Models trained on specific datasets often perform poorly on other datasets, requiring retraining or extensive fine-tuning. This limits the practical application value of the models, especially when dealing with 4D-STEM data acquired using different materials, experimental conditions, or devices, where model performance degrades significantly. The lack of a pre-trained base model capable of learning general, robust feature representations makes cross-domain applications difficult. Summary of the Invention

[0005] In view of this, the purpose of this invention is to propose a self-supervised learning method and system for STEM diffraction images based on multi-loss function fusion, so as to solve the problems of existing technologies lacking a general feature extraction base model, making it difficult to capture multi-level features at the same time, and lacking the ability to generalize feature representations across datasets and tasks.

[0006] To achieve the above objectives, this invention provides a self-supervised learning method for STEM diffraction images based on multi-loss function fusion, comprising the following steps: Step S1, preprocessing the 4D-STEM diffraction image to obtain a preprocessed diffraction image; Step S2, constructing a base model based on the Vision Transformer architecture, using a teacher-student framework, and updating the teacher model parameters through exponential moving average; Step S3, based on the preprocessed diffraction image obtained in Step S1, performing self-supervised pre-training on the base model using fused multi-loss functions, wherein the multi-loss functions include at least DINO loss for learning global features at the image level, iBOT loss for learning local block-level detail features, and Gram loss for capturing correlations between features; Step S4, using the self-supervised pre-trained base model to extract feature representations of the input diffraction image.

[0007] Preferably, in step S1, the data preprocessing includes: S11, physical center-based centering transformation: extracting the physical center coordinates of each diffraction mode, calculating the affine transformation matrix that translates it to the image center, and applying bilinear interpolation for image transformation; S12, normalized cropping transformation: normalizing the image after centering transformation, normalizing the image pixel values ​​to the range of [0,1], multiplying by a magnification factor, and then cropping to the range of [0,1]; S13, logarithmic scaling transformation: applying log1p transformation to the image after normalized cropping transformation, and then re-normalizing it to the range of [0,1].

[0008] Preferably, the base model includes a patch embedding layer, a position encoding layer, multiple self-attention blocks, and a normalization layer.

[0009] Preferably, the multiple loss functions further include KoLeo loss for regularizing feature distribution and preventing feature collapse.

[0010] Preferably, in step S3, the self-supervised pre-training process includes a data augmentation step: applying DINO-style data augmentation to the preprocessed diffraction image to generate at least two global cropped images and multiple local cropped images; wherein, strong data augmentation is applied to the global cropped images and weak data augmentation is applied to the local cropped images.

[0011] Preferably, the self-supervised pre-training process includes: S31, inputting the global cropped image and the local cropped image into the student model and the teacher model respectively, extracting features and obtaining outputs via the corresponding DINO head and iBOT head; S32, calculating the weighted sum of the DINO loss, iBOT loss, Gram loss, and optionally KoLeo loss as the total loss; S33, updating the parameters of the student model through backpropagation based on the total loss, and updating the parameters of the teacher model through an exponential moving average strategy.

[0012] Preferably, step S3 further includes: periodically updating a Gram teacher network dedicated to Gram loss calculation, with the same update strategy as the teacher model.

[0013] Preferably, when calculating DINO loss and iBOT loss, the Sinkhorn-Knopp method or Softmax method is used to center the output of the teacher model.

[0014] Preferably, the feature representation of the extracted input diffraction image is used for downstream tasks, which include at least one of grain segmentation, orientation recognition, defect detection, or anomaly detection.

[0015] This invention also provides a self-supervised learning system for STEM diffraction images based on multi-loss function fusion, comprising: a preprocessing module for targeted data preprocessing of 4D-STEM diffraction images to obtain preprocessed diffraction images; a model building module for constructing a base model based on the Vision Transformer architecture, wherein the base model adopts a teacher-student framework and updates the teacher model parameters through exponential moving average; a training module for self-supervised pre-training of the base model based on the preprocessed diffraction images using fused multi-loss functions, wherein the multi-loss functions include at least DINO loss for learning global features at the image level, iBOT loss for learning local block-level detail features, and Gram loss for capturing correlations between features; and an application module for loading the self-supervised pre-trained base model and extracting feature representations of the input diffraction images.

[0016] The beneficial effects of this invention are as follows: 1. This invention is the first to construct a general feature extraction backbone model in the 4D-STEM field. This model, through large-scale self-supervised pre-training, can learn high-quality, transferable feature representations. Unlike existing methods that require retraining for each task, the backbone model provided by this patent can be directly applied to various downstream tasks (such as grain segmentation, orientation recognition, defect detection, anomaly detection, etc.), significantly improving research efficiency and resource utilization. This backbone model fills the gap in the lack of a general backbone in the 4D-STEM field, providing an important infrastructure for the development of the entire field; 2. This invention innovatively integrates three loss functions: DINO loss, iBOT loss, and Gram loss. DINO loss is responsible for learning global feature representations at the image level, iBOT loss focuses on learning detailed features at the local patch level, and Gram loss captures the correlation between features through the Gram matrix. This multi-level loss function design enables the pedestal model to learn both global and local features simultaneously, significantly improving the richness and discriminative power of feature representations and providing a stronger feature foundation for downstream tasks. 3. The general pedestal model constructed in this invention learns feature representations with strong generalization capabilities through large-scale self-supervised pre-training. This pedestal model can adapt to 4D-STEM data acquired from different materials, experimental conditions, and devices, demonstrating good transfer capabilities between different datasets and tasks. This eliminates the need for researchers to retrain the model for each new dataset or task; they can simply use the pre-trained pedestal model for feature extraction or minor fine-tuning, greatly improving the practicality and application scope of the method. 4. This invention designs a data preprocessing workflow specifically for STEM diffraction images, including physical center-based centering transformation, normalized cropping transformation, and logarithmic scaling transformation. These transformations fully consider the central symmetry and intensity distribution characteristics of STEM diffraction images, effectively preserving the physical meaning of the images while enhancing the generalization ability of the model; 5. This invention constructs a complete self-supervised learning system for STEM diffraction images. While maintaining high accuracy, this system improves training efficiency, providing a more efficient and accurate tool for materials science and crystallography research. Attached Figure Description

[0017] To more clearly illustrate the technical solutions in this invention or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are only for this invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0018] Figure 1 is a schematic diagram of the STEM diffraction image self-supervised learning method based on multi-loss function fusion according to an embodiment of the present invention; Figure 2 is a diagram of the STEM diffraction image self-supervised learning architecture according to an embodiment of the present invention; Figure 3 is a visualization result of kNN (k-nearest neighbor) classification on a given dataset based on the trained model according to an embodiment of the present invention; Figure 4 is a visualization of UMAP features at different iterations during the training process according to an embodiment of the present invention; Figure 5 is an interactive grain segmentation result based on the trained model according to an embodiment of the present invention. Detailed Implementation

[0019] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be further described in detail below with reference to specific embodiments.

[0020] It should be noted that, unless otherwise defined, the technical or scientific terms used in this invention should have the ordinary meaning understood by one of ordinary skill in the art to which this invention pertains. The terms "first," "second," and similar terms used in this invention do not indicate any order, quantity, or importance, but are merely used to distinguish different components. Terms such as "comprising" or "including" mean that the element or object preceding the word encompasses the elements or objects listed following the word and their equivalents, without excluding other elements or objects. Terms such as "connected" or "linked" are not limited to physical or mechanical connections, but can include electrical connections, whether direct or indirect. Terms such as "upper," "lower," "left," and "right" are used only to indicate relative positional relationships; when the absolute position of the described object changes, the relative positional relationship may also change accordingly. Examples

[0021] As shown in Figure 1, this embodiment of the specification provides a self-supervised learning method for STEM diffraction images based on multi-loss function fusion, including the following steps: Step S1: Preprocess the 4D-STEM diffraction image to obtain a preprocessed diffraction image; specifically, the data preprocessing includes: S11: Centering transformation based on the physical center: For each diffraction pattern, extract its physical center coordinates center(center_x, center_y), calculate the affine transformation matrix that translates the physical center to the image center, and apply bilinear interpolation for image transformation. The transformation formula is: translation = -center + image_size / 2; normalized_translation = (translation / image_size) * 2; where image_size represents the preset size of the output image, translation represents the pixel translation amount required to move the physical center to the center of the canvas, and normalized_translation represents the normalized translation parameter. By using bilinear interpolation, the original image is resampled using the above affine transformation matrix to generate a new image. This ensures that the centers of all diffraction patterns are aligned to the image center, preserving the physical meaning.

[0022] S12. Normalized Crop Transformation: The centered image is normalized, with pixel values ​​normalized to the range [0,1]. Then, it is multiplied by a multiply factor and finally cropped to the range [0,1]. This transformation can enhance the contrast of weak signals while preventing numerical overflow.

[0023] S13. Logarithmic scaling transformation: Apply the log1p transformation (log(1+x)) to the normalized cropped image, and then renormalize it to the [0, 1] range. This transformation can compress the dynamic range and highlight detailed information, making it particularly suitable for processing the large dynamic range characteristics of STEM diffraction images.

[0024] Step S2: Construct a base model based on the Vision Transformer architecture. This base model includes a patchbedding layer, a positional encoding layer (RoPE), multiple self-attention blocks, and a normalization layer. The network contains both a student model and a teacher model, with the teacher model updated from the student model using EMA (exponential moving average). This backbone network learns general, transferable feature representations through large-scale self-supervised pre-training, providing a foundation for various downstream tasks in the 4D-STEM domain.

[0025] Step S3: Based on the preprocessed diffraction image obtained in Step S1, the pedestal model is self-supervised pre-trained using a fusion of multiple loss functions. These multiple loss functions include at least DINO (Self-Supervised Visual Representation Learning) loss, iBOT (Image BERT Pre-training) loss, and Gram loss; where DINO loss is used to learn image-level global feature representations. In the teacher-student framework, the teacher output undergoes centering processing (Sinkhorn-Knopp or softmax), and the student output undergoes softmax processing, then cross-entropy loss is calculated. DINO loss is applied to both global crops and local crops, learning invariance by comparing feature consistency between different crops.

[0026] The iBOT loss is used to learn detailed features at the local patch level. The loss is calculated only on masked patches, and representations of local details are learned by comparing the outputs of the student and teacher on the masked patches. The iBOT loss uses the same centering method as the DINO loss. Masked patches are small, randomly masked image patches. Specifically, the input STEM diffraction image is first segmented into a series of regular, non-overlapping small squares, each called a patch. Random masking: A certain proportion of these patches are randomly selected, and their content is replaced with a special, learnable mask marker. These hollowed-out patches are the masked patches.

[0027] Gram loss is used to capture the correlation between features. The Gram matrix (feature similarity matrix) between the patch features of the student and teacher is calculated, and then the MSE loss is calculated. Gram loss can selectively remove negative similarity values, retaining only positive similarities to highlight positive correlations. In materials science and crystallography, each spot in a 4D-STEM diffraction image corresponds to a point in the reciprocal lattice space of the crystal, its position determined by the crystal's orientation and interplanar spacing, and its intensity related to the types and arrangement of atoms. Therefore, the relative positional relationships between spots constitute specific diffraction patterns (such as Kikuchi zones and diffraction rings), reflecting the crystal's symmetry and orientation; while the intensity distribution pattern of the spot group encodes the atomic-level structural information of the material. Traditional analysis methods heavily rely on expert experience to manually calibrate these relationships. This embodiment introduces GramLoss to explicitly constrain and preserve the correlation between feature channels by calculating the feature Gram matrix. In STEM image processing, this means that the model is guided during training to discover and maintain the spatial and intensity correlation rules hidden between diffraction spots that correspond to the essence of the crystal structure. This enables the base model of this embodiment to achieve excellent cross-dataset and cross-task generalization capabilities, and can be directly used for downstream tasks such as grain segmentation and orientation recognition.

[0028] As one implementation method, the centering process in this step employs either the Sinkhorn-Knopp method or the Softmax method. The Sinkhorn-Knopp method achieves uniform distribution through iterative optimization. First, Q = exp(teacher_output / teacher_temp) is calculated, where teacher_output is the output of the teacher network, specifically the vector obtained after the input image passes through the teacher ViT network and then through an MLP network called the "projection head," and teacher_temp is a scalar temperature parameter used to control the sharpness of the output distribution. Q = exp(teacher_output / teacher_temp) scales the original score using the temperature parameter, resulting in a matrix where each row corresponds to one image in a batch, and each column corresponds to a prototype. Then, row normalization (each row sum is 1 / K) and column normalization (each column is 1 / B) are iteratively performed, where K is the number of prototypes and B is the batch size. This method is suitable for high-dimensional feature spaces and large-batch training, ensuring that the model learns rich rather than degenerate feature representations, enabling the entire self-supervised training to proceed stably without crashing.

[0029] The Softmax method directly applies softmax to the centered output, i.e., softmax((teacher_output-center) / teacher_temp). The center is updated via EMA. By using a learnable center parameter, the output distribution is shifted globally to prevent excessive bias to one side. This method is more computationally efficient and suitable for low-dimensional feature spaces and mini-batch training.

[0030] As one implementation, multiple loss functions can also include KoLeo loss: used to regularize the feature distribution and prevent feature collapse. By maximizing the minimum distance between features, it encourages features to be uniformly distributed in the feature space. Other possible loss functions include contrastive loss, triplet loss, etc.

[0031] Gram loss can be replaced by other correlation losses, such as covariance loss, mutual information loss, etc.

[0032] The specific pre-training process is as follows: S31, Data loading: Load the preprocessed diffraction pattern data using a custom dataset class.

[0033] S32. Apply DINO-style data augmentation, including global crops (two large global crops) and local crops (multiple small local crops). Global crops use stronger augmentations (random cropping, color jittering, Gaussian blur, etc.), while local crops use weaker augmentations. For an input 4D-STEM diffraction map, the data augmentation module randomly generates two large global views: possibly randomly cropping different parts of the image and applying strong color jittering, Gaussian blur, and other transformations.

[0034] Multiple small local views: randomly crop smaller portions of the image, apply a weaker transformation, or simply scale the size.

[0035] For DINO Loss, the main goal is to ensure that the global semantic features learned by the student network from any cropped view (especially a local view) are consistent with the features obtained by the teacher network from the global view.

[0036] iBOT Loss, on the other hand, randomly masks off parts of the image in the local view, allowing the student network to predict the masked features based on the context and align them with the features of the corresponding positions in the teacher network.

[0037] S33. Input global crops and local crops into the student and teacher networks respectively, extract CLStoken features and patch token features, and obtain the output through the corresponding heads (DINO head and iBOT head).

[0038] S34. Based on the configured loss weights, calculate the weighted sum of DINO loss, iBOT loss, Gram loss, and KoLeo loss as the total loss: L total = w1·L DINO + w2·L iBOT + w3·L Gram + w4·L KoLeo Where L total L represents the total loss. DINO L iBOT L Gram L KoLeo These represent the DINO loss, iBOT loss, Gram loss, and KoLeo loss, respectively, with w1, w2, w3, and w4 being weights. The weights of the loss functions can be dynamically adjusted according to task requirements, or adaptively adjusted using a learning rate scheduler.

[0039] Backpropagation is performed on the total loss to update the student network parameters. Then, the teacher network parameters are updated via EMA: teacher_param = m * teacher_param + (1-m) * student_param, where m is the EMA momentum coefficient, and teacher_param and student_param represent the corresponding parameters in the teacher and student networks, respectively.

[0040] As one implementation, if Gram loss is enabled, the Gram teacher network can be updated periodically using the same update strategy as the teacher network. The Gram teacher network is an optional auxiliary technique component used to stabilize and improve the training performance of Gram Loss. Instead of directly comparing the student model and the main teacher model, Gram Loss compares the student model and this dedicated Gram Teacher network. This allows the main teacher model to focus on providing stable learning objectives for DINO Loss and iBOTLoss. The Gram Teacher network then specifically provides the objective for Gram Loss. The optimization objective of Gram Loss becomes clearer and more stable, no longer indirectly affected by parameter updates made by the main teacher model to adapt to other losses.

[0041] S4. Using the self-supervised pre-trained pedestal model, extract feature representations of the input diffraction image. Specifically, extract CLS token features as image-level feature representations or extract patchtoken features as local feature representations through the pedestal model network. The features provided by this pedestal model are universal and transferable, and can be directly applied to various downstream tasks without retraining.

[0042] Downstream applications include: Grain segmentation: Using features extracted from the pedestal model, regions belonging to the same grain are identified by calculating feature similarity at different locations. Cosine similarity is used to measure feature distance, and a similarity threshold is set for segmentation. The general features provided by the pedestal model significantly improve segmentation accuracy.

[0043] Orientation identification: Using features extracted from the pedestal model, the most similar training samples are found through feature retrieval (such as KNN), thereby determining the corresponding crystal orientation. General feature representation significantly improves retrieval accuracy.

[0044] Anomaly detection: Using features extracted from the pedestal model, anomalous diffraction patterns are identified through feature distance. General features are better able to capture the differences between normal and anomalous patterns.

[0045] Other applications: The pedestal model can also be applied to various 4D-STEM related tasks such as material classification, defect analysis, and structural prediction, demonstrating its versatility and practicality.

[0046] As an alternative implementation, besides Sinkhorn-Knopp and softmax, other centering methods can be used, such as direct normalization and L2 normalization. The centering method can be dynamically switched according to the training phase; for example, softmax can be used in the early stages of training, while Sinkhorn-Knopp can be used in the later stages.

[0047] In the data preprocessing stage, centering transformations can use other methods, such as centroid-based centering or maximum-based centering. Normalization methods can use other normalization strategies, such as Z-score normalization or Min-Max normalization. Logarithmic scaling can use other nonlinear transformations, such as square root transformation or power transformation.

[0048] As an alternative implementation, in addition to the currently used Vision Transformer architecture, other Vision Transformer variants, such as DeiT, can be used as the backbone of the base model.

[0049] In addition to Vision Transformer, CNN architectures (such as ResNet, EfficientNet, etc.) can also be used as the backbone of general base models.

[0050] Position coding can use other methods, such as learnable position coding, sinusoidal position coding, etc.

[0051] During pre-training, besides EMA updating the teacher, other update strategies can be used, such as direct copying and momentum updates. Data augmentation can use other augmentation strategies, such as MixUp, CutMix, and AutoAugment. The optimizer can use other optimization algorithms, such as AdamW, LAMB, and RAdam. Example

[0052] This embodiment provides a 4D-STEM self-supervised learning architecture based on Vision Transformer, as shown in Figure 2. It employs a Student-Teacher framework and integrates multiple loss functions. Main components: Input layer (top): The input is a 4D-STEM image, including global cropping and local cropping; Network layers: Student network (left): Vision Transformer, trainable; Teacher network (right): Vision Transformer, updating parameters via EMA (Exponential Moving Average); Student parameters are updated to Teacher parameters via EMA.

[0053] Feature Extraction Layer: Student and Teacher output CLS Token and Patch Tokens respectively. Head Layer: Student side: DINO Head, iBOT Head; Teacher side: DINO Head, iBOT Head; Loss Function Layer (four): DINO Loss: Based on global features, connecting the DINO Head of Student and Teacher; iBOT Loss: Based on local features, connecting the iBOT Head of Student and Teacher; Gram Loss: Feature correlation loss, connecting the feature layer of Student and Teacher; KoLeo Loss: Feature regularization loss, connecting the feature layer of Student and Teacher; Total Loss (bottom): Weighted sum of the four loss functions.

[0054] The data flow of the above architecture is as follows: the input image is simultaneously fed into the Student and Teacher networks. The two networks extract features and process them through their respective Headers. Each loss function calculates the difference between the Student and Teacher networks. All losses are summed into a total loss, which is used to optimize the Student network. The Student parameters are updated in the Teacher network through EMA.

[0055] This architecture combines global contrastive learning from DINO, local mask learning from iBOT, and feature constraints from Gram and KoLeo for self-supervised representation learning of 4D-STEM images.

[0056] Figure 3 shows the visualization results of kNN (k-nearest neighbor) classification based on a trained model on a given dataset, used to evaluate the quality of the learned feature representations.

[0057] Layout: Each row contains 1 test sample and 5 samples obtained from nearest neighbor retrieval, with cosine distance used to measure feature similarity.

[0058] Results: The five nearest neighbors of all five test samples matched correctly, with small class distance values ​​(0.002-0.027) and high feature similarity.

[0059] Example: Test sample 1 (category 374): The 5 nearest neighbors are all of category 374, with distances of 0.0078-0.0276. Test sample 3 (category 489): The 5 nearest neighbors are all of category 489, with distances of 0.0024-0.0117. Test sample 5 (category 260): The 5 nearest neighbors are all of category 260, with distances of 0.0022-0.0053. This result shows that the feature representation learned by the model can effectively distinguish 4D-STEM image patterns of different categories. Samples of the same category are clustered in the feature space, and different categories are well separated, verifying the effectiveness of the self-supervised learning method.

[0060] Figure 4 shows the UMAP feature visualizations obtained from experiments using Mg twinning data. This visualization is used to observe how the feature representation learned by the model changes during training. The overall layout is a grid: the left side is a bright-field image, serving as a spatial reference for 4D-STEM scanning; the right side shows the UMAP visualization results from different training iterations, with each column containing two rows. The top row is the UMAP projection map, reducing the high-dimensional features to 2D; the bottom row is the scan color map, mapping the UMAP colors back to the original scan space. Each pixel's color corresponds to its position color in the UMAP, used to observe the spatial distribution pattern of the features.

[0061] As can be seen from the combined visualization, self-supervised learning gradually learns meaningful feature representations on 4D-STEM data. In the early stages of training, the UMAP projection points are scattered, and the color changes in the scanned color map are random, indicating that the features have not yet formed a clear structure. As training progresses, the UMAP projections begin to show local clustering, and continuous regions appear in the scanned color map, indicating that the model is beginning to capture spatial patterns. In the later stages of training, the UMAP projections form clearer clusters and separations, and the scanned color map shows spatial structures corresponding to the brightfield image. Regions with similar colors are spatially continuous, indicating that the features have encoded spatial information. The correspondence between the scanned color map and the left-hand brightfield image further validates the effectiveness of the method.

[0062] In later iterations, the spatial patterns of the scanned color maps corresponded to the structural features (such as grain boundaries, defects, and different material regions) of the bright-field images, indicating that the features learned by the model not only preserved spatial information but also reflected the physical structure of the material. This correspondence was not obvious in early iterations but gradually became clearer in later iterations, suggesting that the training process gradually focused the model on meaningful physical features. The feature space evolved from disordered to ordered, from random distribution to meaningful clustering; the scanned color maps evolved from random colors to correspondences with spatial structures; and the feature representations improved from low quality to high quality. This progressive improvement demonstrates the effectiveness of the self-supervised learning framework. The model gradually learns the latent structures in 4D-STEM data through contrastive learning, providing a high-quality feature representation foundation for subsequent downstream tasks (such as material classification and defect detection).

[0063] Figure 5 shows the interactive grain segmentation results based on the trained model, demonstrating the model's performance in feature extraction and similarity calculation. The data used in Figure 5 is Au simulated data obtained through a simulation library, visually showcasing the model's effectiveness in the grain segmentation task. The segmentation is accurate, the boundaries are clear, and the similar diffraction pattern recognition is accurate, making it suitable for grain analysis and verification in materials science. The first figure shows a visualization of the original grain data. The second figure shows the segmentation mask obtained based on the similarity of clicked point features, highlighting pixels identified as belonging to the same grain. The third figure shows the diffraction pattern (heatmap form) at the clicked location, i.e., the original diffraction image fed into the model for feature calculation.

[0064] Those skilled in the art will recognize that the units and algorithm steps of the various examples described in conjunction with the embodiments disclosed in this application can be implemented in electronic hardware, or a combination of computer software and electronic hardware. Whether these functions are implemented in hardware or software depends on the specific application and design constraints of the technical solution. Those skilled in the art can use different methods to implement the described functions for each specific application, but such implementation should not be considered beyond the scope of this application.

[0065] In the embodiments provided in this application, it should be understood that the disclosed devices / terminal equipment and methods can be implemented in other ways. For example, the device / terminal equipment embodiments described above are merely illustrative. For instance, the division of modules or units is only a logical functional division, and in actual implementation, there may be other division methods. For example, multiple units or components may be combined or integrated into another system, or some features may be ignored or not executed. Furthermore, the displayed or discussed mutual coupling or direct coupling or communication connection may be through some interfaces; the indirect coupling or communication connection between devices or units may be electrical, mechanical, or other forms.

[0066] Furthermore, the functional units in the various embodiments of this application can be integrated into one processing unit, or each unit can exist physically separately, or two or more units can be integrated into one unit. The integrated unit can be implemented in hardware or as a software functional unit.

[0067] If the integrated module / unit is implemented as a software functional unit and sold or used as an independent product, it can be stored in a computer-readable storage medium. Based on this understanding, all or part of the processes in the methods of the above embodiments can also be implemented by a computer program instructing related hardware. The computer program can be stored in a computer-readable storage medium, and when executed by a processor, it can implement the steps of the various method embodiments described above. The computer program includes computer program code, which can be in the form of source code, object code, executable files, or certain intermediate forms. The computer-readable medium can include: any entity or device capable of carrying the computer program code, recording media, USB flash drives, portable hard drives, magnetic disks, optical disks, computer memory, read-only memory (ROM), random access memory (RAM), electrical carrier signals, telecommunication signals, and software distribution media, etc. It should be noted that the content included in the computer-readable medium can be appropriately added or removed according to the requirements of legislation and patent practice in the jurisdiction. For example, in some jurisdictions, according to legislation and patent practice, computer-readable media do not include electrical carrier signals and telecommunication signals.

[0068] The implementation of all or part of the processes in the methods of the above embodiments can also be accomplished by a computer program product. When the computer program product is run on a terminal device, the terminal device can implement the steps in the various method embodiments described above.

[0069] The embodiments described above are only used to illustrate the technical solutions of this application, and are not intended to limit it. Although this application has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some of the technical features; and these modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the spirit and scope of the technical solutions of the embodiments of this application, and should all be included within the protection scope of this application.

Claims

1. A self-supervised learning method for STEM diffraction images based on multi-loss function fusion, characterized in that, The method includes the following steps: Step S1, preprocessing the 4D-STEM diffraction image to obtain a preprocessed diffraction image; Step S2, constructing a base model based on the Vision Transformer architecture, using a teacher-student framework, and updating the teacher model parameters through exponential moving average; Step S3, based on the preprocessed diffraction image obtained in Step S1, performing self-supervised pre-training on the base model using a fusion of multiple loss functions, wherein the multiple loss functions include at least DINO loss for learning global features at the image level, iBOT loss for learning local block-level detail features, and Gram loss for capturing correlations between features; Step S4, using the base model completed by self-supervised pre-training, extracting the feature representation of the input diffraction image.

2. The self-supervised learning method for STEM diffraction images based on multi-loss function fusion according to claim 1, characterized in that, In step S1, the data preprocessing includes: S11, physical center-based centering transformation: extracting the physical center coordinates of each diffraction mode, calculating the affine transformation matrix that translates it to the image center, and applying bilinear interpolation for image transformation; S12, normalized cropping transformation: normalizing the image after centering transformation, normalizing the image pixel values ​​to the range of [0,1], multiplying by a magnification factor, and then cropping to the range of [0,1]; S13, logarithmic scaling transformation: applying log1p transformation to the image after normalized cropping transformation, and then re-normalizing it to the range of [0,1].

3. The self-supervised learning method for STEM diffraction images based on multi-loss function fusion according to claim 1, characterized in that, The base model includes a patch embedding layer, a position encoding layer, multiple self-attention blocks, and a normalization layer.

4. The self-supervised learning method for STEM diffraction images based on multi-loss function fusion according to claim 1, characterized in that, The multiple loss functions also include KoLeo loss for regularizing feature distribution and preventing feature collapse.

5. The self-supervised learning method for STEM diffraction images based on multi-loss function fusion according to claim 1, characterized in that, In step S3, the self-supervised pre-training process includes a data augmentation step: applying DINO-style data augmentation to the preprocessed diffraction image to generate at least two global cropped images and multiple local cropped images; wherein, strong data augmentation is applied to the global cropped images and weak data augmentation is applied to the local cropped images.

6. The self-supervised learning method for STEM diffraction images based on multi-loss function fusion according to claim 5, characterized in that, The self-supervised pre-training process includes: S31, inputting the global cropped image and the local cropped image into the student model and the teacher model respectively, extracting features and obtaining outputs through the corresponding DINO head and iBOT head; S32, calculating the weighted sum of the DINO loss, iBOT loss, Gram loss, and optionally KoLeo loss as the total loss; S33, updating the parameters of the student model through backpropagation based on the total loss, and updating the parameters of the teacher model through an exponential moving average strategy.

7. The self-supervised learning method for STEM diffraction images based on multi-loss function fusion according to claim 1, characterized in that, Step S3 also includes: periodically updating a Gram teacher network dedicated to Gram loss calculation, with the same update strategy as the teacher model.

8. The self-supervised learning method for STEM diffraction images based on multi-loss function fusion according to claim 1, characterized in that, When calculating DINO loss and iBOT loss, the Sinkhorn-Knopp method or Softmax method is used to center the output of the teacher model.

9. The self-supervised learning method for STEM diffraction images based on multi-loss function fusion according to claim 1, characterized in that, In step S4, the feature representation of the extracted input diffraction image is used for downstream tasks, which include at least one of grain segmentation, orientation recognition, defect detection, or anomaly detection.

10. A self-supervised learning system for STEM diffraction images based on multi-loss function fusion, characterized in that, include: The preprocessing module is used to perform targeted data preprocessing on the 4D-STEM diffraction image to obtain the preprocessed diffraction image; The model building module is used to build a base model based on the Vision Transformer architecture. The base model adopts a teacher-student framework and updates the teacher model parameters through exponential moving average. The training module performs self-supervised pre-training on the pedestal model based on the preprocessed diffraction image using a fusion of multiple loss functions. The multiple loss functions include at least the DINO loss for learning global features at the image level, the iBOT loss for learning local block-level detail features, and the Gram loss for capturing the correlation between features. The application module loads the pedestal model that has been self-supervised and pre-trained, and extracts the feature representation of the input diffraction image.

Images