A small sample image defect target detection method based on self-supervised pre-training

Abstract

The application discloses a small sample image defect target detection method based on self-supervised pre-training, adopts a transfer learning paradigm in small sample target detection, that is, pre-training is carried out on base class samples, and fine-tuning is carried out on new class samples; a large number of base class normal picture samples are used for pre-training, that is, base class pre-training is carried out; the trained network model is fine-tuned on a small number of new class defect picture samples, that is, continues to be trained; and finally, testing is carried out on a test picture set, and a new class, that is, a defect class, is detected. The technical scheme of the application has strong robustness and generalization ability in the scene of small sample target detection, and can improve the detection precision of small sample image target detection. The application can be applied to high-speed rail infrastructure image processing and target detection.

Description

Technical Field

[0001] This invention relates to a deep learning self-supervised pre-training image processing method and a deep learning few-shot image target detection method. Specifically, it relates to a few-shot image defect target detection model based on self-supervised pre-training, incorporating a multi-scale attention mechanism and a contextual semantic fusion module. This model can be applied to image processing and target detection of high-speed rail infrastructure and belongs to the field of computer vision technology. Background Technology

[0002] High-speed rail infrastructure is a crucial component in ensuring the safe operation of high-speed trains. Therefore, regular inspections and defect detection of high-speed rail infrastructure are essential for maintaining stable high-speed rail operations. In recent years, with the rapid development of deep learning, an increasing number of deep learning-based models have been applied to railway infrastructure defect detection. Wei et al. proposed an intelligent method for online detection of pantograph sliding plate status based on deep learning and image processing technologies. Subsequently, Ye et al. proposed a target detection method based on differential feature fusion convolutional neural networks for railway target detection. Furthermore, Liu et al. proposed a method for detecting high-speed railway support sleeve screws based on an improved Faster RCNN. While existing deep learning-based models have achieved good results in railway infrastructure defect detection, these methods require a large number of labeled samples. However, many defect samples in high-speed rail infrastructure are difficult to find and collect, such as loose or missing tower bolts, missing bolts in bridge steel structures, and missing fasteners along the track. Additionally, high-speed rail infrastructure defects often present the problem of small target sizes, making detection difficult. Therefore, traditional models based on a large number of labeled samples struggle to detect these defects with small sample sizes. Summary of the Invention

[0003] To address the problems of the existing technologies, this invention implements a few-shot image target detection method based on self-supervised pre-training for high-speed rail infrastructure defect detection. It incorporates a multi-scale attention mechanism and a contextual semantic fusion module, which improves the detection accuracy of few-shot image targets. The few-shot image target detection method of this invention adopts the transfer learning paradigm in few-shot target detection, i.e., pre-training on base class samples and fine-tuning (i.e., continuing training) on ​​new class samples.

[0004] The technical solution provided by this invention is:

[0005] A method for detecting defect targets in few-sample images based on self-supervised pre-training, comprising the following steps:

[0006] 1) Perform self-supervised pre-training to construct a few-shot target detection network model; use the SlotCon self-supervised pre-training method to train the backbone network in the few-shot target detection network model, and obtain the weights of the backbone network after self-supervised pre-training, which are used as the initial weights of the backbone network of the few-shot target detection network model.

[0007] The SlotCon self-supervised pre-training method is combined with few-shot target detection. A large amount of unlabeled high-speed rail infrastructure data is used to pre-train on the self-supervised method SlotCon. The weights of the backbone network after self-supervised pre-training are used to replace the weights of the supervised pre-training network as the initial weights of the backbone network of the few-shot target detection network model constructed in this invention.

[0008] The small sample target detection network model constructed in this invention includes a backbone network, a gradient decoupling layer (GDL layer), a region generation network (RPN), a region of interest pooling structure (RoIPooling), a context semantic fusion module, a classifier, and a bounding box regressor.

[0009] In specific implementation, this invention employs the SlotCon self-supervised method to train the ResNet101 backbone network. Contrastive learning is performed from data-driven semantic slots to jointly handle semantic grouping and representation learning (automatically learning effective features and improving model performance). Semantic grouping assigns pixels to a set of learnable prototypes, which can aggregate features to adapt to each sample and form new slots. Based on the learned data-dependent slots, contrastive target representation learning is used to enhance feature discriminability. The trained ResNet101 backbone network weights replace the supervised pre-trained weights for subsequent object detection.

[0010] 2) This invention employs a transfer learning paradigm for training in small-sample target detection. First, a large number of normal base class samples are used for pre-training (base class pre-training). The pre-trained network model is then fine-tuned on a small number of new defect class samples (continued training). Finally, the model is tested on a test set to detect the new class (defect class). The same network model structure is used in both the pre-training and fine-tuning stages. Specifically, images of high-speed rail infrastructure are fed into a few-shot object detection network model. First, the images pass through the backbone network of the few-shot object detection model, which extracts image features. This backbone network consists of a ResNet101 residual network (backbone network) and a Feature Pyramid Network (FPN). At the same time, an SENet (squeeze and excitation network) attention mechanism is added to the last layer of the residual network, forming a multiscale attention mechanism based on SENet (SE-MAM). The feature map after passing through the backbone network is fed into the GDL layer for forward propagation. The output is then fed into the Region Generation Network (RPN) and the Region of Interest Pooling (RoIPooling) structure. The RPN network provides regression boxes that may contain objects, generating a feature vector with proposal information containing object scores and bounding box regression offsets. RoIPooling uses pooling methods to obtain a fixed-size output feature map from regions of interest of different sizes in the input feature map.

[0011] 3) The output feature map obtained after RoIPooling is passed through the context semantic fusion module (CSF) and output to the classifier and bounding box regressor. The classifier calculates the probability of each object in the candidate box for each category using the softmax function. The final predicted category is the category with the highest output probability. The bounding box regressor uses the smooth L1 loss function to calculate the loss value of the offset between the predicted value and the true value. The original candidate box is corrected by the offset to obtain the final predicted box coordinates, that is, the position of the identified target, thus realizing small sample target detection.

[0012] To address the issue of supervised pre-training methods lacking specificity in providing features within the complex context of railways, this invention employs the self-supervised pre-training method SlotCon, using the weights obtained from self-supervised pre-training as the backbone network weights of the few-shot detector. To improve the model's ability to recognize small targets and enhance its sensitivity to channel features, this invention proposes a multi-scale attention mechanism, comprising a Feature Pyramid Network (FPN) and a SENet attention mechanism. The feature map output from the multi-scale attention network is fed into a Gradient Decoupling Layer (GDL), which adjusts the degree of decoupling between different modules. During forward propagation, an affine transformation layer A is used to simply enhance the feature representation. During backward propagation, the GDL obtains gradients from subsequent layers, multiplies them by coefficients λ∈[0,1], and then passes them to the previous layer. The GDL can be viewed as a pseudo-function G defined by two equations. (A，λ) These two equations describe its forward and backward propagation behavior, as follows:

[0013] G (A，λ) (x)=A(x) (1)

[0014]

[0015] in is the Jacobian matrix of the affine transformation layer; x is the input feature map; A(x) indicates that the feature map is input into the affine transformation layer.

[0016] The Semantic Context Fusion (CSF) module fuses features at different scales and learns well-developed global and local features. The output of the CSF module is fed into a classifier and a bounding box regressor. The classifier obtains the predicted class score, and the bounding box regressor obtains the predicted coordinates, thus achieving few-shot object detection. Implementation results demonstrate that this invention exhibits strong robustness and generalization ability in few-shot object detection scenarios, outperforming state-of-the-art (SOTA) few-shot object detection models. Attached Figure Description

[0017] Figure 1 This is a schematic diagram of the SlotCon framework, the self-supervised method used in this invention.

[0018] Figure 2 This is a schematic diagram of the network structure of the small sample image defect target detection model constructed in this invention, which is based on a multi-scale attention mechanism and a contextual semantic fusion module. Detailed Implementation

[0019] The model structure of this invention includes: (1) a backbone network, which includes a residual network structure, a feature pyramid structure (FPN), and a SENet attention mechanism; (2) an RPN with a gradient decoupling layer and RoIPooling; (3) a context semantic fusion module, the output of which is fed into a classifier and a bounding box regressor. The classifier obtains the predicted class score, and the bounding box regressor obtains the predicted coordinates.

[0020] The self-supervised pre-training method SlotCon framework used in this invention is as follows: Figure 1 As shown, for an unlabeled image dataset D, the self-supervised pre-training method SlotCon aims to obtain a set of prototypes S to classify pixels in the image. Pixels in the same prototype set have similar feature representations. SlotCon uses a pixel-level deep clustering method to obtain prototypes S. Specifically, the SlotCon self-supervised method consists of two networks with the same structure but different parameters: a student network and a teacher network. The student network includes an encoder f. θ and the mapping layer (Projector) g θ And K learnable prototypes The weight set of the teacher network is ξ, and the network weights are updated using an exponential moving average. Given an input image x, two augmented views v are generated using two random augmentation methods. l ∈{v 1 ,v 2 The encoders f of the student network and the teacher network output feature maps respectively. (The feature map output from the enhanced image after passing through the encoders of the student and teacher networks contains three dimensions: height, width, and channels), and then the feature map is obtained through a multilayer perceptron (MLP). (The feature map obtained after the encoder passes through a multilayer perceptron, containing three dimensions: height, width, and channels). Then, the prototype S is used... θ Calculation features Assignment Make it connect with the teacher network in another view v l 'Generated assignment' Matching. At the pixel level, for overlapping regions in two views, pixels at the same location are consistently assigned to prototypes. In object-level representation learning, pixels with semantically similar characteristics on the feature map are aggregated together. Contrastive learning is performed on semantic slots across different views. This mutually promotes and optimizes both directions. In the self-supervised pre-training phase, we use the SlotCon self-supervised method to train the ResNet101 backbone network, performing contrastive learning from data-driven semantic slots for joint semantic grouping and representation learning. Semantic grouping assigns pixels to a set of learnable prototypes, which can aggregate features to adapt to each sample and form new slots. Based on the learned data-dependent slots, contrastive targets are used for representation learning, enhancing the discriminability of features. The trained ResNet101 backbone network weights replace the supervised pre-trained weights for subsequent object detection.

[0021] like Figure 2As shown, the backbone network of this invention includes: a multi-scale feature pyramid to solve the problem of small target size of defect samples. The feature pyramid structure includes a bottom-up path, a top-down path, and lateral connections. Specifically, the bottom-up stage employs a typical CNN model. Through a series of convolutional operations, feature maps of different sizes are obtained. Feature maps of the same size are grouped into one stage. Each extracted feature is the output of the last layer of each stage, thus forming a feature pyramid. Finally, four stages are generated from the bottom up: conv2 (C2), conv3 (C3), conv4 (C4), and conv5 (C5). Taking a 224×224×3 input image as an example, after a series of convolutional operations, the feature map dimensions output by C2, C3, C4, and C5 are 56×56×256, 28×28×512, 14×14×1024, and 7×7×2048, respectively. In the top-down stage, the high-level feature maps are upsampled. This upsampling is a 2x upsampling, ensuring that the height and width of the feature maps are the same. Lateral connections are used for addition and fusion operations. The upsampling algorithm uses nearest neighbor interpolation. Then, this feature is laterally connected to the previous layer's feature map to enhance the higher-level features. A 1×1 convolutional kernel is used for these lateral connections, primarily to adjust the number of output channels in different feature layers, ensuring each layer has 256 channels to facilitate the summation of other upsampled features. Each lateral connection fuses feature maps of the same spatial size from both the bottom-up and top-down paths. After fusion, a 3×3 convolutional kernel is used to convolve each fused result to eliminate the aliasing effect of upsampling. Assume the generated feature maps are P2, P3, P4, and P5, corresponding one-to-one with the original bottom-up convolutional results C2, C3, C4, and C5.

[0022] SENet enhances the learning of convolutional features by explicitly modeling channel dependencies through compressed-activation blocks, thereby increasing the network's sensitivity to information available in subsequent transformations. Let the output after convolution be U. To alleviate the problem that the convolution kernel uses a local receptive field during convolution, preventing each unit of output U from utilizing contextual information outside that region, global average pooling is used to compress global spatial information into the channel dimension. Formally, statistics... It is generated by reducing the spatial dimension of U to H×W, and the c-th element of z can be calculated using the following formula:

[0023]

[0024] To utilize the information aggregated during the compression operation, a second operation is used to capture channel dependencies. This operation employs a simple gating mechanism with a sigmoid activation function:

[0025]

[0026] Where δ is the ReLU activation function. To reduce model complexity and improve generalization ability, a bottleneck structure with two fully connected (FC) layers is adopted. The first FC layer performs dimensionality reduction with a reduction coefficient of r, followed by a ReLU activation function. The final FC layer restores the original dimensionality. The final output is obtained by multiplying the learned activation values ​​of each channel by the original features on U.

[0027]

[0028] in F scale (u c ,s c ) is the indicator quantity s c With feature map u c ∈R H×W The channel product between them.

[0029] The SE-MAM module adds SENet after the last stage of ResNet101, which does not destroy the complete structure of the residual network, but allows the feature layers at each scale to obtain sensitivity to channel features.

[0030] like Figure 2 As shown, this invention proposes a contextual semantic fusion (CSF) module to address the information loss problem that can easily occur during training due to single pooling operations. We introduce three different pooling scales after RoIPooling and perform semantic fusion to better capture global and local features. Specifically, the CSF module does not use a fixed resolution size; instead, it selects three different resolutions—6, 12, and 18—for parallel pooling operations to obtain a more comprehensive feature representation. Larger resolutions focus more on global information, while smaller resolutions focus more on local information, thus better utilizing both global and local information for object detection. After pooling, each resolution's features undergo semantic fusion through two branches. The first branch contains a fully connected layer; the second branch contains a global average pooling layer, a fully connected layer, and an upsampling layer. The three features from the three different resolutions are then merged, and finally, a 1×1 convolution is used for upsampling to restore the feature map to its original size for output.

[0031] The output after passing through the CSF module is further processed through two fully connected layers to extract features. Finally, a bounding box regressor and a classifier are used to perform bounding box regression and class prediction, thereby realizing defect target detection of small sample images based on self-supervised pre-training.

[0032] In a specific implementation of this invention, an image dataset collected by a drone is used to implement and evaluate the method of this invention, as further described below:

[0033] 1. Experiments were conducted on a dataset of high-speed rail infrastructure images captured by drones. As shown in Table 1, the drone aerial image dataset we collected contains 16 categories of high-speed rail infrastructure. First, the 5091 images were divided into a training image set and a test image set at a ratio of 1.2:1. Then, following the dataset partitioning method for small sample object detection, the categories of the high-speed rail infrastructure dataset were divided into a base class and a new class, with the normal category as the base class and the defect category as the new class. Then, according to the partitioning method, K = 1, 2, 3, 5, 10 samples were randomly selected from each category of the new class as fine-tuning samples.

[0034] This experiment uses average precision (AP) as the evaluation metric to assess the detection performance of different algorithms. mAP50 is the average AP50 across all categories. AP50 refers to the threshold of 0.5 used when calculating average precision (AP) for IoU (Intersection over Union, representing the degree of overlap between detected objects and real objects).

[0035] Table 1. Data distribution of the UAV dataset

[0036]

[0037] The small-sample image target detection model of this invention employs an end-to-end stochastic gradient descent algorithm to optimize parameters during training. Stochastic gradient descent involves randomly selecting a small batch of samples in each iteration to calculate the gradient of the loss function, and then using this gradient to update the parameters. The experimental setup is shown in Table 2.

[0038] Table 2 Experimental parameter settings

[0039]

[0040] The experimental environment is shown in Table 3.

[0041] Table 3 Experimental Environment

[0042]

[0043] The number of iterations for the base class pre-trained model and the number of iterations for fine-tuning the model using the new class are shown in Table 4:

[0044] Table 4 shows the model's performance during training and fine-tuning.

[0045]

[0046] 2. First, the ResNet101 backbone network is trained using the self-supervised method SlotCon, and the trained weights are used as the initial weights of the backbone network of the few-shot image object detection model. Then, images are fed into the few-shot image object detection network model for training and inference. Images are fed into the network model in batches. The deep learning framework PyTorch is used to convert the input images from NumPy format to tensor format before input. The images pass through the backbone network, which consists of a residual network structure and a multi-scale attention mechanism based on SENet. The output feature maps are then fed into the gradient decoupling layer (GDL), and the outputs are fed into the Region of Interest Pooling (RPN) and the Region of Interest Pooling (RoIPooling) structure. The feature maps obtained after RoIPooling are then processed through the Semantic Context Fusion (CSF) module and output to the classifier and bounding box regressor. The classifier obtains the final predicted class score, and the bounding box regressor obtains the final predicted coordinates.

[0047] 3. Experiments were conducted on small-sample defect target detection based on a self-supervised pre-trained small-sample target detection model, and the experimental results were obtained. The mean accuracy (mAP50) of the six new classes in the UAV dataset is shown in Table 5.

[0048]

[0049] Clearly, the self-supervised pre-trained few-shot defect detection model proposed in this invention outperforms other state-of-the-art (SOTA) methods in detecting new classes. Among all compared methods, this invention achieves the best high-speed rail infrastructure defect detection results except for 10 shots (mAP50 of 27.4%, 30.5%, 33.6%, and 34.0% on 1, 2, 3, and 5 shots, respectively).

[0050] Finally, it should be noted that the purpose of disclosing the embodiments is to help further understand the present invention. However, those skilled in the art will understand that various substitutions and modifications are possible without departing from the scope of the present invention and the appended claims. Therefore, the present invention should not be limited to the content disclosed in the embodiments, and the scope of protection of the present invention is defined by the claims.

Claims

1. A small sample image defect target detection method based on self-supervised pre-training, characterized in that, The model is trained by using a transfer learning paradigm in small sample target detection, that is, pre-training on base class samples and fine-tuning on new class samples; wherein the normal class is taken as the base class and the defect class is taken as the new class; including the following steps: 1) Self-supervised pre-training is performed to construct a small sample target detection network model; a backbone network in the small sample target detection network model is trained by using a self-supervised pre-training method, and the weight of the self-supervised pre-trained backbone network is obtained as the initial weight of the backbone network of the small sample target detection network model; The constructed small sample target detection network model includes a backbone network, a gradient decoupling layer (GDL), a region proposal network (RPN), a region of interest pooling structure (RoI Pooling), a context semantic fusion module (CSF), a classifier and a bounding box regressor; 2) The model is trained by using a transfer learning paradigm in small sample target detection: first, pre-training is performed on a large number of base class normal image samples, that is, base class pre-training; then, the trained network model is fine-tuned on a small number of new class defect image samples, that is, further training; finally, the model is tested on a test image set to detect the new class, that is, the defect class; the same network model structure is used in the pre-training and fine-tuning stages; The specific process includes: The image is input into the small sample target detection network model, and the image is input into the backbone network of the small sample target detection model; the backbone network is used to extract the features of the image, and the backbone network is composed of a residual network (ResNet101) and a feature pyramid network (FPN); at the same time, an SE-Net attention mechanism is added to the last layer of the residual network to form an SE-Net-based multi-scale attention mechanism (SE-MAM); The feature map after passing through the backbone network is input into the GDL layer for forward propagation, and the output is input into the region proposal network (RPN) and the region of interest pooling structure (RoI Pooling); wherein the RPN network is used to provide a regression box that may contain a target, and to generate a suggestion information feature vector with a target score and a boundary box regression offset; the RoI Pooling is used to obtain a fixed-size output feature map from different sizes of the region of interest in the input feature map by using a pooling method; 3) The output feature map obtained after the RoI Pooling is output to the classifier and the bounding box regressor through the context semantic fusion module; the classifier calculates the probability of the object in the candidate box belonging to each class, and outputs the class with the maximum probability as the predicted class; the bounding box regressor uses a loss function to calculate the loss value of the offset between the predicted value and the true value, and modifies the candidate box through the offset to obtain the predicted box coordinates, that is, the position of the recognized target; The trained small sample target detection network model is used to realize small sample target detection.

2. The small sample image defect target detection method based on self-supervised pre-training according to claim 1, characterized in that, Specifically, the SlotCon self-supervised pre-training method is used to train the backbone network ResNet101, which is used for joint semantic grouping and representation learning through comparison learning from data-driven semantic slots. The semantic grouping concentrates feature adaptation on each sample by assigning pixels to a set of learnable prototypes, and forms new slots; based on the learned data-dependent slots, representation learning is performed using contrastive objectives to enhance feature discriminability.

3. The small sample image defect target detection method based on self-supervised pre-training according to claim 2, characterized in that, The multi-scale attention mechanism includes a feature pyramid network (FPN) and a SE-Net attention mechanism, and the feature map output after the multi-scale attention mechanism network enters a gradient decoupling layer (GDL) for adjusting the decoupling degree between different modules. During forward propagation, an affine transformation layer is employed enhanced feature representation; When backpropagating, the GDL takes the gradient from the subsequent layer, multiplies it by the coefficient The backpropagation is passed to the previous layer; GDL is a pseudo-function defined by two equations for describing the behavior of forward and backward propagation, denoted as: (1) (2) wherein, is the Jacobian matrix of the affine transformation layer; x is the input feature map; represents inputting the feature map into the affine transformation layer.

4. The small sample image defect target detection method based on self-supervised pre-training according to claim 1, characterized in that, The classifier is specifically a softmax function for calculating the probability of an object in a candidate frame belonging to each class; and the frame regressor is specifically a smooth L1 loss function for calculating the loss value of the offset between the predicted value and the true value.

5. The small sample image defect target detection method based on self-supervised pre-training according to claim 1, characterized in that, An SE-Net attention mechanism is added to the last layer of the residual network. The SE-Net models the channel interdependence through a compression-excitation block to enhance the learning of convolutional features. The process includes: Let the output after convolution be U, by reducing the spatial dimension of U generate statistics ; the first element of the vector is calculated by the following formula: (3) wherein is a feature map; A gating mechanism with a sigmoid activation function is used to capture channel dependencies, represented as: (4) wherein is a ReLU activation function, , ; is a dimensionality reduction coefficient; A bottleneck structure with two fully connected layers (FC) is used for dimension reduction and dimension recovery, respectively; the learned activation values of each channel are multiplied by the features on U to obtain the final output, represented as: (5) wherein, , is an index quantity, is an index quantity channel product between the feature map and the feature map

Images