A Method for Constructing and Applying a Metal Surface Defect Classification Model Based on a Structured State-Space Model

Abstract

This invention belongs to the field of industrial defect detection technology and discloses a method for constructing and applying a metal surface defect classification model based on a structured state-space model. The invention constructs a model structure combining deep convolution and bidirectional Mamba modules. Deep convolution is used to capture local spatial features of image data in parallel across spatial dimensions, while the long-sequence modeling capability of bidirectional Mamba is used to improve the upper limit and capacity of the model. Furthermore, a multi-scale feature fusion module M2SF with a learnable scaling factor is proposed, enabling the network to be robustly trained on datasets of various sizes simultaneously. The advantages of the proposed method lie in its simultaneous focus on local and global features, and the significant reduction in the adverse effects of unidirectional causal dependencies through the bidirectional Mamba structure, thus improving the accuracy of the classification model. The introduction of a pyramid structure and the fusion of multi-scale features fully consider the different sizes of the training dataset, enhancing the model's generalization ability.

Description

Technical Field

[0001] This invention belongs to the field of industrial defect detection technology, and more specifically, relates to a method for constructing a metal surface defect classification model based on a structured state space model and its application. Background Technology

[0002] Industrial defect detection is a crucial step in modern manufacturing, especially in metal processing, where it plays a vital role in controlling metal quality. Deep learning is widely used in industrial defect detection in modern production. By constructing deep neural networks to extract deep features from images, it enables automated defect detection. Deep neural networks possess powerful fitting capabilities, making them well-suited for handling irregular and variable defect samples. No specialized knowledge is required; efficient, accurate, and robust defect detection can be achieved simply by collecting sample data. Automated defect detection mainly consists of two tasks: defect classification and defect localization. Defect classification aims to determine the presence and category of defects in an image; defect localization analyzes the location of defects in the image, marking one or more defects with bounding boxes and assigning confidence scores. Complete defect detection comprises both defect localization and defect classification. Efficiently and accurately completing defect classification is fundamental to achieving real-time defect detection.

[0003] Existing neural networks used for defect classification mainly include Convolutional Neural Networks (CNNs), Generative Adversarial Networks (GANs), and Transformer networks. CNNs have strong local feature extraction capabilities, but their ability to understand global features is limited due to inductive bias. Visual architectures based on Transformers include Vision Transformer (ViT) and Swin-Transformer; ViT divides the image into small patches and processes these patches through a self-attention mechanism to capture and integrate global contextual information; Swin-Transformer combines local and global features through a sliding window mechanism. Transformer-based networks can achieve higher classification performance with fewer parameters and can handle extremely complex defect situations; however, the secondary complexity of the self-attention mechanism introduces significant computational overhead, and dividing the image into fixed-size patches leads to insufficient ability to capture local details. Furthermore, the metal defect dataset contains a large number of defects, and defects of the same type may have significant differences in appearance, such as tilted scratches and horizontal scratches. In addition, there are differences in lighting, background and materials during shooting, which may also cause the grayscale of defects of the same type to change. At the same time, defects between classes also have a certain degree of similarity. Directly using neural networks to classify defects also has the problem of low classification accuracy. Summary of the Invention

[0004] In view of the above-mentioned defects or improvement needs of the existing technology, the present invention provides a method and application for constructing a metal surface defect classification model based on a structured state space model, which aims to improve the accuracy and efficiency of metal surface defect classification and reduce production costs.

[0005] To achieve the above objectives, this invention provides a method for constructing a metal surface defect classification model based on a structured state-space model, comprising:

[0006] A neural network based on a structured state-space model is pre-trained using an image dataset, and then fine-tuned using a metal surface defect dataset to obtain a metal surface defect classification model; wherein, the neural network based on the structured state-space model includes:

[0007] The Patch layer is used to downsample the input image samples and increase the number of channels;

[0008] A multi-scale feature extraction module includes deep convolutional layers and bidirectional Mamba layers at different scales. The deep convolutional layers are used to capture local spatial feature maps of different scales in the output image of the Patch layer in parallel along the spatial dimension. The bidirectional Mamba layer includes two Mamba structures with fully connected layers removed and one fully connected layer. One Mamba structure is used to obtain forward output features based on one-dimensional data after downsampling and flattening the local spatial feature maps, and the other Mamba structure is used to obtain inverse output features based on data after flipping the one-dimensional data. The fully connected layer is used to linearize the features obtained by fusing the forward and inverse output features to obtain the feature map output by the bidirectional Mamba layer.

[0009] The multi-scale feature fusion module is used to fuse the feature map output by the deep convolutional layer and the feature map output by the bidirectional Mamba layer at multiple scales to obtain the fused feature map.

[0010] A classification head is used to generate classification results based on the fused feature map.

[0011] Furthermore, the deep convolutional layer is an MBConv layer, and the multi-scale feature extraction module includes four stages connected in a pyramid structure: Stage0-Stage3. The feature map scale output by Stage0-Stage3 gradually decreases, and the number of feature channels is doubled accordingly.

[0012] Stage 0 includes two sequentially connected MBConv layers, Stage 1 includes six sequentially connected MBConv layers, Stage 2 includes fourteen sequentially connected bidirectional Mamba layers, and Stage 3 includes two sequentially connected bidirectional Mamba layers.

[0013] The multi-scale feature fusion module is used to fuse the feature maps output from Stage0 to Stage3 to obtain the fused feature map.

[0014] Furthermore, the positive output feature y forward for:

[0015] y forward =x proj ⊙SiLU(z)

[0016] x proj =SSM(SiLU(Conv1d(x)))

[0017] The inverse output feature y backward for:

[0018]

[0019] Where SSM(·) represents the output of the main path of the Mamba structure after removing the fully connected layer, SiLU(·) is the SiLU activation function, Conv1d(·) represents a one-dimensional convolution operation, and ⊙ represents a dot product operation; x and z represent the input of the main path and the input of the residual path in the Mamba structure obtained by mapping the one-dimensional data through a linear layer, respectively; x b z b The inputs of the main path and the residual path in the Mamba structure are obtained by mapping the data after the one-dimensional data is flipped through a linear layer.

[0020] Furthermore, the multi-scale feature fusion module is a multi-scale feature fusion module M2SF with a learnable scaling factor, specifically including: an SE layer, a normalization layer, an average pooling layer, a scaling factor learned through training, and a merging layer;

[0021] The SE layer is used to adaptively adjust the feature map output by each Stage. The adjusted feature map is then processed by the corresponding normalization layer and average pooling layer to obtain feature maps of the same size. The feature map is multiplied by the corresponding scaling factor and then merged in the channel dimension by the merging layer to obtain the fused feature map.

[0022] Furthermore, the metal surface defect dataset is a dataset obtained by merging multiple metal surface defect datasets of different sizes.

[0023] Furthermore, based on the metal surface defect dataset, the pre-trained neural network was fine-tuned using transfer learning to obtain a metal surface defect classification model.

[0024] This invention also provides a method for classifying metal surface defects, including:

[0025] The image of the metal surface defect to be classified is input into the metal surface defect classification model to obtain the defect classification result; wherein, the metal surface defect classification model is constructed by any of the metal surface defect classification model construction methods described above.

[0026] The present invention also provides an electronic device, including a computer-readable storage medium and a processor;

[0027] The computer-readable storage medium is used to store executable instructions;

[0028] The processor is used to read executable instructions stored in the computer-readable storage medium to execute the metal surface defect classification model construction method described above, or to execute the metal surface defect classification method described above.

[0029] The present invention also provides a computer-readable storage medium having a computer program stored thereon, wherein the program, when executed by a processor, implements the metal surface defect classification model construction method as described in any of the preceding claims, or implements the metal surface defect classification method as described above.

[0030] The present invention also provides a computer program product, including a computer program that, when the computer program is run on a computer, causes the computer to execute the metal surface defect classification model construction method described above, or to execute the metal surface defect classification method described above.

[0031] In summary, the above-described technical solutions conceived in this invention can achieve the following beneficial effects:

[0032] (1) This invention combines CNN with Mamba, a structured state-space model with linear computational complexity and long sequence modeling capabilities, to construct a neural network model structure that integrates deep convolution and bidirectional Mamba modules. Specifically, this invention considers that when Mamba is used to understand metal surface defect images, its inherent causal mode in unidirectional modeling has a negative impact on the understanding of global features. To mitigate this disadvantage, the bidirectional Mamba designed in this invention processes the input from two directions. Compared to the causal mode of the original Mamba, the bidirectional Mamba in this invention is closer to the fully visible mode of Transformer, which avoids the limitations of causal constraints to a certain extent, is more suitable for global image understanding, and avoids the large computational overhead and insufficient ability to capture local details caused by using a Transformer-based network structure. Therefore, this invention utilizes deep convolution to capture local spatial features of image data in parallel in the spatial dimension, and then integrates global information through the long sequence modeling advantage of bidirectional Mamba. This allows the model structure to focus on both local and global features, greatly improving the accuracy of the metal surface defect classification model while balancing computational efficiency, and reducing production costs.

[0033] (2) Preferably, the present invention connects Mamba modules and deep convolutional modules of different scales together in a pyramid structure to build a new network architecture containing four stages. It has the best local two-dimensional spatial feature capture capability and the best global long sequence modeling capability, and avoids the secondary computational complexity of the attention mechanism. It achieves efficient and accurate classification performance on the metal surface defect dataset.

[0034] (3) As a preferred option, the multi-scale features extracted from the feature pyramid are fused through the multi-scale feature fusion module M2SF with a learnable scaling factor, which mixes features at different levels and adaptively adjusts the weights of each level for datasets of different scales, thereby improving the generalization ability of the classification model.

[0035] (4) As a preferred option, in order to cope with the variable situation in the industrial production process, the collected images may not be guaranteed to be of uniform size. The metal surface defect dataset used in this invention is a dataset obtained by merging multiple metal surface defect datasets of different sizes, which further improves the generalization ability of the model.

[0036] (5) Due to issues such as cost and standardization, industrial defect datasets are generally small in size. Using pre-trained models generated on general datasets for transfer learning reduces the training time and computational resource cost of metal surface defect classification models, accelerates model convergence, and is more suitable for industrial applications. Attached Figure Description

[0037] Figure 1 This is a flowchart of the method for constructing a metal surface defect classification model based on a structured state-space model in an embodiment of the present invention.

[0038] Figure 2 A schematic diagram of the Mamba structure and the bidirectional Mamba structure (Bi-Mamba) provided in the embodiments of the present invention.

[0039] Figure 3 A diagram of a neural network structure combining depthwise convolution and bidirectional Mamba provided for embodiments of the present invention.

[0040] Figure 4 The diagram shows the structure of the multi-scale feature fusion module M2SF provided in this embodiment of the invention.

[0041] Figure 5 The confusion matrix for defect classification in the ResNet18 model.

[0042] Figure 6 The confusion matrix for defect classification in the ViT-B model.

[0043] Figure 7 This is the model defect classification confusion matrix without the M2SF module in this embodiment of the invention.

[0044] Figure 8 This is the model defect classification confusion matrix that includes the M2SF module in this embodiment of the invention. Detailed Implementation

[0045] 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 and not intended to limit the invention. Furthermore, the technical features involved in the various embodiments of this invention described below can be combined with each other as long as they do not conflict with each other.

[0046] Example 1

[0047] like Figures 1-2 As shown, the method for constructing a metal surface defect classification model based on a structured state-space model in this embodiment of the invention includes:

[0048] A neural network based on a structured state-space model was pre-trained using an image dataset, and the pre-trained neural network was fine-tuned using a metal surface defect dataset to obtain a metal surface defect classification model.

[0049] In this embodiment of the invention, the image dataset used for pre-training the neural network is the ImageNet1K dataset. In other embodiments, other image datasets may also be selected to pre-train the neural network based on the structured state-space model.

[0050] As a preferred embodiment, this invention takes into account that some existing methods perform well on defect datasets of the same size, but their accuracy decreases and their generalization ability is weak when dealing with defect data of various sizes. To cope with the variable conditions in industrial production processes, the acquired images may not be guaranteed to be of uniform size. Therefore, the metal surface defect dataset used in this invention is a dataset obtained by merging multiple metal surface defect datasets of different sizes, in order to improve the generalization ability of the model.

[0051] In this embodiment of the invention, two labeled metal surface defect datasets of different sizes are merged and divided into training and validation sets proportionally. The datasets are sourced from the Northeastern University Surface Defect Database, a relatively authoritative publicly available dataset in the field of metal surface defect detection. One dataset, NEU-CLS, has a uniform pixel size of 200×200 and collects six typical surface defects of hot-rolled steel strips: rolling scale (Rs), patches (Pa), cracks (Cr), pitting (Ps), inclusions (In), and scratches (Sc). Each category contains 300 images, totaling 1800 image samples. The other dataset, NEU-CLS-64, has a uniform pixel size of 64×64 and is an extension of the NEU-CLS dataset, expanding the defect types to nine categories (Cr, Gg, In, Pa, Ps, Rp, Rs, Sc, Sp), where Gg represents grinding marks, Rp represents indented scale, and Sp represents scale peeling.

[0052] The neural network in this embodiment of the invention is a model structure based on a combination of a deep convolutional MBConv module and a bidirectional Mamba module, specifically including a patch layer, a feature extraction module, a multi-scale feature fusion module, and a classification head.

[0053] The Patch layer is used to downsample the input raw image and increase the number of channels.

[0054] The multi-scale feature extraction module includes deep convolutional layers of different scales and bidirectional Mamba layers. In this embodiment, the deep convolutional layer is an MBConv layer, i.e., mobile inverted bottleneck convolution (MBConv). The deep convolutional layers of different scales are used to capture local spatial features of the image processed by the Patch layer in parallel in the spatial dimension through deep convolution. After obtaining sufficient local two-dimensional spatial information, the obtained local spatial features are downsampled and flattened into one-dimensional data, which is then input to the bidirectional Mamba layer (bidirectional Mamba encoder). The bidirectional Mamba layer includes: two Mamba structures after removing the fully connected layers and one fully connected layer. One Mamba structure is used to input the flattened one-dimensional data to obtain the forward output feature. The other Mamba structure is used to input the data after flipping the flattened one-dimensional data to obtain the reverse output feature. The fully connected layer is used to linearize the feature after fusing the forward and reverse output features to obtain the feature map output by the bidirectional Mamba layer.

[0055] The multi-scale feature fusion module is used to fuse the feature maps output by the deep convolutional layer and the feature maps output by the bidirectional Mamba layer to obtain the fused feature map.

[0056] A classification head is used to generate classification results based on the fused feature maps. In this embodiment of the invention, the classification head is a linear output layer that outputs predicted values ​​in nine directions.

[0057] In this embodiment of the invention, Mamba is a structured state-space model, introduced into sequence modeling tasks in deep learning. In image understanding tasks, Transformer series models typically segment image data into multiple fixed-size patches, inputting the patches as tokens into the model and establishing global dependencies through a self-attention mechanism. When Mamba is used for image understanding, this setup can also be followed, but its inherent causal pattern in unidirectional modeling imposes additional causal constraints on token mixing, which has a detrimental effect on global understanding. To mitigate this disadvantage, this invention designs a bidirectional Mamba to process the input from two directions. Specifically, the recursive process of the existing unidirectional Mamba is set as y′=SSM(x). Assuming the input sequence is input (in this embodiment, the one-dimensional data flattened after the MBConv layer), after being input into the Mamba module, it is mapped through a linear layer to obtain the input x of the main path and the input z of the residual path. The calculation process of the main path is as follows:

[0058] x proj =SSM(SiLU(Conv1d(x))) (1)

[0059] Where SiLU(·) is the SiLU activation function, and Conv1d(·) represents a one-dimensional convolution operation.

[0060] After mixing the residual paths, a positive output feature y is obtained. forward :

[0061] y forward =x proj ⊙SiLU(z) (2)

[0062] Here, ⊙ represents the dot product operation.

[0063] The input is then flipped along the sequence dimension to obtain input_b, which is then mapped through a linear layer to obtain x. b and z b Similarly, the reverse output feature y is obtained. backward :

[0064]

[0065] After fusing the features from both directions and passing them through a linear output layer, the output y of the bidirectional Mamba module can be obtained:

[0066]

[0067] Compared to the original Mamba causal mode, bidirectional Mamba is closer to the Transformer's fully visible mode, which avoids the limitations of causal constraints to some extent and is more suitable for global image understanding.

[0068] In this embodiment of the invention, the Patch layer consists of two 3×3 convolutional layers. The stride of the first convolutional layer is 2, which outputs the original 3-channel image as a 48-channel feature map with half the size. The second convolutional layer does not change the number of channels and size of the feature map, further enhancing the expressive power of the network.

[0069] As a preferred option, such as Figure 3As shown, the feature extraction module in this embodiment of the invention includes four stages connected in series. Stage 0 includes two connected MBConv layers, Stage 1 includes six connected MBConv layers, Stage 2 includes fourteen bidirectional Mamba layers, and Stage 3 includes two bidirectional Mamba layers. The MBConv layers are used to capture two-dimensional local spatial features and include a 1×1 convolutional layer, a 3×3 depthwise convolutional layer that expands the number of channels by a factor of 4, an SE module, and a 1×1 convolutional layer that restores the number of channels to its original value. The SE module performs global average pooling on the spatial dimension of the input feature map, generating weights for each channel through two fully connected layers and a non-linear activation function, which are then applied to each channel of the original feature map to achieve adaptive adjustment of channel information. The first MBConv layer of each stage additionally performs downsampling by changing the stride of the first convolutional layer from 1 to 2.

[0070] As a preferred approach, to maintain good performance when processing data of varying sizes, a bottom-up pyramid structure is used to build the network. The four stages are connected in a pyramid structure, with the feature map scales processed by Stages 0-3 gradually decreasing, and the number of feature channels doubling accordingly. Arranging the four stages in a pyramid structure allows features to flow between depthwise convolutional layers (MBConv layers) of varying widths and bidirectional Mamba modules, acquiring features at different resolutions and semantic levels. This enables the model to handle both small and large objects simultaneously, enhancing its generalization ability and classification performance.

[0071] As a further design of the present invention, the multi-scale feature fusion module in this embodiment is a multi-scale feature fusion module M2SF with a learnable scaling factor, such as... Figure 4 As shown, it specifically includes: an SE layer (Squeeze-and-Excitation Layer), a normalization layer, an average pooling layer, a scaling factor learned through training, and a merging layer.

[0072] The SE layer is used to adaptively adjust the feature maps output by each Stage. In this embodiment, there are four SE layers, each used to adaptively adjust the feature maps output by one of the four Stages. The adjusted feature maps of different scales are then processed by a normalization layer and an average pooling layer. The average pooling layer uses kernels of different sizes to average pool the feature maps of different scales to the same size, resulting in feature maps of the same size. Each Stage is associated with a learnable scaling factor. The scaling factor is multiplied by the corresponding feature map, and the merging layer is used to merge the multiplied feature maps along the channel dimension to obtain the final fused feature map. The fused feature map is then subjected to global average pooling and input into the classification head to generate the classification result.

[0073] In this embodiment of the invention, the multi-scale feature fusion module M2SF with learnable scaling factors utilizes the channel adaptive enhancement of the SE layer to enhance the feature representation capability of each stage. The enhanced features are scaled to adapt to data of different scales. If the target is small, the scaling factor corresponding to the lower-level stage (e.g., Stage 0, Stage 1) is increased to gain more attention. If the target is large, the scaling factor corresponding to the higher-level stage (e.g., Stage 2, Stage 3) is increased.

[0074] The network of the pyramid structure constructed in this embodiment of the invention is shown in Table 1.

[0075] Table 1. Network of the pyramid structure constructed in the embodiments of the present invention.

[0076]

[0077]

[0078] In Table 1, E and D represent the scale and number of channels of the feature map, respectively.

[0079] Preferably, this invention employs transfer learning, using a pre-trained model from ImageNet1K to fine-tune the model on a defect classification dataset. Since the metal surface defect classification dataset is relatively small, to enhance model performance, the model is pre-trained on the ImageNet1K image classification dataset, transferring the learned knowledge to the metal surface defect classification task. Because some general feature representations have already been learned on a large-scale dataset, general low- and mid-level features can be extracted. Therefore, when fine-tuning on the metal surface defect dataset, the model can converge quickly, reducing a significant amount of training time and computational resources. It also effectively avoids problems such as overfitting and difficulty in convergence caused by a small task dataset.

[0080] The ImageNet1K dataset contains 1.28M training images and 50K validation images from 1000 classes. The training settings primarily follow DeiT, specifically applying random cropping, random horizontal flipping, label smoothing regularization, blending, and random erasure as data augmentation. When training on input images with a pixel size of 224×224, this invention sets the momentum to 0.9, the total batch size to 512, and the weight decay to 0.05 to optimize the model. The learning rate uses cosine decay and 3×10⁻⁶. -4The initial learning rate and EMA were used for 300 epochs of training (due to computational resource limitations, the total batch size was reduced from 1024 DeiT to 512, so the initial learning rate was also reduced accordingly). Training was performed on four 4090 GPUs, and the optimal model weights after training were saved to a .pth file.

[0081] Fine-tuning training requires first loading the saved pre-trained model, then modifying the linear output layer of the network, changing the number of output channels from 1000 classes in ImageNet1K to 9 classes of metal surface defects. The metal surface defect dataset is then input for fine-tuning. Because the metal surface defect data differs significantly from the ImageNet1K data, no neural layers are frozen. The mixed metal surface defect dataset undergoes data augmentation operations (random cropping, flipping, and rotating) and normalization, resizing to a uniform size of 224×224. Since it is a grayscale image, the grayscale values ​​of the RGB color channels are the same. The AdamW optimizer is used during training, with a fixed learning rate of 1×10⁻⁶. -4 The batch size was set to 8, the epochs to 25, and fine-tuning training was performed on a single 4090 GPU. The training data input to the model was an 8×224×224×3 matrix, and the output was an 8×9 vector, with each vector corresponding to one defect. The vector with the largest output value was the final classification result.

[0082] The deep neural network in this embodiment of the invention has 24 neural layers with progressively decreasing widths to achieve the above-mentioned nine defect classifications. The Patch layer consists of two 3×3 convolutional layers. The first convolutional layer has a stride of 2, outputting a 48-channel feature map with half the size from the original 3-channel image. The second convolutional layer does not change the number of channels or the size of the feature map, further enhancing the network's expressive power. The tensor shape changes to 8×112×112×48. Stage 0 has two MBConv layers. The first MBConv layer includes a downsampling max-pooling layer, resulting in a tensor shape of 8×56×56×48 after passing through the Stage 0 network layers. Stage 1 has six MBConv layers. The first MBConv layer also includes a downsampling layer, resulting in a tensor shape of 8×28×28×96 after passing through the Stage 1 network layers. Stage 2 has 14 bidirectional Mamba layers. An additional 3×3 convolutional layer is added before the first bidirectional Mamba layer. The convolutional downsampling layer downsamples the tensor to 8×14×14×192. To adapt to the one-dimensional data input format of the Mamba layer, the tensor is flattened from two dimensions to one dimension (8×196×192) and then normalized. After 14 bidirectional Mamba layers, the shape remains unchanged. Stage 3 has two bidirectional Mamba layers that restore the one-dimensional data shape to a two-dimensional format. After downsampling, the tensor shape is 8×7×7×384, which is then flattened to 8×49×384. After two bidirectional Mamba layers, the shape remains unchanged. The outputs of Stage 2 and Stage 3 are transformed into two-dimensional feature maps, which are then processed simultaneously with the feature maps output from Stage 0 and Stage 1 using the SE layer and normalized using BatchNorm. Average pooling is then performed to obtain four feature maps of the same size but different numbers of channels. Each feature map is multiplied by its corresponding scaling factor, and they are fused along the channel dimension to obtain a tensor shape of 8×7×7×720, which flattens to 8×49×720. After layer normalization, average pooling is performed in the second dimension to obtain an 8×720 vector. This vector is then passed through an output linear layer to output an 8×9 vector.

[0083] The performance of the classification model is evaluated using a validation set. The validation set data is only resized and normalized using the mean and standard deviation of the training set to ensure a stable data distribution during validation, consistent with real-world usage scenarios. To visually demonstrate the classification results, this invention uses accuracy, precision, recall, and F1 score as evaluation criteria, expressed by the following formulas:

[0084]

[0085]

[0086] TP, FP, FN, and TN represent the number of true positives, false positives, false negatives, and true negatives, respectively.

[0087] The various indicators of the metal surface defect classification model constructed in this embodiment of the invention are shown in Table 2, where support represents the number of defect samples of this type participating in the test.

[0088] Table 2. Indicators of the classification model for 9 types of metal surface defects

[0089]

[0090] It can be observed that the models provided in this embodiment of the invention perform well in classifying the nine defect categories, with an overall accuracy of 0.983228. Comparing the performance of the proposed model with ResNet18 and ViT, the model of this invention also achieves the highest classification accuracy. Since Rp-class defects are the most difficult to distinguish among all models, the three classification metrics for Rp are also compared separately, as shown in Table 3.

[0091] Table 3 Comparison of classification performance of various models

[0092]

[0093]

[0094] Analysis shows that the model of this invention achieved good results in all three metrics. When the M2SF module was removed, the recall metric was the best. Although the model accuracy was not as good as ResNet18 and ViT, the computational efficiency of the model was higher. With the M2SF module added, the model of this invention achieved the best results, proving that the M2SF proposed in this invention improves the model's generalization and classification capabilities.

[0095] In addition, Figures 5-8 Confusion matrices were plotted for four models: ResNet18, ViT-B, the model of this invention without the M2SF module, and the model of this invention with the M2SF module. Each column of the confusion matrix represents the predicted class, and the total number in each column represents the number of data points predicted to belong to that class. Each row represents the true class of the data, and the total number of data points in each row represents the number of data instances in that class. The value in each column represents the number of actual data points predicted to belong to that class. Darker colors on the main diagonal indicate better classification performance. It can be seen that the model of this invention achieves better classification results.

[0096] This invention focuses on defect classification tasks and proposes a novel network architecture to improve classification performance and provide reliable support for the complete defect detection process. Addressing the limitations of CNNs and Transformers, this invention combines CNNs with Mamba, a structured state-space model with linear computational complexity and long sequence modeling capabilities, to construct a model structure that integrates deep convolution and bidirectional Mamba modules. Deep convolution captures local spatial features of image data in parallel along the spatial dimension, while Mamba's long sequence modeling capabilities enhance the model's upper limit and capacity. The bidirectional Mamba design aims to reduce the adverse effects of unidirectional causal dependencies caused by Mamba's RNN-like structure on global understanding. This model structure simultaneously focuses on local and global features, significantly improving the accuracy of the classification model while maintaining computational efficiency. During network construction, the network is divided into four parts, each with different feature scales that decrease progressively. M2SF is used to fuse these features along the channel dimension, fully utilizing the feature information at each scale and assigning adaptive weights to these features. This mitigates the impact of insufficient defect samples to some extent and enables the network to be robustly trained on datasets of varying sizes, enhancing the model's generalization ability.

[0097] Example 2

[0098] This invention provides a method for classifying metal surface defects, including:

[0099] The images of metal surface defects to be classified are input into the metal surface defect classification model to obtain the defect classification results. The metal surface defect classification model is constructed using the metal surface defect classification model construction method based on the structured state-space model described in Example 1.

[0100] The relevant technical solutions are described in Example 1 and will not be repeated here.

[0101] Example 3

[0102] This invention provides an electronic device, including a memory and a processor. The memory stores a computer program, and when the processor executes the computer program, it implements the steps of the metal surface defect classification model construction method based on the structured state space model in Embodiment 1 or the metal surface defect classification method in Embodiment 2.

[0103] The electronic device can be a desktop computer, laptop, handheld computer, or cloud server, etc. The processor can be a Central Processing Unit (CPU), or other general-purpose processors, digital signal processors (DSPs), application-specific integrated circuits (ASICs), field-programmable gate arrays (FPGAs), or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, etc. The memory can be used to store computer programs and / or modules. The processor performs various functions of the electronic device by running or executing the computer programs and / or modules stored in the memory, and by accessing data stored in the memory.

[0104] The relevant technical solutions are the same as above, and will not be repeated here.

[0105] Example 3

[0106] This invention provides a computer-readable storage medium storing a computer program thereon. When the computer program is executed by a processor, it implements the steps of the metal surface defect classification model construction method based on the structured state space model in Embodiment 1 or the metal surface defect classification method in Embodiment 2.

[0107] Specifically, the memory may include high-speed random access memory, as well as non-volatile memory, such as hard disks, RAM, plug-in hard disks, smart media cards (SMC), secure digital cards (SD), flash cards, at least one disk storage device, flash memory device, or other volatile solid-state storage devices.

[0108] The relevant technical solutions are the same as above, and will not be repeated here.

[0109] Example 4

[0110] This application provides a computer program product, including a computer program that, when run on a computer, causes the computer to execute the steps of the metal surface defect classification model construction method based on the structured state space model in Embodiment 1 or the metal surface defect classification method in Embodiment 2.

[0111] The relevant technical solutions are the same as above, and will not be repeated here.

[0112] Those skilled in the art will readily understand that the above description is merely a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, and improvements made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.

Claims

1. A method for constructing a metal surface defect classification model based on a structured state-space model, characterized in that, include: A neural network based on a structured state-space model is pre-trained using an image dataset, and then fine-tuned using a metal surface defect dataset to obtain a metal surface defect classification model; wherein, the neural network based on the structured state-space model includes: The Patch layer is used to downsample the input image samples and increase the number of channels; A multi-scale feature extraction module includes deep convolutional layers and bidirectional Mamba layers at different scales. The deep convolutional layers are used to capture local spatial feature maps of different scales in the output image of the Patch layer in parallel along the spatial dimension. The bidirectional Mamba layer includes two Mamba structures with fully connected layers removed and one fully connected layer. One Mamba structure is used to obtain forward output features based on one-dimensional data after downsampling and flattening the local spatial feature maps, and the other Mamba structure is used to obtain inverse output features based on data after flipping the one-dimensional data. The fully connected layer is used to linearize the features obtained by fusing the forward and inverse output features to obtain the feature map output by the bidirectional Mamba layer. The multi-scale feature fusion module is used to fuse the feature map output by the deep convolutional layer and the feature map output by the bidirectional Mamba layer at multiple scales to obtain the fused feature map. A classification head is used to generate classification results based on the fused feature map; The deep convolutional layer is an MBConv layer, and the multi-scale feature extraction module includes four stages connected in a pyramid structure: Stage 0-Stage 3. The feature map scale output by Stage 0-Stage 3 gradually decreases, and the number of feature channels is doubled accordingly. Stage 0 includes two sequentially connected MBConv layers, Stage 1 includes six sequentially connected MBConv layers, Stage 2 includes fourteen sequentially connected bidirectional Mamba layers, and Stage 3 includes two sequentially connected bidirectional Mamba layers. The multi-scale feature fusion module is used to fuse the feature maps output from Stage0 to Stage3 to obtain the fused feature map.

2. The method for constructing a metal surface defect classification model according to claim 1, characterized in that, The positive output feature for: The reverse output feature for: in, This represents the output of the main path of the Mamba structure after removing the fully connected layer. for Activation function This represents a one-dimensional convolution operation. This represents the dot product operation; , These represent the input of the main path and the input of the residual path in the Mamba structure obtained by mapping the one-dimensional data through a linear layer, respectively. , The inputs of the main path and the residual path in the Mamba structure are obtained by mapping the one-dimensional data after it has been flipped through a linear layer.

3. The method for constructing a metal surface defect classification model according to claim 2, characterized in that, The multi-scale feature fusion module is a multi-scale feature fusion module M2SF with a learnable scaling factor, specifically including: an SE layer, a normalization layer, an average pooling layer, a scaling factor learned through training, and a merging layer. The SE layer is used to adaptively adjust the feature map output by each Stage. The adjusted feature map is then processed by the corresponding normalization layer and average pooling layer to obtain feature maps of the same size. The feature map is multiplied by the corresponding scaling factor and then merged in the channel dimension by the merging layer to obtain the fused feature map.

4. The method for constructing a metal surface defect classification model according to claim 1 or 3, characterized in that, The metal surface defect dataset is a dataset obtained by merging multiple metal surface defect datasets of different sizes.

5. The method for constructing a metal surface defect classification model according to claim 4, characterized in that, Based on a metal surface defect dataset, a metal surface defect classification model was obtained by fine-tuning the pre-trained neural network using transfer learning.

6. A method for classifying metal surface defects, characterized in that, include: The image of the metal surface defect to be classified is input into the metal surface defect classification model to obtain the defect classification result; wherein, the metal surface defect classification model is constructed by the metal surface defect classification model construction method according to any one of claims 1-5.

7. An electronic device, characterized in that, Includes computer-readable storage media and processors; The computer-readable storage medium is used to store executable instructions; The processor is used to read executable instructions stored in the computer-readable storage medium to execute the metal surface defect classification model construction method according to any one of claims 1-5, or to execute the metal surface defect classification method according to claim 6.

8. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the program is executed by the processor, it implements the metal surface defect classification model construction method as described in any one of claims 1-5, or the metal surface defect classification method as described in claim 6.

9. A computer program product, characterized in that, The method includes a computer program that, when run on a computer, causes the computer to execute the metal surface defect classification model construction method according to any one of claims 1-5, or to execute the metal surface defect classification method according to claim 6.

Images