Lightweight steel surface defect segmentation model based on efficient convolution and attention mechanism

Abstract

The invention discloses a lightweight steel surface defect segmentation model based on efficient convolution and an attention mechanism, and belongs to the field of machine vision. The method mainly comprises the following modules: a lightweight feature extraction network, a lightweight feature enhancement module, a hierarchical adaptive upsampling fusion module and a lightweight segmentation head module. According to the model, in a feature extraction stage, a lightweight feature extraction module is adopted to efficiently extract space detail information and semantic information, then in a decoder stage, features of a specific scale are further enhanced through a lightweight local feature enhancement module, and a feature graph with the size being 1 / 4 of that of an original graph is selected to perform feature enhancement in the aspect of space details; and then the enhanced features are fused through a hierarchical adaptive upsampling fusion module to recover the original features. According to the method, the defect image can be effectively segmented, and good balance is achieved before real-time performance and accuracy. And a new solution is provided for steel surface defect segmentation in an actual scene.

Description

Technical Field

[0001] The technical field to which the present invention belongs is the field of computer vision, specifically the direction of industrial defect detection. It is a lightweight steel surface defect segmentation model based on efficient convolution and attention mechanism. Background Art

[0002] In the steel production process, surface defect detection is a critical step in ensuring product quality. Accurately detecting surface defects in steel in real time can effectively prevent inferior products from entering the market, reducing production costs and improving market competitiveness.

[0003] In the field of surface defect detection, early manual inspection methods consumed significant time and human resources. In recent years, computer vision-based defect detection methods have gained widespread application due to their outstanding performance, providing a reliable solution for ensuring product quality. In particular, deep learning-based methods have achieved remarkable success in industrial surface defect detection. Among these, methods based on convolutional neural networks have been widely studied due to their outstanding performance in image processing tasks. These methods demonstrate high accuracy and robustness in complex and dynamic industrial environments. One commonly used approach is surface defect classification, while another widely used method is area-based defect detection. Neither defect detection technique provides detailed location and shape information for defects, but they can struggle with defects exhibiting complex textures or shape variations. Pixel-level detection methods focus on accurately classifying each pixel in an image, aiming to more precisely capture the edges and details of defects. Through this fine-grained recognition, pixel-level methods can achieve comprehensive and accurate detection of surface defects, providing a more efficient means of quality control in industrial production.

[0004] With the advancement of deep learning, numerous network models have been proposed. However, due to their high number of parameters and computational overhead, they are challenging to meet the requirements of practical applications that require rapid inference results. Although current lightweight detection networks improve inference speed by simplifying their design to reduce model parameters and computational load, the resulting inference acceleration often fails to meet the demanding real-time requirements of certain practical applications. Furthermore, their performance in complex defect scenarios is generally poor. Furthermore, current lightweight semantic segmentation networks often face a trade-off between the performance of small-scale features and segmentation capabilities, as they cannot effectively process both low-level details and high-level semantics simultaneously.

[0005] For the task of steel surface defect segmentation, it is necessary to achieve the goals of lightweight, real-time, and high segmentation accuracy. The current general lightweight surface defect segmentation model cannot achieve ideal application effects.

[0006] In light of the above issues, this paper proposes a lightweight steel surface defect segmentation model based on an efficient convolution and attention mechanism. This model adopts an asymmetric encoder-decoder structure. In the encoder structure, depthwise separable convolution and ghost convolution are used as the basis for constructing the encoder's basic residual module. In the decoder structure, skip connections are reduced, and a lightweight attention mechanism is fully utilized to enhance features. Summary of the Invention

[0007] Technical problem: In response to the shortcomings of existing surface defect segmentation technology, the present invention provides a lightweight steel surface defect segmentation model based on efficient convolution and attention mechanism to solve the problems existing in the above background technology.

[0008] Technical Solution: First, the defect images and labels used for training in the dataset are preprocessed into PyTorch Tensors and resized to a specific size. The preprocessed images are then fed into an encoder to extract local and global features layer by layer, generating multi-scale feature maps. The shallow spatial feature maps and deep semantic feature maps are then enhanced using a feature enhancement module, and the features are further fused through step-by-step decoding and skip connections. Finally, the fused feature maps are upsampled to the original image size to produce the segmentation result. The segmentation result is then compared with the actual segmentation map and the loss function is calculated, continuously updating the model parameters. After multiple rounds of training, the trained segmentation model is capable of accurately segmenting defects. The model utilizes a self-designed lightweight module in the encoder network design. In the decoder network design, the number of model parameters is significantly reduced by eliminating skip connections. The result is a lightweight semantic segmentation model with an asymmetric structure.

[0009] The present invention proposes a lightweight steel surface defect segmentation model based on efficient convolution and attention mechanism, which mainly includes a lightweight feature extraction network, a lightweight feature enhancement module integrated with attention mechanism, a hierarchical adaptive upsampling fusion module, and a segmentation head module.

[0010] The lightweight feature extraction network is used to efficiently extract spatial detail features and high-level semantic features of images.

[0011] The feature enhancement module enhances features by further processing of convolution and attention mechanism, including a spatial detail enhancement module and a semantic information enhancement module.

[0012] The hierarchical adaptive fusion upsampling module is used to adaptively fuse multi-scale features in the decoding stage.

[0013] The lightweight segmentation head converts the input feature map into a segmentation result with the same resolution as the input image through upsampling and convolution operations.

[0014] Furthermore, the lightweight feature extraction network is constructed with 5 feature extraction blocks, which can be regarded as 5 stages of feature extraction, respectively extracting feature maps of different scales. Each feature extraction block is composed of three structures of feature extraction bottleneck blocks and double Ghost modules (Double Ghost). The three structures of feature extraction bottleneck block structures are ordinary feature extraction bottleneck block (GDFB), feature extraction bottleneck block with downsampling (Down_GDFB), and wide field feature extraction bottleneck block (Wide_GDFB). The first feature extraction block and the second feature extraction block are respectively composed of a Down_GDFB and a GDFB. Down_GDFB downsamples the feature map and performs feature extraction, and GDFB performs further feature extraction on the feature map. Feature maps of 1 / 2 and 1 / 4 original image scales are obtained respectively. The third and fourth feature extraction blocks are composed of a Down_GDFB, a GDFB, and a Wide_GDFB. They downsample the input feature map first and then perform further feature extraction to obtain feature maps with a scale of 1 / 8 and 1 / 16 of the original image, respectively. The fifth feature extraction block is composed of a Down_GDFB, a GDFB, and a Double Ghost module. It first downsamples the input feature map and then uses GDFB and a lightweight Double Ghost module to perform further feature extraction to obtain a feature map with a scale of 1 / 32 of the original image.

[0015] Furthermore, the basic modules for constructing each feature extraction block in the lightweight feature extraction network: feature extraction bottleneck block (GDFB), feature extraction bottleneck block with downsampling (Down_GDFB), wide field feature extraction bottleneck block (Wide_GDFB) module, and double Ghost module (Double Ghost) are specifically implemented as follows:

[0016] The feature extraction bottleneck block consists of a 1x1 ordinary convolution layer, a normalization layer, a 3x3 depth convolution layer, a Ghost convolution module, and two ReLu activation functions. In particular, it adopts a residual structure and connects the input and output with residuals to alleviate the impact of gradient disappearance in deep networks.

[0017] The feature extraction bottleneck block with downsampling adds a maximum pooling layer for downsampling on the basis of the feature extraction bottleneck block. Overall, this module consists of a maximum pooling layer, a 1x1 ordinary convolution layer, a normalization layer, a 3x3 depth convolution layer, a Ghost convolution module, and two ReLu activation functions. In particular, it adopts a residual structure and performs residual connection between the input and output.

[0018] The wide-field feature extraction bottleneck block consists of a 1x1 ordinary convolution layer, a normalization layer, a 5x5 depth convolution layer, a Ghost Module, and two ReLu activation functions. In particular, it adopts a residual structure, which connects the input and output with residual connections. The 5x5 convolution kernel is used to quickly obtain a larger receptive field.

[0019] The Dual Ghost Module consists of two Ghost Modules, a 3x3 deep convolutional layer, and a ReLU activation function. It aims to generate features in a simple way at the deep network layer, reducing the number of parameters.

[0020] Furthermore, the lightweight feature extraction network building block Ghost Module is used to generate features in a simple way in a deep network. The module consists of two branches. The first branch includes a 1x1 convolution layer, a normalization layer, and a ReLU activation function layer for generating real features. The other branch includes a 1x1 convolution layer, a normalization layer, and a ReLU activation function layer for generating cheap features. Finally, the features of the two branches are merged as the overall feature output.

[0021] Furthermore, the spatial detail enhancement module includes two Sobel filters (sobel_x and sobel_y) to calculate detail information in the x and y directions, respectively. sobel_x and sobel_y are predefined 3x3 convolution kernels, which are used to detect details in the horizontal and vertical directions, respectively. It is used to normalize the extracted edge features and generate an edge attention map. It includes a Sigmoid activation function to compress the input value into the range of [0,1], which is used to adjust the weight of the feature map. It includes a 1x1 convolution layer to reduce the number of channels of the enhanced feature map from the number of input channels to the number of output channels.

[0022] Furthermore, the semantic feature enhancement module is composed of a 3x3 depth convolution layer, an average pooling layer, a 1x1 convolution layer, and a Sigmoid activation function layer.

[0023] Furthermore, the hierarchical adaptive fusion upsampling module receives two-scale feature maps, first upsampling the smaller-scale feature map, then fusing the upsampled feature map with the other feature map to generate a new feature map. Channel attention weights are calculated on this new feature map, and finally the attention weights are multiplied by the new feature map for output. This module consists of a bilinear interpolation upsampling layer, a normalization layer, two ReLU activation function layers, an average pooling layer, two convolutional layers, and a sigmoid activation function layer.

[0024] Furthermore, the segmentation head module is used to map the fused features to the original image size through upsampling. It consists of a 1x1 convolution layer, a 3x3 depth convolution layer, and a 3x3 normal convolution layer.

[0025] After processing the above modules in sequence, the real-time segmentation task of surface defects can be completed.

[0026] Compared with the prior art, the present invention has the following beneficial effects:

[0027] 1. The present invention focuses on the lightweight design of the defect segmentation model, aiming to solve the limitations of traditional defect segmentation in terms of computing resource consumption, model parameter scale, etc., so as to better adapt to resource-constrained application scenarios such as mobile terminals and embedded devices. A lightweight steel surface defect segmentation model based on efficient convolution and attention mechanism is proposed. Studies have shown that the network achieves a good balance between inference speed and accuracy, and maintains the lightweight of the model. Compared with the mature semantic segmentation model U-Net (parameters 31.39M, FLOPS 42.75G), the proposed network model has only 0.57M parameters, a reduction of 98.18%, and only uses 0.30G floating-point operations (FLOPS), reducing the computational load by 98.67%.

[0028] 2. The present invention designs a lightweight feature extraction network (GD-Net), which is constructed based on depthwise separable convolution and Ghost convolution. Depthwise separable convolution decomposes the traditional standard convolution operation into two independent steps: depthwise convolution and pointwise convolution. While effectively extracting features, it significantly reduces the amount of calculation and the number of parameters. Ghost convolution is based on the principle of feature redundancy, using partial convolution operations to generate the original feature map, and then obtaining a "ghost" feature map with similar information through a simple linear transformation, further reducing the computational overhead and parameter scale of the convolution operation. By rationally organizing these two efficient convolution methods, the lightweight feature extraction network constructed can achieve efficient and lightweight modeling while ensuring feature extraction capabilities.

[0029] 3. This paper designs a feature enhancement module that combines a lightweight attention mechanism. This module abandons the complex and computationally intensive operations of the traditional attention mechanism. By introducing a lightweight design concept, it significantly reduces computing resource consumption while retaining the attention mechanism's advantage of focusing on key features and suppressing irrelevant information. This module can automatically learn the importance weights of different positions and channels in the feature map, strengthen important features, and suppress irrelevant or minor features, thereby achieving effective feature enhancement and providing more discriminative feature representations for subsequent model processing.

[0030] 4. The present invention designs a hierarchical adaptive fusion upsampling module, which focuses on information integration and upsampling operations in the feature fusion process. In traditional feature fusion methods, fixed fusion strategies are often adopted, which are difficult to adapt to the differences and complexities between features at different levels. However, this module can dynamically adjust the fusion weights according to the characteristics of features at different levels through a hierarchical adaptive mechanism, thereby achieving seamless fusion and upsampling between features at different levels. This adaptive fusion method can not only make full use of the information of features at each level, but also effectively avoid the loss and redundancy of feature information, thereby achieving efficient feature enhancement and fusion. BRIEF DESCRIPTION OF THE DRAWINGS

[0031] Figure 1 Schematic diagram of the overall network structure of the model of the present invention

[0032] Figure 2 Schematic diagram of the lightweight feature extraction network structure of the model of the present invention

[0033] Figure 3 Schematic diagram of the lightweight extraction network basic block structure of the model of the present invention

[0034] Figure 4 Schematic diagram of the structure of the detail feature enhancement module and semantic feature enhancement module of the model of the present invention

[0035] Figure 5 Schematic diagram of the hierarchical adaptive fusion upsampling module structure of the model of the present invention

[0036] Among them, Figure 1 , attached Figure 2 The winning bid numbers are as follows:

[0037] 000-model input, 100-lightweight feature extraction network, 110-shallow feature extraction block, 120-deep feature extraction block, 130 deepest semantic feature extraction block, 200-spatial detail feature enhancement module, 210-semantic feature enhancement module, 300-level adaptive fusion upsampling module, 400-segmentation head module, 500-output feature map. DETAILED DESCRIPTION

[0038] The overall structure of the model of the invention is as follows Figure 1 As shown, completing the training and testing of the model requires model construction, input data processing, training logic construction, and model testing.

[0039] 1. Model Construction

[0040] The invention proposes the construction of a lightweight steel surface defect segmentation model based on efficient convolution and attention mechanism, which includes the construction of a lightweight feature extraction network and a decoder network.

[0041] 1. Construction of lightweight feature extraction network (GD-Net)

[0042] The depth-wise separable convolution basic blocks _DWConv3x1_1x3, _DWConv5x1_1x5 encapsulate convolution, batch normalization, and activation functions in one module, which simplifies the network construction process and improves code readability and maintainability.

[0043] (1) Depthwise Separable Convolution Basic Block_DWConv3x1_1x3 Construction

[0044] _DWConv3x1_1x3 is a depthwise separable convolution module designed to reduce computational complexity and parameter requirements while maintaining good feature extraction capabilities. This module implements depthwise separable convolution, breaking down the standard convolution into two steps: depthwise convolution and pointwise convolution. It also adds optional batch normalization and activation functions. Batch normalization and ReLU activation functions can be optionally applied after the convolution operation.

[0045] Depthwise convolution is implemented using nn.Conv2d with a kernel size of kernel_size=3. groups=in_channels indicates that each input channel is convolved independently. Pointwise convolution is implemented using nn.Conv2d with a kernel size of 1.

[0046] In the forward propagation method, the input feature map x is accepted, forward propagated through the defined convolution sequence, and the convolution result is returned.

[0047] This module significantly reduces the amount of computation and parameters through depth-wise separable convolution while maintaining good feature extraction capabilities.

[0048] (2) Depthwise Separable Convolution Basic Block_DWConv5x1_1x5 Construction

[0049] _DWConv5x1_1x5 is a depthwise separable convolution module designed to expand the receptive field of deep networks. The module mainly consists of depthwise separable convolution, pointwise convolution, batch normalization and activation function.

[0050] Depthwise separable convolutions are constructed using two convolution kernels with different orientations (5×1 and 1×5) to implement depthwise convolution. This effectively expands the receptive field while maintaining computational efficiency and capturing richer contextual information. Pointwise convolutions use 1×1 convolution kernels to perform channel mixing on the output of the depthwise convolution and adjust the number of channels. Batch normalization normalizes the output to reduce internal covariate shift and improve training stability. Finally, the ReLU6 activation function is used to prevent the vanishing gradient effect.

[0051] (3) Channel adjustment convolution basic fast ConvBNReLU_1x1 construction

[0052] ConvBNReLU_1x1 is a convolutional module for adjusting the number of channels. It consists of 1×1 convolution, batch normalization, and an optional ReLU6 activation function, providing an efficient and flexible channel adjustment operation. It encapsulates convolution, batch normalization, and activation functions in a single module, simplifying the network construction process and improving code readability and maintainability.

[0053] (4) Construction of dual GhostModule modules with residual structure

[0054] GhostModule is a module derived from GhostNet that aims to improve computational efficiency by reducing redundant features. It primarily performs three operations: primary feature extraction, ghost feature generation, and feature combination. Primary feature extraction uses a standard convolutional layer to extract the main features. Ghost feature generation uses a cheap operation to generate additional features. These features are obtained by performing deep convolution operations on the primary features to simulate more feature representations. Finally, a cheap operation is used to generate additional features. The primary and ghost features are concatenated along the channel dimension to form the final output features. To minimize computational resource consumption while maintaining good feature extraction capabilities, GhostModule is used here to construct a dual GhostModule module with a residual structure.

[0055] The dual GhostModule module with residual structure uses two GhostModules to achieve feature expansion and recovery. The first GhostModule in the dual GhostModule is used to expand the number of channels of the input feature. By generating more feature representations, the expressive power of the feature is enhanced. The second GhostModule is used to restore the feature channel to its original size. By reducing redundant features, computational efficiency is maintained. The residual connection is then used to add the input features to the features processed by the two GhostModules. This design is to alleviate the gradient vanishing problem in deep networks and improve the training efficiency of the network. Finally, ReLU6 is used as the activation function. The module structure is as follows Figure 3 (e) shown.

[0056] The combination of two GhostModules enables this module to maintain feature representation capabilities while significantly reducing computational resource consumption. The introduction of a residual structure enables more efficient training of deep models, reducing the risk of vanishing gradients. This minimizes computational resource consumption, making it a fast foundation for building the feature extraction bottleneck block, GDFB.

[0057] (5) Feature extraction basic fast GDFB construction

[0058] The Feature Extraction Bottleneck Block (GDFB) is the core module for building a lightweight feature extraction network. By combining the GhostModule and depthwise separable convolution, the module provides an efficient feature extraction method. The design concept of this module is to minimize the consumption of computing resources while maintaining good feature extraction capabilities.

[0059] The feature extraction bottleneck block (GDFB) first uses the ConvBNReLU_1x1 module to adjust the channel to reduce the amount of calculation. Then it uses _DWConv3x1_1x3 to perform standard feature extraction, and then uses the GhostModule to restore the number of channels. By generating a small number of main features and additional Ghost features, the computational efficiency is maintained. Finally, the residual connection is used to add the input features to the features after convolution processing, and activated by the ReLU6 activation function. The module structure is as follows Figure 3 As shown in (a).

[0060] The feature extraction bottleneck block with downsampling (Down_GDFB) adds the MaxPool layer on the basis of GDFB to achieve the downsampling function. The module structure is as follows Figure 3 (b) shown.

[0061] The wide receptive field feature extraction bottleneck block (Wide_GDFB) replaces the _DWConv3x1_1x3 module in GDFB with the _DWConv5x1_1x5 module, expanding the receptive field without significantly increasing the amount of computation and capturing richer contextual information. The module structure is as follows Figure 3 (c) shown.

[0062] (6) Construction of lightweight feature extraction network (GD-Net)

[0063] Lightweight feature extraction network such as Figure 2 As shown in the figure, it is constructed from five feature extraction blocks, which extract feature maps of different scales. The first and second feature extraction blocks are composed of Down_GDFB and GDFB. Down_GDFB downsamples the feature map and performs feature extraction, while GDFB performs further feature extraction on the feature map. Feature maps with a scale of 1 / 2 and 1 / 4 of the original image are obtained, respectively. The third and fourth feature extraction blocks are composed of Down_GDFB, GDFB, and Wide_GDFB. They downsample the input feature map before further feature extraction, obtaining feature maps with a scale of 1 / 8 and 1 / 16 of the original image, respectively. The fifth feature extraction block is composed of Down_GDFB, GDFB, and Double Ghost modules. Down_GDFB is first used for downsampling, and then GDFB and the lightweight Double Ghost module are used for further feature extraction to obtain a feature map with a scale of 1 / 32 of the original image.

[0064] 2. Decoder network construction

[0065] (1) Construction of semantic feature enhancement module (FEB)

[0066] The semantic feature enhancement module combines convolution operation and attention mechanism.

[0067] _DWConv3x1_1x3 performs convolution operations to extract input features. The use of depth-separable convolution is to reduce the amount of calculation while maintaining good feature extraction capabilities. Then, the global information of the feature map is extracted through global average pooling (AdaptiveAvgPool2d), and the global information is processed using 1×1 convolution to generate attention weights. The Hardsigmoid activation function is then used to map the attention weights to the [0,1] range. Finally, the attention weights are multiplied by the feature map to enhance important features and suppress unimportant features. The module structure is as follows Figure 4 As shown in (a).

[0068] Through the attention mechanism, the module automatically selects and enhances important features. This selective enhancement helps improve model performance. Using depthwise separable convolutions and a simple attention mechanism, the module reduces computational resource consumption while maintaining good feature representation capabilities.

[0069] (2) Construction of Detail Feature Enhancement Module (FSEB)

[0070] This module is a module for enhancing spatial detail features. It uses the Sobel filter to extract edge information to improve the model's sensitivity to edges and details. First, the Sobel filter is used to extract the edge information of the input feature map. The Sobel filter is a commonly used edge detection operator that can effectively capture edge features in images. By calculating the attention weight of the edge feature, the edge information in the input feature map is enhanced. Then, the Sigmoid activation function is used to map the edge feature to the [0,1] range to generate the attention weight. The edge attention weight is then multiplied by the input feature map, and the edge information in the feature map is enhanced by weighted sum. Finally, the 1×1 convolution adjusts the number of channels of the enhanced feature map to the number of channels of the output. The module structure is as follows Figure 4 (b) shown.

[0071] The spatial feature enhancement module provides an efficient edge feature enhancement method. This module is built to minimize the consumption of computing resources while maintaining good feature extraction capabilities.

[0072] (3) Construction of hierarchical adaptive fusion upsampling module (UFB)

[0073] The Hierarchical Adaptive Fusion Upsampling module is a feature fusion module that combines upsampling, attention mechanism, and feature fusion to provide an efficient feature fusion method. Its design philosophy is to minimize computing resource consumption while maintaining good feature extraction capabilities.

[0074] First, use bilinear interpolation (F.interpolate) to upsample the high-level features so that they match the spatial size of the low-level features, so as to maintain spatial alignment during feature fusion. Secondly, check the number of channels. If the number of input and output channels is inconsistent, use ConvBNReLU_1x1 to adjust the channels. Furthermore, the upsampled high-level features are added to the low-level features for feature fusion, so that the semantic information of the high-level features and the detail information of the low-level features are fully combined. Furthermore, the attention mechanism is used to weight the fused features. This operation is to extract global information through global average pooling (AdaptiveAvgPool2d), generate attention weights through two 1×1 convolutions and ReLU activation functions, and finally map the weights to the [0,1] range through the Sigmoid activation function. Finally, the attention weights are multiplied by the fused features to enhance important features and suppress unimportant features. Its structure is as follows Figure 5 shown.

[0075] (4) Construction of segmentation head module

[0076] The segmentation head module is responsible for converting feature maps into the final segmentation results. By combining channel adjustment, feature extraction, and upsampling, it provides an efficient image segmentation method. Its design philosophy is to minimize computational resource consumption while maintaining good feature extraction capabilities.

[0077] During construction, first, ConvBNReLU_1x1 is used to perform channel adjustment to halve the number of channels of the input feature map. This operation helps reduce the amount of computation while maintaining the expressiveness of the features. Secondly, _DWConv3x1_1x3 is used for depthwise separable convolution to further extract and enhance features. Depthwise separable convolution can maintain good feature extraction capabilities while reducing the amount of computation. Further, a 3×3 convolution layer is used to convert the feature map into the final segmentation result. The number of output channels is equal to the number of target categories n_classes. Further, bilinear interpolation (F.interpolate) is used to upsample the feature map to the specified size (usually the size of the original image). This operation ensures that the output segmentation result is consistent with the size of the input image.

[0078] (5) Construction of decoding path

[0079] After processing the input feature map through the lightweight feature extraction network, multi-scale feature maps are generated: 1 / 2, 1 / 4, 1 / 8, 1 / 16, and 1 / 32 of the original image size. The feature map at 1 / 4 of the original image size is processed by the detail feature enhancement module (FEEB) to obtain feat2_sp. The feature maps at 1 / 8, 1 / 16, and 1 / 32 of the original image size are processed by the adaptive feature enhancement module to obtain feat3, feat4, and feat5. feat5 and feat4 are input to the hierarchical adaptive fusion upsampling module to obtain fuse_feture4. fuse_feture4 and feture3 are input to the hierarchical adaptive fusion upsampling module to obtain fuse_feture3. fuse_feture3 and fe2_sp are input to the hierarchical adaptive fusion upsampling module to obtain fuse_feture2. fuse_feture2 is then processed by seg_head to obtain the final feature map.

[0080] 2. Input Data Loader Construction

[0081] The input data for training includes the original defect image and the defect segmentation image. The data loader needs to process the input image data and defect segmentation image so that they can be input into the segmentation network for the next step of training.

[0082] 3. Steel Surface Defect Segmentation Model Training Steps

[0083] The complete process of training a steel surface defect segmentation model, including data preparation, model training, verification, and saving checkpoints.

[0084] 1. Dataset creation and division:

[0085] Use BasicDataset to create a dataset and split it into a training set and a validation set. There are also parameters that determine the proportion of the validation set.

[0086] 2. Data loader construction:

[0087] Create training and validation data loaders (DataLoader) for batch loading of data. The batch_size parameter determines the number of samples in each batch.

[0088] 3. Training data recording:

[0089] Use wandb to record experiments, training parameters and process indicators.

[0090] 4. Optimizer and loss function:

[0091] Use RMSprop as the optimizer and CrossEntropyLoss as the loss function, depending on the number of output categories of the model.

[0092] Use the ReduceLROnPlateau scheduler to dynamically adjust the learning rate based on the validation score.

[0093] Mixed precision training using GradScaler.

[0094] Use BCE loss function and DICE loss function.

[0095] 5. Iterative training:

[0096] Iterate over the training set, calculate the loss and perform backpropagation and optimization.

[0097] Use autocast for mixed precision training to reduce video memory usage.

[0098] Record the loss of each epoch and record it in wandb.

[0099] 6. Verify and save checkpoint:

[0100] Verification is performed at the end of each epoch, the Dice score of the verification set is calculated, and the learning rate is adjusted according to the verification score.

[0101] 7. Save the checkpoint of the model:

[0102] If save_checkpoint is True, the current model state is saved. mIOU is calculated every 5 epochs and the best model is saved.

[0103] 4. Model Validation and Use

[0104] The best trained model is tested under the current model framework. The input is an image that has not been trained, and the output is a complete defect segmentation map. Different colors represent different defect classifications.

[0105] The model building process, model training details, and usage steps listed above are merely specific illustrations of feasible implementation methods of the present invention. Their purpose is to enable those familiar with the technology to implement them based on the present invention. They are not intended to limit the scope of protection of the present invention. Any equivalent or modification that departs from the technology of the present invention shall be included in the scope of protection of the present invention.

Claims

1. A lightweight steel surface defect segmentation model based on efficient convolution and attention mechanism, characterized by It includes a lightweight feature extraction network (GD-Net), a lightweight feature enhancement module integrated with the attention mechanism, a hierarchical adaptive upsampling fusion module, and a lightweight segmentation head module. The lightweight feature extraction network is used to efficiently extract spatial detail features and high-level semantic features of images. The feature enhancement module enhances features through further convolution and attention mechanism processing. The hierarchical adaptive fusion upsampling module is used to adaptively fuse multi-scale features in the decoding stage. The lightweight segmentation head converts the input feature map into a segmentation result with the same resolution as the input image through upsampling and convolution operations.

2. A lightweight steel surface defect segmentation model based on efficient convolution and attention mechanism according to claim 1, characterized in that: The lightweight feature extraction network (GD-Net) is constructed with 5 feature extraction blocks, which can be regarded as 5 stages of feature extraction, respectively extracting feature maps of different scales. Each feature extraction block is composed of three structures of feature extraction bottleneck blocks and double Ghost modules (Double Ghost). The three structures of feature extraction bottleneck block structures are ordinary feature extraction bottleneck block (GDFB), feature extraction bottleneck block with downsampling (Down_GDFB), and wide receptive field feature extraction bottleneck block (Wide_GDFB). The first feature extraction block and the second feature extraction block are respectively composed of a Down_GDFB and a GDFB. Down_GDFB downsamples the feature map and performs feature extraction, and GDFB performs further feature extraction on the feature map. Feature maps of 1 / 2 and 1 / 4 original image scales are obtained respectively. The third and fourth feature extraction blocks are composed of a Down_GDFB, a GDFB, and a Wide_GDFB. They downsample the input feature map first and then perform further feature extraction to obtain feature maps with a scale of 1 / 8 and 1 / 16 of the original image, respectively. The fifth feature extraction block is composed of a Down_GDFB, a GDFB, and a Double Ghost module. It first downsamples the input feature map and then uses GDFB and a lightweight Double Ghost module to perform further feature extraction to obtain a feature map with a scale of 1 / 32 of the original image. The feature enhancement modules are respectively a spatial detail information enhancement module and a semantic information enhancement module, wherein the spatial detail enhancement module is implemented in combination with the Sobel operator, and the semantic information enhancement module first extracts features from the feature map through deep convolution, then calculates the attention weight through the attention mechanism, and finally assigns the weight to the input feature map. The hierarchical adaptive upsampling and fusion module combines upsampling, attention, and feature fusion to provide an efficient feature fusion method. The lightweight attention mechanism in this module draws on the design principles of the SE module but simplifies its implementation by omitting the feature flattening and linear layers after the adaptive average pooling layer in the SE module and directly using 1x1 convolutions for dimensionality increase and reduction.

3. The lightweight steel surface defect segmentation model based on efficient convolution and attention mechanism according to claim 1 is characterized in that: The basic modules for constructing each feature extraction block in the lightweight feature extraction network are: feature extraction bottleneck block (GDFB), feature extraction bottleneck block with downsampling (Down_GDFB), wide receptive field feature extraction bottleneck block (Wide_GDFB) module, and double Ghost module (Double Ghost). The specific implementations are as follows: The feature extraction bottleneck block consists of a 1x1 ordinary convolution layer, a normalization layer, a 3x3 depth convolution layer, a Ghost convolution module, and two ReLu activation functions. In particular, it adopts a residual structure and connects the input and output with residuals to alleviate the impact of gradient disappearance in deep networks. The feature extraction bottleneck block with downsampling adds a maximum pooling layer for downsampling on the basis of the feature extraction bottleneck block. Overall, this module consists of a maximum pooling layer, a 1x1 ordinary convolution layer, a normalization layer, a 3x3 depth convolution layer, a Ghost convolution module, and two ReLu activation functions. In particular, it adopts a residual structure and performs residual connection between the input and output. The wide receptive field feature extraction bottleneck block consists of a 1x1 ordinary convolution layer, a normalization layer, a 5x5 depth convolution layer, a Ghost Module, and two ReLu activation functions. In particular, it adopts a residual structure, which connects the input and output with residual connections. The 5x5 convolution kernel is used to quickly obtain a larger receptive field. The Dual Ghost Module consists of two Ghost Modules, a 3x3 deep convolutional layer, and a ReLU activation function. It aims to generate features in a simple way at the deep network layer, reducing the number of parameters.

4. The lightweight steel surface defect segmentation model based on efficient convolution and attention mechanism according to claim 1 is characterized in that: The lightweight feature extraction network building block Ghost Module is used to generate features in a simple way in a deep network. The module consists of two branches. The first branch includes a 1x1 convolution layer, a normalization layer, and a ReLU activation function layer for generating real features. The other branch includes a 1x1 convolution layer, a normalization layer, and a ReLU activation function layer for generating cheap features. Finally, the features of the two branches are merged as the overall feature output.

5. The lightweight steel surface defect segmentation model based on efficient convolution and attention mechanism according to claim 1 is characterized in that: The spatial detail enhancement module contains two Sobel filters (sobel_x and sobel_y) to calculate detail information in the x and y directions, respectively. sobel_x and sobel_y are predefined 3x3 convolution kernels used to detect details in the horizontal and vertical directions, respectively. They are used to normalize the extracted edge features and generate an edge attention map. A sigmoid activation function is included to compress the input values ​​into the range [0, 1] and adjust the weights of the feature map. Contains a 1x1 convolution layer to reduce the number of channels of the enhanced feature map from the number of input channels to the number of output channels. The semantic feature enhancement module consists of a 3x3 depth convolution layer, an average pooling layer, a 1x1 convolution layer, and a Sigmoid activation function layer.

6. The lightweight steel surface defect segmentation model based on efficient convolution and attention mechanism according to claim 1 is characterized in that: The hierarchical adaptive fusion upsampling module consists of a bilinear interpolation upsampling layer, a normalization layer, two ReLU activation function layers, an average pooling layer, two convolution layers, and a Sigmoid activation function layer.

7. The lightweight steel surface defect segmentation model based on efficient convolution and attention mechanism according to claim 1 is characterized in that: The segmentation head module is used to map the fused features to the original image size through upsampling. It consists of a 1x1 convolutional layer, a 3x3 depthwise convolutional layer, and a 3x3 normal convolutional layer. This design can generate high-quality segmentation output without significantly increasing computational complexity. The depthwise separable convolution used in the module helps reduce computation while maintaining model performance.

Images