Image recognition method and device based on attention mechanism, equipment and storage medium
By applying an image recognition method based on attention mechanism in railway truck fault detection, the problem of defective targets occupying a small proportion in the image and not easy to identify is solved, and higher target detection accuracy and stability are achieved.
Patent Information
- Application Number
- CN202411315760.1
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-09-20
- Publication Date
- 2025-02-14
AI Technical Summary
In the detection of railway truck faults, defect targets account for a small proportion of images and are not easy to identify.
Using an image recognition method based on attention mechanism, through multiple convolution and pooling of the backbone network and the head network, combined with the first and second attention mechanism modules, multi-scale features are acquired and weighted adjustments are performed, and the image to be processed is finally processed. Make identification.
It improves the accuracy and stability of object detection, enhances attention to important information in the image, and significantly improves the detection capabilities of small and dense targets.
Smart Images
Figure CN119445318A_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the field of image recognition technology, and in particular to an image recognition method, device, equipment and storage medium based on an attention mechanism. Background Art
[0002] With the rapid development of machine vision technology, technologies such as computer vision, target detection, and deep learning have been widely used in the field of image detection.
[0003] Railway freight cars currently mainly use traditional fault detection methods, which rely on manual inspection and empirical judgment. This method is time-consuming and prone to missed and misjudgment, and work efficiency cannot be improved.
[0004] With the rapid development of the truck transportation industry, improving operational efficiency and safety has become a necessary requirement, and it has become a trend to improve the efficiency and accuracy of truck fault detection through new technologies.
[0005] However, compared with other fields, railway train fault monitoring faces the problem of small proportion of defective targets in the image and difficult to identify in the application of deep learning.
[0006] Based on this, it is necessary to develop and design an image recognition method based on the attention mechanism. Summary of the invention
[0007] The embodiments of the present invention provide an image recognition method, apparatus, device and storage medium based on an attention mechanism, which are used to solve the problem in the prior art that defective targets account for a small proportion in an image and are not easy to identify.
[0008] In a first aspect, an embodiment of the present invention provides an image recognition method based on an attention mechanism, comprising:
[0009] Get the image to be processed;
[0010] The image to be processed is input into the backbone network, and multiple intermediate features are obtained through multiple convolutions and poolings of the backbone network, wherein the third module of the backbone network is a first attention mechanism module, and the first attention mechanism module obtains features containing multi-scale information by convolving the obtained features with different convolution kernels and fusing them and then sending them to the activation function;
[0011] Inputting the multiple intermediate features into the head network, weighting and adjusting the multiple intermediate features through multiple convolutions and pooling of the head network to obtain multiple key features, wherein the output module of the head network is a second attention mechanism module, and the second attention mechanism module obtains the key features by multiplying the global information of the channel and the global information of the spatial dimension;
[0012] The image to be processed is identified according to the multiple key features.
[0013] In a possible implementation, the first attention mechanism module obtains features containing multi-scale information by convolving the obtained features with different convolution kernels, fusing them, and then sending them to the activation function, including:
[0014] The original input is convolved using two different convolution kernels to obtain two feature maps, where each feature map corresponds to a convolution kernel;
[0015] Performing element-by-element summation on the two feature maps to obtain a first new feature map;
[0016] Performing pooling processing on the first new feature map by using global average pooling to obtain a second new feature map;
[0017] Processing the second new feature map through a fully connected network to obtain a third new feature map;
[0018] Adaptively select multiple information space scales using cross-channel soft attention;
[0019] The multiple information space scales are fused to obtain the feature containing the multi-scale information.
[0020] In a possible implementation manner, the adopting global average pooling to perform pooling processing on the first new feature map to obtain the second new feature map includes:
[0021] The first new feature map is pooled according to a first formula to obtain a second new feature map, wherein the first formula is:
[0022]
[0023] In the formula, s c is the characteristic of the c channel, U c (i, j) is the element in the i-th row and j-th column of the c-channel, H is the total number of rows of the c-channel, and W is the total number of columns of the c-channel.
[0024] In one possible implementation, the activation function of the fully connected network is:
[0025]
[0026] Where Sparsemax(x) is the activation function and x is the variable.
[0027] In one possible implementation, the second attention mechanism module includes a channel attention submodule and a spatial attention submodule;
[0028] The channel attention submodule obtains the channel weight according to the pooling result of the input feature map, and performs channel weighting on the input feature map according to the channel weight to obtain a weighted feature map;
[0029] The spatial attention submodule performs splicing and convolution transformation according to two different pooling results of the weighted feature map to obtain a spatial attention map, and implements spatial weighting according to the spatial attention map and the weighted feature map.
[0030] In a possible implementation, the channel attention submodule obtains a channel weight according to a pooling result of an input feature map, and performs channel weighting on the input feature map according to the channel weight to obtain a weighted feature map, including:
[0031] Performing global average pooling and global maximum pooling on the input feature map to obtain a first pooling result and a second pooling result;
[0032] Transforming the first pooling result and the second pooling result by a shared multi-layer perceptron to obtain a weight of each channel;
[0033] The weight of each channel is multiplied by the input feature map to obtain a weighted feature map.
[0034] In a possible implementation, the spatial attention submodule performs concatenation and convolution transformation according to two different pooling results of the weighted feature map to obtain a spatial attention map, and implements spatial weighting according to the spatial attention map and the weighted feature map, including:
[0035] Performing average pooling and maximum pooling on the weighted feature map to obtain a third pooling result and a fourth pooling result;
[0036] The third pooling result and the fourth pooling result are concatenated together, and transformed using a convolution layer to generate a two-dimensional spatial attention map, wherein the convolution layer includes two processes: depth convolution and point-by-point convolution;
[0037] The two-dimensional spatial attention map is multiplied by the weighted feature map to obtain a spatial weighted map.
[0038] In a second aspect, an embodiment of the present invention provides an image recognition device based on an attention mechanism, which is used to implement the image recognition method based on the attention mechanism as described in the first aspect or any possible implementation of the first aspect, and the image recognition device based on the attention mechanism includes:
[0039] An image acquisition module, used for acquiring an image to be processed;
[0040] An intermediate feature extraction module is used to input the image to be processed into the backbone network, and obtain multiple intermediate features through multiple convolutions and pooling of the backbone network, wherein the third module of the backbone network is a first attention mechanism module, and the first attention mechanism module obtains features containing multi-scale information by convolving the obtained features with different convolution kernels and fusing them and then sending them to the activation function;
[0041] A key feature extraction module, used for inputting the multiple intermediate features into the head network, weighting and adjusting the multiple intermediate features through multiple convolutions and pooling of the head network, and obtaining multiple key features, wherein the output module of the head network is a second attention mechanism module, and the second attention mechanism module obtains the key features by multiplying the global information of the channel and the global information of the spatial dimension;
[0042] as well as,
[0043] An image recognition module is used to recognize the image to be processed according to the multiple key features.
[0044] In a third aspect, an embodiment of the present invention provides an electronic device, comprising a memory and a processor, wherein the memory stores a computer program that can be run on the processor, and when the processor executes the computer program, it implements the steps of the method described in the first aspect or any possible implementation method of the first aspect.
[0045] In a fourth aspect, an embodiment of the present invention provides a computer-readable storage medium, wherein the computer-readable storage medium stores a computer program, and when the computer program is executed by a processor, the steps of the method described in the first aspect or any possible implementation of the first aspect are implemented.
[0046] Compared with the prior art, the embodiments of the present invention have the following beneficial effects:
[0047] The present invention is based on the image recognition method implementation method of the attention mechanism, which first obtains the image to be processed; then the image to be processed is input into the backbone network, and multiple intermediate features are obtained through multiple convolutions and pooling of the backbone network, wherein the third module of the backbone network is the first attention mechanism module, and the first attention mechanism module obtains features containing multi-scale information by convolving the obtained features with different convolution kernels and fusing them into the activation function; then the multiple intermediate features are input into the head network, and the multiple intermediate features are weighted and adjusted through multiple convolutions and pooling of the head network to obtain multiple key features, wherein the output module of the head network is the second attention mechanism module, and the second attention mechanism module obtains key features by multiplying the global information of the channel and the global information of the spatial dimension; finally, the image to be processed is identified according to the multiple key features. After improving and optimizing the head part and the backbone part, the implementation method of the present invention improves the accuracy and stability of target detection compared with the baseline target detection algorithm. After improving the attention mechanism in the head and the backbone part, the algorithm model pays more attention to the important information in the image, so it performs better in terms of accuracy. BRIEF DESCRIPTION OF THE DRAWINGS
[0048] In order to more clearly illustrate the technical solutions in the embodiments of the present invention, the drawings required for use in the embodiments or the description of the prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying creative labor.
[0049] Figure 1 is a flowchart of an image recognition method based on an attention mechanism provided by an embodiment of the present invention;
[0050] Figure 2 It is a schematic diagram of the improvement of the main part provided by the embodiment of the present invention;
[0051] Figure 3 is a functional block diagram of a first attention mechanism module provided by an embodiment of the present invention;
[0052] Figure 4 It is a schematic diagram of comparing the similarity of two curves using the sampling method provided in an embodiment of the present invention;
[0053] Figure 5 is a functional block diagram of a second attention mechanism module provided by an embodiment of the present invention;
[0054] Figure 6 is a functional block diagram of an electronic device provided by an embodiment of the present invention;
[0055] Figure 7 It is a comparison diagram of the image recognition method of the present invention provided by the embodiment of the present invention and YOLOv5 and YOLOv7. DETAILED DESCRIPTION
[0056] In the following description, specific details such as specific system structures and technologies are provided for the purpose of illustration rather than limitation so as to provide a thorough understanding of the embodiments of the present invention. However, it should be clear to those skilled in the art that the present invention may be implemented in other embodiments without these specific details. In other cases, detailed descriptions of well-known systems, devices, and methods are omitted so as not to obscure the description of the present invention with unnecessary details.
[0057] In order to make the purpose, technical solutions and advantages of the present invention more clear, the following will be described through specific implementation methods in conjunction with the accompanying drawings.
[0058] The following is a detailed description of an embodiment of the present invention. This example is implemented based on the technical solution of the present invention, and provides a detailed implementation method and a specific operation process, but the protection scope of the present invention is not limited to the following embodiments.
[0059] Figure 1 A flowchart of an image recognition method based on an attention mechanism provided in an embodiment of the present invention.
[0060] like Figure 1 As shown, it shows a flowchart of an implementation method of an image recognition method based on an attention mechanism provided by an embodiment of the present invention, which is described in detail as follows:
[0061] In step 101, an image to be processed is obtained.
[0062] In step 102, the image to be processed is input into the backbone network, and multiple intermediate features are obtained through multiple convolutions and pooling of the backbone network. The third module of the backbone network is a first attention mechanism module. The first attention mechanism module obtains features containing multi-scale information by convolving the obtained features with different convolution kernels and fusing them into the activation function.
[0063] In some implementations, the first attention mechanism module obtains features containing multi-scale information by convolving the obtained features with different convolution kernels, fusing them, and then feeding them into the activation function, including:
[0064] The original input is convolved using two different convolution kernels to obtain two feature maps, where each feature map corresponds to a convolution kernel;
[0065] Performing element-by-element summation on the two feature maps to obtain a first new feature map;
[0066] Performing pooling processing on the first new feature map by using global average pooling to obtain a second new feature map;
[0067] Processing the second new feature map through a fully connected network to obtain a third new feature map;
[0068] Adaptively select multiple information space scales using cross-channel soft attention;
[0069] The multiple information space scales are fused to obtain the feature containing the multi-scale information.
[0070] In some implementations, the using global average pooling to perform pooling processing on the first new feature map to obtain the second new feature map includes:
[0071] The first new feature map is pooled according to a first formula to obtain a second new feature map, wherein the first formula is:
[0072]
[0073] In the formula, s c is the characteristic of the c channel, U c (i, j) is the element in the i-th row and j-th column of the c-channel, H is the total number of rows of the c-channel, and W is the total number of columns of the c-channel.
[0074] In some embodiments, the activation function of the fully connected network is:
[0075]
[0076] Where Sparsemax(x) is the activation function and x is the variable.
[0077] For example, the backbone of the target detection algorithm is mainly responsible for extracting useful feature information from the input image for use by the subsequent detection head. In the research and analysis of algorithm optimization and improvement, there are small targets and densely covered small targets in the data set of railway freight car fault detection. By optimizing and improving the backbone, the algorithm model can pay more attention to the detection and extraction of feature information in the image at the first time of inputting image information, thereby improving the algorithm model's ability to detect small targets and densely covered targets. The backbone network structure is as follows: Figure 2 shown.
[0078] SKAttention (Selected Kernel Attention Mechanism) is a deep learning model component based on the attention mechanism, specially designed to improve the adaptability and capture ability of convolutional neural networks for multi-scale features. The model can dynamically select appropriate convolution kernels according to the different content and feature sizes of the input image to extract richer and more effective feature representations. The SKAttention mechanism can automatically learn the importance of the input feature map in the channel dimension to adaptively weight the feature map.
[0079] The core idea of SKAttention is to use parallel computing of convolution operations of different sizes to extract features of different scales, and learn the weight of each channel through the selector network to adaptively weight features of different scales.
[0080] As shown in the figure, in the split stage, for the original input C, a 3×3 convolution kernel and a 5×5 convolution kernel are used to introduce receptive fields of different scales. The 5×5 convolution kernel is replaced by a 3×3 dilated convolution (hole convolution) to improve efficiency.
[0081] In the fusion stage, the two feature maps are summed element by element to obtain a new feature map U.
[0082]
[0083] U uses GAP to get S:
[0084]
[0085] In the formula, s c is the characteristic of the c channel, U c (i, j) is the element in the i-th row and j-th column of the c-channel, H is the total number of rows of the c-channel, and W is the total number of columns of the c-channel.
[0086] S plus a full connection to get Z:
[0087] z=F fc (s) = δ(B(Ws))
[0088] In the selection stage, under the guidance of the compact feature descriptor z, cross-channel soft attention is used to adaptively select different information space scales. The feature map V is obtained by the attention weights on various kernels, which fuses the information of multiple receptive fields and has multi-scale information.
[0089] The activation function used by SKAttention is Softmax, which maps the input to a probability distribution in the interval (0,1). Most of the output vectors will have a small non-zero value. When detecting small targets and multi-scale features, the model may focus too much on irrelevant features or background noise and reduce the ability to detect small targets.
[0090] Therefore, in view of the above situation, this paper chooses to apply Sparsemax to SKAttention to further improve the small target detection ability of the model. Sparsemax is a sparse activation function, similar to Softmax, both of which map the input to a probability distribution in the (0,1) interval, but the output of Sparsemax is more sparse, and more output vectors are set to zero, so it has greater advantages in processing multi-label classification and attention mechanisms. Figure 3 Schematic diagram of the SKAttention mechanism based on the Sparsemax activation function.
[0091] The Sparsemax activation function is:
[0092]
[0093] Where Sparsemax(x) is the activation function and x is the variable.
[0094] The implementation mode of the present invention chooses to add the improved SKAttention to the backbone part (Backbone) of the algorithm model. By analyzing the characteristics and advantages of the SKAttention attention mechanism, this article uses SKAttention to improve the CBS module of the third layer in the backbone part (Backbone) at the front position of the backbone part without increasing the complexity of the algorithm model. In the initial stage of the network structure, the use of a larger convolution kernel is helpful to extract the boundary features of the target and improve the performance of the model in locating and segmenting the target. A smaller step size can maintain the size of the input feature map and reduce information loss. It can avoid the problems of information missing and resolution loss, and is conducive to retaining details and local features.
[0095] In step 103, the multiple intermediate features are input into the head network, and the multiple intermediate features are weighted and adjusted through multiple convolutions and pooling of the head network to obtain multiple key features, wherein the output module of the head network is a second attention mechanism module, and the second attention mechanism module obtains the key features by multiplying the global information of the channel and the global information of the spatial dimension.
[0096] In some embodiments, the second attention mechanism module includes a channel attention submodule and a spatial attention submodule;
[0097] The channel attention submodule obtains the channel weight according to the pooling result of the input feature map, and performs channel weighting on the input feature map according to the channel weight to obtain a weighted feature map;
[0098] The spatial attention submodule performs splicing and convolution transformation according to two different pooling results of the weighted feature map to obtain a spatial attention map, and implements spatial weighting according to the spatial attention map and the weighted feature map.
[0099] In some implementations, the channel attention submodule obtains a channel weight according to a pooling result of an input feature map, and performs channel weighting on the input feature map according to the channel weight to obtain a weighted feature map, including:
[0100] Performing global average pooling and global maximum pooling on the input feature map to obtain a first pooling result and a second pooling result;
[0101] Transforming the first pooling result and the second pooling result by a shared multi-layer perceptron to obtain a weight of each channel;
[0102] The weight of each channel is multiplied by the input feature map to obtain a weighted feature map.
[0103] In some embodiments, the spatial attention submodule performs concatenation and convolution transformation according to two different pooling results of the weighted feature map to obtain a spatial attention map, and implements spatial weighting according to the spatial attention map and the weighted feature map, including:
[0104] Performing average pooling and maximum pooling on the weighted feature map to obtain a third pooling result and a fourth pooling result;
[0105] The third pooling result and the fourth pooling result are concatenated together, and transformed using a convolution layer to generate a two-dimensional spatial attention map, wherein the convolution layer includes two processes: depth convolution and point-by-point convolution;
[0106] Multiply the two-dimensional spatial attention map with the weighted feature map to obtain a spatial weighted map.
[0107] For example, the head part of the target detection algorithm is responsible for converting the feature map into the detection result, including the location of the target (detection box coordinates), category and confidence. Optimizing and improving the head part of the algorithm model can make the algorithm model pay more attention to the generation of the detection box, the classification of categories and the setting of confidence before output. It plays an important role in improving the detection accuracy, speed and generalization ability of the algorithm model.
[0108] Attention Mechanism is a commonly used technique in deep learning, which is used for the model to selectively focus on and assign weights to specific parts of the input data. In computer vision tasks, the attention mechanism can help the model select and focus on areas of interest in the image. In order to further improve the accuracy of target detection and reduce the cases of missed detection and false detection of the algorithm model, the generalization ability of the model is improved. In the embodiment of the present invention, the attention mechanism is added to the head part, and the network structure of the algorithm is improved, in order to improve the accuracy of the model, better adapt to small targets and occlusions, improve the generalization ability of the model, and can reduce the amount of calculation and the number of parameters.
[0109] CBAM is an attention mechanism module for convolutional neural networks, which aims to enhance the model's ability to model and distinguish features. It combines two sub-modules, channel attention and spatial attention, to weight and adjust features in the network, which can significantly improve the accuracy of image classification and object detection.
[0110] This module consists of two submodules: channel attention submodule and spatial attention submodule. The channel attention module in CBAM first captures the global information of the channel by performing global average pooling and global maximum pooling on the input feature map respectively. Then, the two pooling results are transformed through a shared multi-layer perceptron (MLP) to obtain the weight of each channel. Finally, these weights are multiplied with the original feature map to achieve channel weighting.
[0111] The spatial attention module performs average pooling and maximum pooling on the output of the channel attention module to obtain global information in the spatial dimension. The two pooled results are then concatenated and transformed through a convolutional layer to generate a two-dimensional spatial attention map. Finally, this spatial attention map is multiplied with the feature map that has been weighted by the channel attention to achieve spatial weighting.
[0112] The channel attention and spatial attention mechanisms in CBAM can improve the model's sensitivity to channel and spatial information respectively, helping the model to better focus on important features and pixels. This method of combining the two attention mechanisms can further improve the performance of the model, making it more accurate and robust when dealing with challenging tasks.
[0113] Adding CBAM to the algorithm model will inevitably increase the complexity of the network structure and the number of parameters, reduce the computational efficiency of the model, and while CBAM enhances key features, it will weaken the algorithm model's ability to detect small targets. In view of the above situation, this paper chooses to apply depthwise separable convolution in CBAM. Depthwise separable convolution decomposes traditional convolution into two processes: depthwise convolution and pointwise convolution. Compared with traditional convolution, it can significantly reduce the amount of calculation and model complexity. When processing small target detection, it integrates the features of different channels through pointwise convolution 1×1, which increases the flexibility of feature extraction and helps capture more detailed information.
[0114] Combining the characteristics of CBAM attention mechanism to improve the channel and spatial sensitivity of the model, this paper uses CBAM based on depthwise separable convolution to improve the ELAN-W module before the three outputs of the head part (Head).
[0115] First, before the bounding box coordinates are output, the output structure is improved so that the regression of the bounding box coordinates can focus more on the distinguishability and accuracy of the detection box, thereby improving the detection rate of small and dense targets.
[0116] Second, to improve the target category output of the target detection algorithm, the CBAM attention mechanism can adaptively adjust the weight of the feature through the channel attention mechanism, so that the algorithm model pays more attention to the feature channel with target information;
[0117] Finally, the confidence output of the algorithm model is improved, and the spatial attention mechanism in the CBAM module can adjust the spatial weight of the feature map, emphasize the target area and suppress background interference, thereby improving the target detection ability. Figure 4 Schematic diagram of adding CBAM based on depthwise separable convolution.
[0118] In step 104, the image to be processed is identified according to the multiple key features.
[0119] Exemplarily, by inputting a plurality of key features into a classifier, for example, a classifier based on an artificial neural network, it is possible to identify whether a target defect exists in an image.
[0120] Comparison and analysis of experimental results:
[0121] In the experimental comparison, the implementation method of the present invention selected YOLOv5 and baseline YOLOv7 as comparison models. YOLOv5 is an efficient and accurate target detection algorithm and is the fifth version of the YOLO series. The dataset images are input into the three algorithm models at a resolution of 640×640, the epochs hyperparameter is set to 200, the batch-size hyperparameter is set to 8, and the three algorithm model training processes are all accelerated by GPU training. In order to speed up and improve the training process, the official pre-trained model is used in the comparison experiment training process, where the YOLOv5 algorithm model uses the YOLOv5m pre-trained model, and the baseline YOLOv7 and the optimized and improved algorithm model both use the YOLOv7 official pre-trained model.
[0122] The comparison results of precision, recall, and mAP@.5:.95 are shown in the table.
[0123] Table 1 Comparison of evaluation indicators
[0124] Model Accuracy Recall mAP@.5:.95 YOLOv5 51.4％ 99.8％ 86.96％ YOLOv7 90.4％ 99％ 76.61％ Ours 99.67％ 99.8％ 90.6％
[0125] Figure 7 The following figure shows the changes in the indicators of the three models during the training process. Combined with the chart and the experimental data, the improved model in this paper has improved the accuracy by 9.27% compared with the baseline YOLOv7. The comparison of the training changes in accuracy in the figure below shows that the improved model is more stable.
[0126] In terms of recall rate, all three models demonstrate a high ability to identify positive sample instances. The optimized and improved algorithm model in this paper performs more stably than the baseline YOLOv7 during training.
[0127] In terms of average precision, the mAP index of the optimized and improved algorithm model in this paper is improved by 3.64% compared with the better performing YOLOv5 algorithm model, and by 13.99% compared with the baseline YOLOv7 algorithm model.
[0128] Based on the above comparison and analysis, on the railway freight car fault detection dataset produced in this paper, after improving and optimizing the head part and the trunk part, the accuracy and stability of target detection are improved compared with the baseline target detection algorithm and YOLOv5 algorithm. After improving the attention mechanism in the head and trunk parts, the algorithm model pays more attention to the important information in the image, so it performs better in terms of accuracy. After improving the normalization method, the algorithm model converges faster and performs more stably during training.
[0129] The present invention is based on the image recognition method implementation method of the attention mechanism, which first obtains the image to be processed; then the image to be processed is input into the backbone network, and multiple intermediate features are obtained through multiple convolutions and pooling of the backbone network, wherein the third module of the backbone network is the first attention mechanism module, and the first attention mechanism module obtains features containing multi-scale information by convolving the obtained features with different convolution kernels and fusing them into the activation function; then the multiple intermediate features are input into the head network, and the multiple intermediate features are weighted and adjusted through multiple convolutions and pooling of the head network to obtain multiple key features, wherein the output module of the head network is the second attention mechanism module, and the second attention mechanism module obtains key features by multiplying the global information of the channel and the global information of the spatial dimension; finally, the image to be processed is identified according to the multiple key features. After improving and optimizing the head part and the backbone part, the implementation method of the present invention improves the accuracy and stability of target detection compared with the baseline target detection algorithm. After improving the attention mechanism in the head and the backbone part, the algorithm model pays more attention to the important information in the image, so it performs better in terms of accuracy.
[0130] It should be understood that the size of the serial numbers of the steps in the above implementation does not mean the order of execution. The execution order of each process should be determined by its function and internal logic, and should not constitute any limitation on the implementation process of the implementation method of the present invention.
[0131] The following is an implementation of the device of the present invention. For details not described in detail, reference may be made to the corresponding method implementation described above.
[0132] Figure 5 is a functional block diagram of an image recognition device based on an attention mechanism provided by an embodiment of the present invention, referring to Figure 5 The image recognition device based on the attention mechanism includes: an image acquisition module 501, an intermediate feature extraction module 502, a key feature extraction module 503 and an image recognition module 504, wherein:
[0133] An image acquisition module 501 is used to acquire an image to be processed;
[0134] The intermediate feature extraction module 502 is used to input the image to be processed into the backbone network, and obtain multiple intermediate features through multiple convolutions and pooling of the backbone network, wherein the third module of the backbone network is a first attention mechanism module, and the first attention mechanism module obtains features containing multi-scale information by convolving the obtained features with different convolution kernels and fusing them and then sending them to the activation function;
[0135] A key feature extraction module 503 is used to input the multiple intermediate features into the head network, and weight and adjust the multiple intermediate features through multiple convolutions and pooling of the head network to obtain multiple key features, wherein the output module of the head network is a second attention mechanism module, and the second attention mechanism module obtains the key features by multiplying the global information of the channel and the global information of the spatial dimension;
[0136] as well as,
[0137] The image recognition module 504 is used to recognize the image to be processed according to the multiple key features.
[0138] Figure 6 is a functional block diagram of an electronic device provided by an embodiment of the present invention. Figure 6 As shown, the electronic device 6 of this embodiment includes: a processor 600 and a memory 601, wherein the memory 601 stores a computer program 602 that can be run on the processor 600. When the processor 600 executes the computer program 602, the steps in the above-mentioned image recognition method based on the attention mechanism and the embodiment are implemented, for example Figure 1 Steps 101 to 104 are shown.
[0139] Exemplarily, the computer program 602 may be divided into one or more modules / units, and the one or more modules / units are stored in the memory 601 and executed by the processor 600 to implement the present invention.
[0140] The electronic device 6 may be a computing device such as a desktop computer, a notebook, a PDA, or a cloud server. The electronic device 6 may include, but is not limited to, a processor 600 and a memory 601. Those skilled in the art will appreciate that Figure 6 It is only an example of the electronic device 6 and does not constitute a limitation of the electronic device 6. It may include more or fewer components than shown in the figure, or a combination of certain components, or different components. For example, the electronic device 6 may also include input and output devices, network access devices, buses, etc.
[0141] The processor 600 may be a central processing unit (CPU), or other general-purpose processors, digital signal processors (DSP), application-specific integrated circuits (ASIC), field-programmable gate arrays (FPGA) or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, etc. A general-purpose processor may be a microprocessor or any conventional processor, etc.
[0142] The memory 601 may be an internal storage unit of the electronic device 6, such as a hard disk or memory of the electronic device 6. The memory 601 may also be an external storage device of the electronic device 6, such as a plug-in hard disk, a smart media card (SMC), a secure digital (SD) card, a flash card, etc. equipped on the electronic device 6. Further, the memory 601 may also include both an internal storage unit and an external storage device of the electronic device 6. The memory 601 is used to store the computer program 602 and other programs and data required by the electronic device 6. The memory 601 may also be used to temporarily store data that has been output or is to be output.
[0143] Those skilled in the art can clearly understand that for the convenience and simplicity of description, only the division of the above-mentioned functional units and modules is used as an example for illustration. In actual applications, the above-mentioned function allocation can be completed by different functional units and modules as needed, that is, the internal structure of the device can be divided into different functional units or modules to complete all or part of the functions described above. The functional units and modules in the implementation method can be integrated into a processing unit, or each unit can exist physically separately, or two or more units can be integrated into one unit. The above-mentioned integrated unit can be implemented in the form of hardware or in the form of software functional units. In addition, the specific names of the functional units and modules are only for the convenience of distinguishing each other, and are not used to limit the scope of protection of this application. The specific working process of the units and modules in the above-mentioned system can refer to the corresponding process in the aforementioned method implementation method, which will not be repeated here.
[0144] In the above-mentioned embodiments, the description of each embodiment has its own emphasis. For parts that are not described or recorded in detail in a certain embodiment, reference can be made to the relevant descriptions of other embodiments.
[0145] Those skilled in the art will appreciate that the units and algorithm steps of each example described in conjunction with the embodiments disclosed herein can be implemented in electronic hardware, or a combination of computer software and electronic hardware. Whether these functions are performed in hardware or software depends on the specific application and design constraints of the technical solution. Professional and technical personnel can use different methods to implement the described functions for each specific application, but such implementation should not be considered to be beyond the scope of the present invention.
[0146] In the embodiments provided by the present invention, it should be understood that the disclosed devices / electronic devices and methods can be implemented in other ways. For example, the device / electronic device embodiments described above are only schematic. For example, the division of the modules or units is only a logical function division. There may be other division methods in actual implementation, such as multiple units or components can be combined or integrated into another system, or some features can be ignored or not executed. Another point is that the mutual coupling or direct coupling or communication connection shown or discussed can be through some interfaces, indirect coupling or communication connection of devices or units, which can be electrical, mechanical or other forms.
[0147] The units described as separate components may or may not be physically separated, and the components shown as units may or may not be physical units, that is, they may be located in one place or distributed on multiple network units. Some or all of the units may be selected according to actual needs to achieve the purpose of the present embodiment.
[0148] In addition, each functional unit in each embodiment of the present invention may be integrated into one processing unit, or each unit may exist physically separately, or two or more units may be integrated into one unit. The above-mentioned integrated unit may be implemented in the form of hardware or in the form of software functional units.
[0149] If the integrated module / unit is implemented in the form of a software functional unit and sold or used as an independent product, it can be stored in a computer-readable storage medium. Based on this understanding, the present invention implements all or part of the processes in the above-mentioned implementation method, and can also be completed by instructing the relevant hardware through a computer program. The computer program can be stored in a computer-readable storage medium. When the computer program is executed by the processor, it can implement the steps of the above-mentioned various methods and device implementation methods. Among them, the computer program includes computer program code, and the computer program code can be in source code form, object code form, executable file or some intermediate form. The computer-readable medium may include: any entity or device that can carry the computer program code, recording medium, U disk, mobile hard disk, disk, optical disk, computer memory, read-only memory (ROM), random access memory (RAM), electrical carrier signal, telecommunication signal and software distribution medium.
[0150] The above-described embodiments are only used to illustrate the technical solutions of the present invention, rather than to limit the same. Although the present invention has been described in detail with reference to the aforementioned embodiments, a person skilled in the art should understand that the technical solutions described in the aforementioned embodiments may still be modified, or some of the technical features may be replaced by equivalents. Such modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the various embodiments of the present invention, and should all be included in the protection scope of the present invention.
Claims
1. An image recognition method based on attention mechanism, characterized in that: include: Get the image to be processed; The image to be processed is input into the backbone network, and multiple intermediate features are obtained through multiple convolutions and poolings of the backbone network, wherein the third module of the backbone network is a first attention mechanism module, and the first attention mechanism module obtains features containing multi-scale information by convolving the obtained features with different convolution kernels and fusing them and then sending them to the activation function; Inputting the multiple intermediate features into the head network, weighting and adjusting the multiple intermediate features through multiple convolutions and pooling of the head network to obtain multiple key features, wherein the output module of the head network is a second attention mechanism module, and the second attention mechanism module obtains the key features by multiplying the global information of the channel and the global information of the spatial dimension; The image to be processed is identified according to the multiple key features.
2. The image recognition method based on the attention mechanism according to claim 1, characterized in that: The first attention mechanism module obtains features containing multi-scale information by convolving the obtained features with different convolution kernels and fusing them into the activation function, including: The original input is convolved using two different convolution kernels to obtain two feature maps, where each feature map corresponds to a convolution kernel; Performing element-by-element summation on the two feature maps to obtain a first new feature map; Performing pooling processing on the first new feature map by using global average pooling to obtain a second new feature map; Processing the second new feature map through a fully connected network to obtain a third new feature map; Adaptively select multiple information space scales using cross-channel soft attention; The multiple information space scales are fused to obtain the feature containing the multi-scale information.
3. The image recognition method based on the attention mechanism according to claim 2, characterized in that: The step of performing pooling processing on the first new feature map by using global average pooling to obtain a second new feature map includes: The first new feature map is pooled according to a first formula to obtain a second new feature map, wherein the first formula is: In the formula, s c is the characteristic of the c channel, U c (i, j) is the element in the i-th row and j-th column of the c-channel, H is the total number of rows of the c-channel, and W is the total number of columns of the c-channel.
4. The image recognition method based on the attention mechanism according to claim 2, characterized in that: The activation function of the fully connected network is: Where Sparsemax(x) is the activation function and x is the variable.
5. The image recognition method based on the attention mechanism according to any one of claims 1 to 4, characterized in that: The second attention mechanism module includes a channel attention submodule and a spatial attention submodule; The channel attention submodule obtains the channel weight according to the pooling result of the input feature map, and performs channel weighting on the input feature map according to the channel weight to obtain a weighted feature map; The spatial attention submodule performs splicing and convolution transformation according to two different pooling results of the weighted feature map to obtain a spatial attention map, and implements spatial weighting according to the spatial attention map and the weighted feature map.
6. The image recognition method based on the attention mechanism according to claim 5, characterized in that: The channel attention submodule obtains the channel weight according to the pooling result of the input feature map, and performs channel weighting on the input feature map according to the channel weight to obtain a weighted feature map, including: Performing global average pooling and global maximum pooling on the input feature map to obtain a first pooling result and a second pooling result; Transforming the first pooling result and the second pooling result by a shared multi-layer perceptron to obtain a weight of each channel; The weight of each channel is multiplied by the input feature map to obtain a weighted feature map.
7. The image recognition method based on the attention mechanism according to claim 5, characterized in that: The spatial attention submodule performs concatenation and convolution transformation according to two different pooling results of the weighted feature map to obtain a spatial attention map, and implements spatial weighting according to the spatial attention map and the weighted feature map, including: Performing average pooling and maximum pooling on the weighted feature map to obtain a third pooling result and a fourth pooling result; The third pooling result and the fourth pooling result are concatenated together, and transformed using a convolution layer to generate a two-dimensional spatial attention map, wherein the convolution layer includes two processes: depth convolution and point-by-point convolution; The two-dimensional spatial attention map is multiplied by the weighted feature map to obtain a spatial weighted map.
8. An image recognition device based on an attention mechanism, characterized in that: Used to implement the image recognition method based on the attention mechanism according to any one of claims 1 to 7, the image recognition device based on the attention mechanism comprises: An image acquisition module, used for acquiring an image to be processed; An intermediate feature extraction module is used to input the image to be processed into the backbone network, and obtain multiple intermediate features through multiple convolutions and pooling of the backbone network, wherein the third module of the backbone network is a first attention mechanism module, and the first attention mechanism module obtains features containing multi-scale information by convolving the obtained features with different convolution kernels and fusing them and then sending them to the activation function; A key feature extraction module, used for inputting the multiple intermediate features into the head network, weighting and adjusting the multiple intermediate features through multiple convolutions and pooling of the head network, and obtaining multiple key features, wherein the output module of the head network is a second attention mechanism module, and the second attention mechanism module obtains the key features by multiplying the global information of the channel and the global information of the spatial dimension; as well as, An image recognition module is used to recognize the image to be processed according to the multiple key features.
9. An electronic device comprising a memory and a processor, wherein the memory stores a computer program that can be run on the processor, characterized in that: When the processor executes the computer program, the steps of the method according to any one of claims 1 to 7 are implemented.
10. A computer-readable storage medium storing a computer program, characterized in that: When the computer program is executed by a processor, the steps of the method as claimed in any one of claims 1 to 7 are implemented.