An image fault fine-grained identification method based on an attention mechanism

Abstract

The application provides an image fault fine-grained identification method based on an attention mechanism, and on the basis of a fine-grained identification network NTS-Net, a soft-pooling double-channel attention network is used, aiming at the black, gray and white backgrounds of the railway data set, and through two channels, more effective feature semantics are obtained from two dimensions. In order to avoid the loss of feature information when passing through the attention network, the soft-pooling method is used to avoid the information loss in the pooling process, which is helpful to the subsequent link. Finally, by changing the network loss function, the interference brought by the imbalance between the fault samples and the normal samples to the result is weakened without the need of data enhancement on the fault pictures. The identification method adopts a weak supervision method, is efficient and fast, has wide application range and low computing power consumption.

Description

Technical Field

[0001] This invention belongs to the field of image recognition technology, specifically relating to a fine-grained image fault recognition method based on an attention mechanism. Background Technology

[0002] This invention focuses more on specific engineering applications, aiming to solve the problem of fault detection in railway freight car components within the current railway system. Because railway freight cars contain a wide variety of hardware and parts, both large and small components can exhibit minute faults. Traditional target detection methods are not ideal for detecting very small faults, such as scratches and cracks. Fine-grained methods focus on more subtle features; therefore, fine-grained image recognition methods are used to detect and provide early warnings for minute faults occurring in railways, thus assisting inspection personnel.

[0003] Fine-grained identification involves analyzing all objects in all subclasses of a given class, such as flower varieties or airplane types. Thanks to the rise of deep learning, fine-grained identification has made significant progress in recent years, but problems such as insufficient feature extraction and inadequate granularity still exist. Furthermore, due to the limited and imbalanced nature of fault samples, as well as the small inter-class differences and large intra-class differences between fault and non-fault samples, fine-grained fault identification still has considerable room for improvement.

[0004] Currently, there are two main types of fine-grained image recognition algorithms: fine-grained models with strong supervision information and fine-grained models with weak supervision information.

[0005] Strongly supervised fine-grained models require additional auxiliary information during training, such as bounding boxes of objects in the image and annotations of parts that are useful for training, in order to achieve higher recognition accuracy.

[0006] While strong supervision has achieved significant progress in fine-grained image recognition, relying on manually labeled tags, bounding boxes, and auxiliary text, and has yielded high accuracy on publicly available fine-grained datasets, the labeling process is tedious and requires substantial manual resources. Weakly supervised fine-grained models, on the other hand, utilize only image category labels during network training, requiring no additional information, yet achieve classification accuracy similar to or higher than that of strongly supervised models. Therefore, research on fine-grained algorithms with weak supervision is currently a mainstream trend.

[0007] While some existing fine-grained recognition methods have achieved good progress on public datasets, this paper focuses on a railway freight car dataset that differs from the fine-grained public dataset CUB-200-2011. Images from the railway freight car dataset are shown below. Figure 2 and Figure 3 As shown, the backgrounds of the image samples are mostly black, white, or gray, and there is no clear boundary between the target object and the background. Due to the uniform color of railway freight cars, ordinary feature extraction networks often extract irrelevant feature information, thus interfering with the final detection. This is extremely detrimental to the learning of fine-grained recognition networks. Furthermore, because the frequency of faults is very low, there are few fault samples, resulting in a highly unbalanced quantitative relationship between fault samples and normal samples. Simply expanding the data by augmenting fault samples still requires manual annotation, which has drawbacks. Summary of the Invention

[0008] Purpose of the invention: To address the technical problems of existing fine-grained image recognition methods, such as the inability to extract useful fine-grained features due to the homogeneity of the dataset background and the poor overall classification results caused by the imbalance between the number of faulty and normal samples, a fine-grained image fault recognition method based on an attention mechanism is proposed.

[0009] To achieve the above objectives, the present invention adopts the following technical solution: a fine-grained image fault recognition method based on an attention mechanism, comprising the following steps:

[0010] Step 1: After preprocessing, the collected samples are input into a pre-trained and optimized Navigator network for feature extraction, resulting in a series of rectangular regions. These rectangular regions are then input into the STAN network for weight redistribution. The K regions with the most information after weight redistribution are mapped from the original input samples to obtain K mapped regions.

[0011] Step 2: Input the K mapping regions into the Scrutinizer network for fine-grained classification and identification to obtain the category of the sample. For data samples with a highly imbalanced relationship between faulty and normal samples, the loss function in the Scrutinizer network is the Focal Loss function when performing fine-grained classification and identification.

[0012] Furthermore, the STAN network includes an input terminal, a Conv layer, a BatchNorm layer, a SoftPool layer, a Sigomid layer, an adder, a multiplier, and an output terminal;

[0013] The input terminals are connected to the Conv layer, the SoftPool layer, and the adder, respectively. The adder is connected to the BatchNorm layer and the multiplier, respectively. The multiplier is connected to the Sigomid layer and the output terminal, respectively. The Conv layer is connected to the BatchNorm layer. The SoftPool layer is connected to another Conv layer, and the other Conv layer is connected to the Sigomid layer.

[0014] Furthermore, the formula for the Focal Loss function is as follows:

[0015] FL(p t )=-(1-p t ) γ log(p t )

[0016] Where, p t y represents the probability that the sample belongs to the correct category, p represents the probability that the predicted sample belongs to 1, y represents the label, and γ is the focusing parameter adjusted using cross-validation, which is also the modulation factor.

[0017] Furthermore, the specific training steps of the Navigator network include:

[0018] Step A1: The preprocessed image samples are input into the Navigator network for feature extraction to obtain a series of rectangular regions. Then, the series of rectangular regions are input into the STAN network for weight redistribution. The M regions with the highest information richness in the redistributed regions are mapped from the fault samples to obtain the M most discriminative regions.

[0019] Step A2: Input the M most discriminative regions into the Teacher network, obtain the confidence scores of the M regions through the Teacher network, and then calculate the loss value of the confidence scores of the M regions using the Teacher network's loss function. Continuously reduce the loss value to optimize the Teacher network. For data samples with a highly imbalanced relationship between faulty and normal samples, the loss function in the Teacher network is Focal when calculating the loss value of the confidence scores of the M regions. The loss function is used to calculate the confidence scores of the M regions after optimization. Then, the confidence scores of the M regions are re-weighted using the optimized Teacher network. Next, the STAN network is used to redistribute the weights of these M regions, resulting in the weighted confidence scores of the M regions. These weighted confidence scores are then fed back to the Navigator network. The Navigator network's loss function calculates the loss value that aligns the information richness of the M most discriminative regions with the confidence scores of the weighted M regions. The Navigator network is trained by gradually decreasing the loss value until the confidence ranking of the weighted M regions matches the information richness ranking of the M most discriminative regions. For data samples with a highly imbalanced mix of faulty and normal samples, the Focal Loss loss function is used in the Navigator network when calculating the loss value that aligns the information richness of the M most discriminative regions with the confidence scores of the weighted M regions.

[0020] Furthermore, the classification function used in step 2 is the softmax function.

[0021] Beneficial Effects: Firstly, a STAN network is introduced based on the NTS-Net network. The STAN network, through exponential weighting, maximizes the preservation of information from the original feature maps during the downsampling stage, avoiding information loss and better serving subsequent operations. Secondly, when the input data samples are highly imbalanced between faulty and normal samples, the loss functions of the Navigator, Teacher, and Scrutinizer networks are all replaced with Focal Loss. This addresses the imbalance in the proportion of faulty and normal samples among different types of faults in the private railway freight car dataset, improving overall classification accuracy. The algorithm of this invention is an effective fine-grained image fault recognition algorithm that can effectively improve the overall classification accuracy of faults while meeting the real-time requirements of industrial production. Attached Figure Description

[0022] Figure 1 This is a schematic diagram of the network structure of the image fault fine-grained recognition method based on the attention mechanism of the present invention;

[0023] Figure 2 This is a fault diagram of the upper pull rod malfunctioning in the training and testing datasets of this invention;

[0024] Figure 3 This is a fault diagram of mitigating valve lever loss in the training and test datasets of this invention;

[0025] Figure 4 This is a schematic diagram of the STAN attention module structure of the present invention;

[0026] Figure 5 This is a schematic diagram of the Scrutinizer network structure for the present invention.

[0027] Figure 6 This is a result image of the first four regions with the highest information richness when the upper pull rod extends out in this invention;

[0028] Figure 7 This is a result diagram of the four most information-rich areas in the present invention for mitigating valve rod loss. Detailed Implementation

[0029] The invention will now be further explained with reference to the accompanying drawings.

[0030] This invention provides a fine-grained image fault recognition method based on an attention mechanism, comprising:

[0031] Step 1: After preprocessing, the collected samples are input into a pre-trained and optimized Navigator network for feature extraction, resulting in a series of rectangular regions. These rectangular regions are then input into the STAN network for weight redistribution. The K regions with the most information after weight redistribution are mapped from the original input samples to obtain K mapped regions.

[0032] Step 2: Input the K mapping regions into the Scrutinizer network for fine-grained classification and identification to obtain the category of the sample. For data samples with a highly imbalanced relationship between faulty and normal samples, the loss function in the Scrutinizer network is the Focal Loss function when performing fine-grained classification and identification.

[0033] The software environment required to run the above method includes a PyTorch environment, a compatible CUDA version, and a GeForce RTX 3090 graphics card.

[0034] like Figure 1 As shown, in step 1, the input fault samples, i.e., the input images, are input. The size of the input images is normalized to 448×448, and data augmentation operations such as random cropping and horizontal flipping are performed. Before the images are fed into the network for training, the mean of the images is subtracted, and a normalization operation is performed. Information from the training and test sets is also read from a txt file. Through data augmentation and normalization, the generalization ability of the model can be improved and the influence of common parts can be eliminated.

[0035] The preprocessed input image is fed into the Navigator network. Within the Navigator network, a ResNet-50 pre-trained model generates a series of rectangular regions based on different scales and sizes. These rectangular regions are then fed into the STAN network for weight reassignment. Next, the rectangular regions are sorted by information richness, and the M regions with the highest information content are selected and mapped from the original image. The Navigator network uses a ResNet-50 pre-trained model, which is used for feature extraction from the preprocessed input image.

[0036] All rectangular regions are fed into the STAN network. Within the STAN network, a set of convolutional networks is used to increase the receptive field of each convolutional layer. Then, normalization is applied to stabilize the distribution of input data in each layer. These are then added to the original features to obtain a cumulative feature sum. Next, the original features are processed using a SoftPool, compressing each feature channel into a global pixel. This pixel is then subjected to a one-dimensional convolution with a configurable kernel size, and the weights of each channel are output using the Sigmoid function. Finally, these weights are multiplied by the corresponding elements of the cumulative feature sum to obtain the final output feature map. The feature map is then reweighted, selecting some feature maps with higher weights (i.e., the reweighted regions). These high-weight feature maps are then cropped from the input fault samples to obtain discriminative regions (i.e., cropped images). This process weakens the weights of the background and other similar features, achieving a good weight redistribution effect. Simultaneously, the SoftPool, through exponential weighting, ensures that the network retains as much information as possible from the original feature map during the downsampling stage, avoiding information loss and better serving subsequent operations.

[0037] During Navigator network training, input image samples undergo preprocessing. The Navigator network extracts features from the preprocessed fault samples, resulting in a series of rectangular regions. All these rectangular regions are then input into the STAN network for weight redistribution. The regions with redistributed weights are then sorted by information richness, and the M regions with the highest information richness are selected. These M regions are mapped from the original image to obtain the M most discriminative regions, which are then fed into the Teacher network.

[0038] The Teacher network evaluates the confidence of M regions, then calculates the loss value of the confidence of the M regions using the loss function in the Teacher network. The Teacher network is continuously optimized by reducing the loss value. The optimized Teacher network is then used to obtain the optimized confidence of the M regions, which is then input into the STAN network. The STAN network then reallocates the weights of the optimized confidence of the M regions, obtaining the reallocated confidence of the M regions, further improving the accuracy of the confidence of the M regions.

[0039] In this embodiment, when the ratio of faulty samples to normal samples is balanced, the loss function in the Teacher network is the Teaching loss. The loss function in the Teacher network calculates the prediction loss of the confidence scores of the M regions of the Teacher network. By continuously reducing this loss, the confidence prediction of the Teacher network becomes more and more accurate. The mathematical expression of the Teaching loss is:

[0040]

[0041] Among them, R i This represents the i-th most discriminative region feature generated by the Navigator network.

[0042] The Teacher network feeds back the confidence scores of the M reassigned regions to the Navigator network. The Navigator network's loss function calculates the information richness of the M most discriminative regions and the loss value of the confidence scores of the M reassigned regions, optimizing the ranking learning algorithm until the information richness of the M most discriminative regions and the confidence scores of the M reassigned regions are aligned. This process trains the Navigator network, resulting in a pre-trained and optimized Navigator network that can identify more discriminative local regions. The number of discriminative regions can be set to any value, but in this invention, experimental results show that a value of M = 6 yields the best recognition performance. The information richness of a region represents the accuracy of fine-grained object classification; the more discriminative the region, the better the classification result. Both the Navigator and Teacher networks are optimized using SGD networks.

[0043] In this embodiment, when the ratio of faulty samples to normal samples is balanced, the loss function in the Navigator network is the Navigation loss. The loss function in the Navigator network calculates the information richness of the M most discriminative regions and the confidence loss of the M reassigned regions. By continuously reducing this loss, the Navigator network is trained more thoroughly. The mathematical expression for Navigation loss is:

[0044]

[0045] Where f is the non-increasing function, I represents the information richness of the M regions, and C represents the confidence function. i C represents the confidence level of the i-th region. s The confidence level of the s-th region.

[0046] STAN networks such as Figure 4As shown, the STAN network includes an input terminal, a Conv layer, a BatchNorm layer, a SoftPool layer, a Sigomid layer, an adder, a multiplier, and an output terminal. The input terminal is connected to the Conv layer, the SoftPool layer, and the adder, respectively, and the adder is connected to both the BatchNorm layer and the multiplier. The multiplier is connected to both the Sigomid layer and the output terminal. The Conv layer is connected to the BatchNorm layer. The SoftPool layer is connected to another Conv layer, and this other Conv layer is connected to the Sigomid layer.

[0047] The STAN network, based on DCAN, changes the pooling method to SoftPool. By replacing the pooling function, the STAN network retains the information of the original feature map to the greatest extent during the downsampling stage, avoiding information loss and better serving subsequent operations.

[0048] Mapping the rectangular region from the original image to obtain the mapped region involves sorting the local regions extracted by the Navigator network from highest to lowest information richness and selecting the M regions with the highest information richness. In this embodiment of the invention, M is set to 6, which yields the best results.

[0049] like Figure 5 As shown, in step 2, K mapping regions are input into the Scrutinizer network for fine-grained classification and identification to obtain the category of the sample. When the input K mapping regions are data samples with a highly imbalanced relationship between faulty samples and normal samples, the loss function of the Scrutinizer network is changed to the Focal Loss loss function during fine-grained classification and identification.

[0050] In this embodiment, when the ratio of faulty samples to normal samples is balanced, the loss function in the Scrutinizer network is the Scrutinizing loss. The loss function of the Scrutinizer network calculates the loss of the feature map resulting from the fusion of the K local regions generated by the Navigator network and the features of the original input samples. By continuously reducing this loss, the class prediction of the Scrutinizer network becomes increasingly accurate. The mathematical expression for the Scrutinizing loss is:

[0051] L S = -log S(X,R1,R2,…,R) K )

[0052] Where S represents the Scrutinizer network, X represents the features of the original input sample, and R1, R2, ..., R KThis represents the K most informative regional features generated by the Navigator network.

[0053] When processing input samples using a loss function during implementation, the following two scenarios are included:

[0054] The first approach is to use the cross-entropy loss function of the baseline network when the ratio of faulty samples to normal samples is balanced, and then calculate the loss value using the loss function of the baseline network.

[0055] The loss function expression for the baseline network is:

[0056]

[0057] Make p t Let represent the probability that a sample belongs to the correct class. Then the formula can be simplified to:

[0058] CE(p,y)=CE(p t ) = -log(p t )

[0059] Where p represents the probability of predicting that the sample belongs to 1, and y represents the label.

[0060] The second approach is to use the Focal Loss function when the ratio of faulty samples to normal samples is unbalanced. The loss value is calculated using the Focal Loss function.

[0061] The formula for Focal Loss is:

[0062] FL(p t )=-(1-p t ) γ log(p t )

[0063] Here, γ is the focusing parameter adjusted using cross-validation, which is also the modulation factor. By adjusting γ, the weights of positive and negative samples can be adjusted, and the weights of samples classified as easy or difficult can be controlled.

[0064] Focal Loss is achieved by reducing the weights. Weighting is a technique that reduces the impact of easy samples on the loss function, thereby allowing more attention to be paid to difficult samples.

[0065] In the embodiments of the present invention, when the proportion of faulty samples and normal samples is balanced, the corresponding loss functions used by the Navigator network, Teacher network, and Scrutinizer network are Navigation loss, Teaching loss, and Scrutinizing loss, respectively. When the proportion of faulty samples and normal samples is balanced, the loss function used by the Navigator network, Teacher network, and Scrutinizer network is Focal Loss. By adjusting the size of the γ parameter, the contribution of faulty samples and normal samples to the overall loss is adjusted, thereby solving the problem of the imbalance between the proportion of faulty samples and normal samples in different types of faults in the private railway freight car dataset and improving the overall classification accuracy.

[0066] The overall network structure is optimized by calculating the total loss of avigation loss, teaching loss, and scrutinizing loss, and the contribution weights of the Navigator network, Teacher network, and Scrutinizer network to the total loss are controlled. The total loss can be expressed as:

[0067] L total =L I +λ·L S +μ·L c

[0068] Here, λ and μ are hyperparameters, which are generally set to λ = μ = 1.

[0069] There are two main methods for feature fusion: feature concatenation and pixel-by-pixel addition. Feature concatenation increases the number of feature channels, allowing for the acquisition of more feature information. Pixel-by-pixel addition, however, does not change the size or number of channels of the feature map, and this method can pass information from different levels down the hierarchy.

[0070] In embodiments of this invention, feature fusion is performed using feature concatenation, with the fused features fed into a fully connected layer. Feature concatenation requires that the feature maps have the same resolution for subsequent stitching operations. The feature tensor can be represented as:

[0071]

[0072] Where X and Y represent different feature tensors, R represents the matrix space, H and W represent the height and width of the feature tensor, respectively, and C X C represents the number of channels in the X tensor. Y This represents the number of channels in the Y tensor.

[0073] The feature cascading fusion method can obtain more feature information by increasing the number of feature channels.

[0074] The training process is end-to-end, resulting in a trained model.

[0075] The classification function is the softmax function.

[0076] The loss function in the Navigator network optimizes the network by continuously reducing the information richness of the M most discriminative regions and the confidence of the M reassigned regions. The loss function in the Teacher network optimizes the prediction loss by continuously reducing the confidence of the M regions. The loss function in the Scrutinizer network optimizes the network by continuously reducing the loss of the feature map obtained by fusing features from the K local regions and the original input samples. After optimization, a trained fault model with the suffix .ckpt is obtained.

[0077] The pre-trained and optimized Navigator network, Teacher network, Scrutinizer network, and STAN module constitute the overall fault model.

[0078] The fault model checking steps specifically include two parts: predicting the image category and calculating the overall classification accuracy.

[0079] Predicting the image category involves loading the image and a trained ckpt model, applying the softmax classification function to obtain the predicted probability of belonging to all categories, and then taking the category with the highest predicted probability as the final prediction result.

[0080] In the process of calculating the overall classification accuracy, the prediction results are classified by a classification function to obtain positive and negative samples, and the overall classification accuracy of the positive and negative samples is calculated.

[0081] The expression for calculating the overall classification accuracy is:

[0082]

[0083] Wherein, OA represents the proportion of correctly classified samples out of all test samples, TP represents the number of positive samples predicted as positive by the model, TN represents the number of negative samples predicted as negative, FN represents the number of positive samples predicted as negative, and FP represents the number of negative samples predicted as positive.

[0084] The effectiveness of training a fault model is measured by predicting image categories and calculating overall classification accuracy.

[0085] Under the ckpt model test, the result images of the top four information richest regions for the upper pull rod malfunction and the relief valve pull rod loss in the embodiment of the present invention are as follows: Figure 6 and Figure 7 As shown.

[0086] In the embodiments of the present invention, the loaded image refers to a dataset of railway freight cars.

[0087] The embodiments of the present invention address the imbalance between faulty and normal samples in different types of faults in railway freight car datasets by replacing the network loss function with Focal Loss, thereby improving the overall classification accuracy.

[0088] This invention can quickly identify fault images and assist in achieving fine-grained fault image classification.

[0089] Compared with existing weakly supervised fine-grained identification algorithms that incorporate attention mechanisms, this invention improves the accuracy of fault identification in industry while also enabling faster detection, thus meeting the real-time requirements of industry.

[0090] Compared with the basic network NTS-Net, this invention is an effective fine-grained image fault identification algorithm that can effectively improve the overall fault classification accuracy while meeting the actual operation scenarios and application needs in industry.

[0091] The above description is only a preferred embodiment of the present invention. It should be noted that for those skilled in the art, several improvements and modifications can be made without departing from the principle of the present invention, and these improvements and modifications should also be considered within the scope of protection of the present invention.

Claims

1. A fine-grained image fault recognition method based on an attention mechanism, characterized in that, Includes the following steps: Step 1: After preprocessing, the collected samples are input into a pre-trained and optimized Navigator network for feature extraction, resulting in a series of rectangular regions. These rectangular regions are then input into the STAN network for weight redistribution. The K regions with the most information after weight redistribution are mapped from the original input samples to obtain K mapped regions. The specific training steps for the Navigator network include: Step A1: The preprocessed image samples are input into the Navigator network for feature extraction to obtain a series of rectangular regions. Then, the series of rectangular regions are input into the STAN network for weight redistribution. The M regions with the highest information richness in the redistributed regions are mapped from the fault samples to obtain the M most discriminative regions. Step A2: Input the M most discriminative regions into the Teacher network, obtain the confidence scores of the M regions through the Teacher network, and then calculate the loss value of the confidence scores of the M regions using the loss function in the Teacher network. Continuously reduce the loss value to optimize the Teacher network. For data samples with a highly imbalanced relationship between faulty and normal samples, the loss function in the Teacher network is Focal when calculating the loss value of the confidence scores of the M regions. The loss function is used, and then the optimized Teacher network is used to re-obtain the confidence scores of the M regions. Next, the STAN network is used to redistribute the weights of the optimized M regions, resulting in the weighted confidence scores of the M regions. These weighted confidence scores are then fed back to the Navigator network. The loss function in the Navigator network is used to calculate the loss value that aligns the information richness of the M most discriminative regions with the confidence scores of the weighted M regions. The Navigator network is trained by gradually decreasing the loss value until the confidence ranking of the weighted M regions matches the information richness ranking of the M most discriminative regions. For data samples with a highly imbalanced mix of faulty and normal samples, the Focal Loss loss function is used in the Navigator network when calculating the loss value that aligns the information richness of the M most discriminative regions with the confidence scores of the weighted M regions. Step 2: Input the K mapping regions into the Scrutinizer network for fine-grained classification and identification to obtain the category of the sample. For data samples with a highly imbalanced relationship between faulty and normal samples, the loss function in the Scrutinizer network is the Focal Loss function when performing fine-grained classification and identification.

2. The image fault fine-grained recognition method based on attention mechanism according to claim 1, characterized in that, The STAN network includes an input terminal, a Conv layer, a BatchNorm layer, a SoftPool layer, a Sigomid layer, an adder, a multiplier, and an output terminal. The input terminals are connected to the Conv layer, the SoftPool layer, and the adder, respectively. The adder is connected to the BatchNorm layer and the multiplier, respectively. The multiplier is connected to the Sigomid layer and the output terminal, respectively. The Conv layer is connected to the BatchNorm layer. The SoftPool layer is connected to another Conv layer, and the other Conv layer is connected to the Sigomid layer.

3. The image fault fine-grained recognition method based on attention mechanism according to claim 1, characterized in that, The formula for the Focal Loss function is as follows: in, This represents the probability that a sample belongs to the correct category. This represents the probability that the predicted sample belongs to the category of 1. Indicates a label, It is the focusing parameter adjusted using cross-validation, which is also the modulation factor.

4. The image fault fine-grained recognition method based on attention mechanism according to claim 1, characterized in that, The classification function used in step 2 is the softmax function.

Images