A small target detection method based on enhanced feature extraction
By adding a coordinate attention mechanism and improving the feature pyramid structure to the YOLOv5 network, and using the ShuffleNet v2 backbone network with pruned invalid computations, the problem of low accuracy in small target detection was solved, achieving higher detection accuracy and real-time performance.
Patent Information
- Application Number
- CN202211515877.5
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2022-11-29
- Publication Date
- 2026-02-03
- Estimated Expiration
- 2042-11-29
AI Technical Summary
Existing deep learning algorithms do not perform well in detecting small targets in aerial images, especially because small targets are small in scale, have a small coverage area, and are easily mixed with background information, resulting in low detection accuracy.
Based on the YOLOv5 network, a coordinate attention mechanism is added to improve the feature pyramid structure. ShuffleNet v2 is used as the backbone network and lightweighted to prune invalid computations and enhance feature extraction capabilities.
It improves the accuracy and real-time performance of small target detection by enhancing feature extraction capabilities and reducing model complexity, thereby improving the model's detection performance.
Smart Images
Figure CN115984172B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the fields of image processing and computer vision technology, and in particular to a small target detection method based on enhanced feature extraction. Background Technology
[0002] Object detection has long been an important research direction in the field of computer vision, and it is now widely used in various fields such as medicine, aerospace, industry and agriculture, transportation, and autonomous driving, which has a great impact on improving people's quality of life.
[0003] Traditional object detection algorithms typically employ manual feature extraction and specialized classifiers designed for different objects. These methods suffer from low accuracy, complex design, and poor generalization ability, thus failing to achieve ideal results. With the advent of the big data era and the improvement of hardware computing power, deep learning-based convolutional neural networks (CNNs) have gradually come into focus, providing a significant impetus for the development of object detection. Deep learning-based object detection algorithms can be divided into two-stage detection algorithms based on candidate boxes and one-stage detection algorithms based on regression. Two-stage algorithms first extract candidate regions and then classify the samples using a neural network. Examples of two-stage detection methods include R-CNN, SPP-NET, Fast R-CNN, Faster R-CNN, Mask R-CNN, and Cascade R-CNN. One-stage detection algorithms, on the other hand, directly regress the object's category and location. Examples of one-stage detection methods include the YOLO series and the SSD algorithm. One-stage detection algorithms offer higher real-time performance and have therefore received more widespread attention. Pei et al. proposed an improved SSD aerial image detection algorithm by fusing shallow location features with deep semantic features of the network, but the increased network depth further increased the model's complexity. Xu et al. used deformable convolution to improve aerial image detection accuracy by extracting multi-scale feature information under different receptive fields, but they neglected the contribution of feature maps at different levels to small targets.
[0004] Small targets are small in scale and cover a small area in images, making them easy to mix with background information. Therefore, they are more complex to detect than regular targets. Current deep learning algorithms do not perform well in detecting small targets, such as those in aerial images. Summary of the Invention
[0005] To address the shortcomings of existing technologies in small object detection, this invention proposes a small object detection algorithm based on enhanced feature extraction, building upon the YOLOv5 network. A coordinate attention mechanism is added to the YOLOv5 backbone network to improve the feature pyramid (FPN) structure and increase shallow feature reuse. Simultaneously, lightweighting and pruning are employed to reduce the number of network parameters and computational cost, minimizing the real-time performance loss caused by increased model complexity.
[0006] To achieve the above objectives, this invention proposes a small target detection method based on enhanced feature extraction, comprising:
[0007] (1) Collect images to construct a dataset, and divide the dataset into a training set and a test set;
[0008] (2) Constructing an improved YOLOv5 network model:
[0009] (2.1) The backbone network is replaced by ShuffleNet V2: it adopts a dual convolution kernel processing and connection method, and merges continuous element-wise operations at the same time;
[0010] (2.2) Add coordinate attention mechanism: embed position information into channel attention, integrate global image information, and weaken background information interference;
[0011] (2.3) Improve the feature extraction structure: deepen the feature pyramid FPN, increase the reuse of shallow features, retain richer location information, and enhance the feature extraction capability of small targets. At the same time, due to the characteristics of the target scale distribution in the dataset, the low-resolution target detection layer is deleted to avoid invalid computation.
[0012] (3) The improved YOLOv5 network model is trained using the dataset to obtain a small object detection model.
[0013] Furthermore, step (2.1) of replacing the backbone network with ShuffleNet V2 includes: replacing the convolutional modules of layers 3-8 and the C3 module in the original network with the convolutional modules of ShuffleNet V2, and adopting a channel splitting and channel shuffling strategy.
[0014] Furthermore, the channel segmentation involves dividing the c feature channel input of each unit into two branches, each having c-c' and c' channels, where c' takes the value c / 2; then, the two branches obtained by convolution are connected to maintain the same number of channels.
[0015] The channel shuffling ensures that information can be exchanged between different branches.
[0016] Furthermore, step (2.2) of adding a coordinate attention mechanism includes adding coordinate attention modules after layers 2, 5, 8, and 12 of the backbone network.
[0017] Furthermore, the coordinate attention module is implemented as follows:
[0018] (2.2.1) Coordinate attention embedding: Global pooling is decomposed to obtain one-to-one one-dimensional feature encoding; For input of C×H×W dimension, where C is the number of channels, H is the height and W is the width, each channel is encoded according to the horizontal and vertical coordinates using a pooling kernel of dimension (H,1) or (1,W);
[0019] (2.2.2) Coordinate attention generation: The representation obtained in step (2.2.1) is transformed using a 1×1 convolution kernel transformation function, and the transformation result is decomposed into two separate tensors along the spatial dimension. Finally, the convolution transformation is used to transform the input and the input into tensors with the same number of channels, where the convolution transformation uses the sigmoid activation function.
[0020] (2.2.3) Expand the result in step (2.2.2) as attention weights.
[0021] Furthermore, step (2.3) of improving the feature extraction structure includes: the feature map output by layer 14 in the network structure is upsampled by nearest and then fused with the feature map output by layer 8 by concat; the fused feature map is upsampled by nearest and then fused with the feature map output by layer 5 by concat; the above result is upsampled by nearest again and then fused with the feature map output by layer 2 by concat to obtain the fused feature map; this feature map is then passed through module C3 to finally obtain the output results of each prediction layer; this result includes a newly added high-resolution target detection layer for small targets, which improves the small target detection effect, while deleting the original low-resolution target detection layer to avoid invalid computation.
[0022] Furthermore, training the improved YOLOv5 network model using the training set data includes:
[0023] (3.1) Calculate the corresponding initial anchor box values for the dataset using the Kmeans++ method;
[0024] (3.2) Perform input data preprocessing on the training set data; that is, splice the data by random scaling, random cropping and random arrangement to achieve Mosaic data augmentation;
[0025] (3.3) Input the preprocessed training set data into the improved YOLOV5 network model for training to obtain a small target detection model;
[0026] (3.4) After N iterations, the trained small target detection model is validated and evaluated using the validation set data.
[0027] Furthermore, the evaluation metric uses the mean precision, which is related to precision and recall.
[0028] The accuracy formula is as follows:
[0029]
[0030] The recall rate formula is as follows:
[0031]
[0032] Average precision refers to the area enclosed by the curve formed by recall on the horizontal axis and precision on the vertical axis, as shown in the formula below:
[0033]
[0034] The mean precision is the average precision of all classes in the dataset, and the formula is as follows:
[0035]
[0036] Where: TP represents true positives, FP represents false positives, TN represents true negatives, FN represents false negatives, AP value refers to the area under the PR curve, and N represents the total number of categories detected.
[0037] The beneficial effects of this invention are:
[0038] 1. This invention uses the Kmeans++ clustering method to reset the initial anchor box values, which intuitively enhances the model effect.
[0039] 2. This invention integrates global information of an image by adding a coordinate attention mechanism, which weakens background information interference and enhances feature expression ability; it improves positioning accuracy by deepening the feature pyramid structure to retain rich positional information; and it enhances the feature extraction ability of the model by using a multi-scale fusion strategy to aggregate shallow feature maps with deep feature maps with rich semantic information.
[0040] 3. Since almost all targets in this dataset are small to medium sized, the detection layer for large targets in the original network has almost lost its original function, only increasing the computational load and model complexity. Therefore, this invention removes this detection layer and further uses the ShuffleNet v2 lightweight network to replace the backbone network in the original network, reducing the model size to meet real-time requirements.
[0041] 4. The recognition accuracy of this invention is better than that of the original algorithm on the RSOD dataset, and it has better detection effect. Attached Figure Description
[0042] Figure 1 This is a schematic diagram of the small target detection method based on enhanced feature extraction according to an embodiment of the present invention.
[0043] Figure 2 This is a schematic diagram of the ShuffleNet v2 module structure in an embodiment of the present invention.
[0044] Figure 3 This is a schematic diagram of the coordinate attention module structure in an embodiment of the present invention.
[0045] Figure 4 This is a schematic diagram of the improved feature extraction structure according to an embodiment of the present invention.
[0046] Figure 5 This is a comparison chart of the detection results of the model in this embodiment and the detection results of the original YOLOv5 model. Detailed Implementation
[0047] The present invention will now be further described with reference to the accompanying drawings and embodiments.
[0048] like Figure 1 As shown, this embodiment of the invention provides a small target detection method based on enhanced feature extraction, including the following steps:
[0049] S101. Collect images to construct a dataset, and divide the dataset into a training set and a test set;
[0050] This embodiment takes aerial image target detection as an example. It uses the open-source aerial target detection dataset RSOD and selects two categories, airport and oiltank, as the training set and detection samples. The training set and test set are divided in a 9:1 ratio, and the training data and test data are strictly independent.
[0051] S102. Constructing an improved YOLOv5 network model:
[0052] The overall structure of the improved YOLOv5 network model is shown in Table 1:
[0053] Table 1 Overall Network Structure
[0054]
[0055]
[0056] The "from" column indicates which layer the input comes from, -1 indicates it comes from the previous layer, the "module" column is the module of this layer, and the "arguments" column contains module parameter information, including the number of input channels, kernel size, stride, etc.
[0057] The network structure described above is as follows: In the YOLOv5 model, the backbone network is replaced by Shufflenet v2, a CA attention module is added, the network structure is improved, shallow feature reuse is increased, and invalid detection layers are pruned.
[0058] (1) The backbone network is replaced by ShuffleNet V2: it adopts a double convolution kernel processing and connection method, and merges continuous element-by-element operations at the same time;
[0059] like Figure 2 As shown, ShuffleNet v2 employs channel splitting and channel shuffling strategies. Channel splitting divides the c feature channel input of each unit into two branches, with c-c' and c' channels respectively, where c' is c / 2. The two branches obtained from convolution are then concatenated to maintain the same number of channels. The channel shuffling strategy ensures information exchange between different branches. Furthermore, to address the issue of element-wise operations consuming significant time, ReLU and depthwise convolution operations are placed in the same branch, and the three consecutive element-wise operations—concat, channel shuffling, and channel splitting—are combined into a single element-wise operation.
[0060] (2) Add coordinate attention mechanism: embed position information into channel attention, integrate global image information, and weaken background information interference;
[0061] With the Focus module as layer 0, coordinate attention modules are added after layers 2, 5, 8, and 12 of the backbone network;
[0062] like Figure 3 As shown, the implementation of the CA attention module is as follows:
[0063] (2.1) Coordinate Attention Embedding. Global pooling is decomposed to obtain one-to-one encoding of one-dimensional features. For inputs of C×H×W dimensions, where C is the number of channels, H is the height, and W is the width, each channel is encoded according to the horizontal and vertical coordinates using a pooling kernel of dimension (H,1) or (1,W).
[0064] (2.2) Coordinate attention generation. The representation obtained in step 1) is transformed using a 1×1 convolution kernel transformation function, and the transformation result is decomposed into two separate tensors along the spatial dimension. Finally, the two tensors are transformed into tensors with the same number of channels as the input using convolution transformation, where the convolution transformation uses the sigmoid activation function.
[0065] (2.3) The results in (2.2) are expanded and used as attention weights.
[0066] (3) Improve the feature extraction structure: deepen the feature pyramid FPN, increase the reuse of shallow features, retain richer location information, enhance the feature extraction capability of small targets, and at the same time, due to the characteristics of the target scale distribution in the dataset, delete the low-resolution target detection layer to avoid invalid computation;
[0067] like Figure 4 As shown, shallow feature reuse is implemented as follows: the feature map output from layer 14 in the network structure is upsampled by the nearest feature map and then fused with the feature map output from layer 8 through concat; the fused feature map is upsampled by the nearest feature map and then fused with the feature map output from layer 5 through concat; the above result is upsampled by the nearest feature map again and then concatted with the feature map output from layer 2 to obtain the fused feature map; this feature map is then passed through module C3 to finally obtain the output results of each prediction layer; invalid detection layers in the prediction layers are pruned. The overall process is to combine the feature map extracted from the backbone network with the Feature Pyramid Network (FPN) and the Path Aggregation Network (PAN) for feature aggregation. The former conveys strong semantic features from top to bottom, and the latter conveys strong localization features from bottom to top, and obtains the output results of each prediction layer based on the aggregated feature map; this result includes a newly added high-resolution target detection layer for small targets, which improves the detection effect of small targets, while deleting the original low-resolution target detection layer to avoid invalid computation.
[0068] S103. Train the improved YOLOv5 network model using the dataset.
[0069] Specifically, the following steps are included:
[0070] (1) Calculate the corresponding initial anchor box values for the dataset using the Kmeans++ method;
[0071] The initialization principle of the Kmeans++ algorithm is to maximize the distance between the initial cluster centers in order to obtain better cluster centers.
[0072] (2) Perform input data preprocessing on the training set data;
[0073] This involves splicing the data using random scaling, random cropping, and random arrangement to achieve Mosaic data augmentation.
[0074] (3) Set the parameters as shown in the table below, and input the preprocessed training set data into the improved YOLOV5 network model for training;
[0075] Table 2 Parameter Settings
[0076]
[0077] (4) After N iterations, a trained small target detection model is obtained, and the trained small target detection model is then validated and evaluated using the validation set data.
[0078] The evaluation metric used is the mean average precision (mAP). The mean average precision is related to precision and recall.
[0079] Precision refers to the proportion of correctly predicted samples out of the total number of samples in a dataset, as shown in the formula below:
[0080]
[0081] Recall rate refers to the proportion of correctly predicted positive samples in the predicted dataset to the actual number of positive samples. The formula is as follows:
[0082]
[0083] Average precision refers to the area enclosed by the curve formed by recall on the horizontal axis and precision on the vertical axis, as shown in the formula below:
[0084]
[0085] The mean precision is the average precision of all classes in the dataset. The formula is shown below:
[0086]
[0087] Where TP represents true positives, FP represents false positives, TN represents true negatives, FN represents false negatives, AP value refers to the area under the PR curve, and N represents the total number of categories detected.
[0088] In this embodiment, the results are mainly compared with those of the original YOLOv5 network structure, the network structure using the ShuffleNetv2 lightweight module, and the network structure of the present invention. The experimental results are shown in Table 3.
[0089] The table shows that compared to the original YOLOv5 architecture, using a lightweight model significantly reduces the model size and the number of parameters. Specifically, the model size decreases from 3.9MB to 1.0MB, and the number of parameters decreases from 1,761,871 to 201,931, demonstrating the effectiveness of the ShuffleNetv2 module in improving model speed. However, the model accuracy also decreases from 97.1% to 94.7%, a reduction of 2.4 percentage points.
[0090] Compared to the original network model, the number of parameters in our proposed model is 752,959, only 42.7% of the original model. Simultaneously, the model size is reduced from 3.9MB to 2.7MB, fully demonstrating the advantages of our proposed model in terms of real-time performance. Furthermore, compared to the original network model, our proposed model achieves a 1.2% improvement in mAp, proving that the high-resolution object detection layer added in this embodiment effectively enhances the model's feature extraction capability for small targets.
[0091] Table 3 Performance Comparison of Different Algorithms
[0092]
[0093] To better illustrate the feasibility of the model presented in this paper, a portion of the test set images were selected for testing. The results are as follows: Figure 5 As shown. Comparison Figure 5 (a) and Figure 5 (b) In densely distributed scenarios, YOLOv5's complex background information interfered with the correct detection of the target, while the proposed model accurately distinguished between the background and the target. Furthermore, the confidence level of the proposed model's detection performance was also higher than that of the original model. (Comparison) Figure 5 (c) and Figure 5 (d) In cases where the background is simple and the target is singular, YOLOv5 incorrectly detects the same target repeatedly, while the model in this paper identifies the target in the image.
[0094] Based on the comparison, the model presented in this paper achieves better feature fusion results by retaining more shallow spatial information and integrating global information, thereby enhancing the ability to extract target feature information and effectively improving the problems of false detection and re-detection in the YOLOv5 model.
[0095] Specific embodiments of the present invention have been described above with reference to the accompanying drawings. It should be understood that the present invention is not limited to the specific embodiments described above, and those skilled in the art can make various modifications or variations within the scope of the claims, which do not affect the essence of the present invention.
Claims
1. A small target detection method based on enhanced feature extraction, characterized in that, include: (1) Collect images to construct a dataset, and divide the dataset into a training set and a test set; (2) Constructing an improved YOLOv5 network model: (2.1) The backbone network is replaced by ShuffleNet V2: it adopts a dual convolution kernel processing and connection method, and merges continuous element-wise operations at the same time; The backbone network is replaced with ShuffleNet V2, which includes replacing the convolutional modules and C3 modules in layers 3-8 of the original network with the convolutional modules of ShuffleNet V2, and adopting a channel splitting and channel shuffling strategy. (2.2) Add coordinate attention mechanism: embed position information into channel attention, integrate global image information, and weaken background information interference; The addition of coordinate attention mechanisms includes adding coordinate attention modules after layers 2, 5, 8, and 12 of the backbone network; (2.3) Improve the feature extraction structure: deepen the feature pyramid FPN, increase the reuse of shallow features, retain richer location information, enhance the feature extraction capability of small targets, and at the same time, due to the characteristics of the target scale distribution in the dataset, delete the low-resolution target detection layer to avoid invalid computation; The improved feature extraction structure includes: the feature map output from layer 14 in the network structure is upsampled by the nearest neighbor and then fused with the feature map output from layer 8 through concat; the fused feature map is upsampled by the nearest neighbor and then fused with the feature map output from layer 5 through concat; the above result is upsampled by the nearest neighbor again and then concatted with the feature map output from layer 2 to obtain the fused feature map; this feature map is then passed through module C3 to finally obtain the output results of each prediction layer; this result includes a newly added high-resolution target detection layer for small targets, which improves the detection effect of small targets, while the original low-resolution target detection layer is removed to avoid invalid computation; (3) Use the dataset to train the improved YOLOV5 network model to obtain a small target detection model.
2. The small target detection method based on enhanced feature extraction according to claim 1, characterized in that: The channel segmentation involves dividing the c feature channel input of each unit into two branches, each with c-c' and c' channels, where c' takes the value c / 2; then, the two branches obtained by convolution are connected to maintain the same number of channels. The channel shuffling ensures that information can be exchanged between different branches.
3. The small target detection method based on enhanced feature extraction according to claim 1, characterized in that, The coordinate attention module is implemented as follows: (2.2.1) Coordinate attention embedding: Global pooling is decomposed to obtain one-to-one one-dimensional feature encoding; For input of C×H×W dimension, where C is the number of channels, H is the height and W is the width, each channel is encoded according to the horizontal and vertical coordinates using a pooling kernel of dimension (H,1) or (1,W); (2.2.2) Coordinate attention generation: The representation obtained in step (2.2.1) is transformed using a 1×1 convolution kernel transformation function, and the transformation result is decomposed into two separate tensors along the spatial dimension. Finally, the convolution transformation is used to transform the input and the input into tensors with the same number of channels, where the convolution transformation uses the sigmoid activation function. (2.2.3) Expand the results in step (2.2.2) as attention weights.
4. The small target detection method based on enhanced feature extraction according to claim 1, characterized in that, Training the improved YOLOv5 network model using the training set data includes: (3.1) Calculate the corresponding initial anchor box values for the dataset using the Kmeans++ method; (3.2) Perform input data preprocessing on the training set data; that is, splice the data by random scaling, random cropping and random arrangement to achieve Mosaic data augmentation; (3.3) Input the preprocessed training set data into the improved YOLOV5 network model for training to obtain a small target detection model; (3.4) After N iterations, the trained small target detection model is validated and evaluated using the validation set data.
5. The small target detection method based on enhanced feature extraction according to claim 4, characterized in that: The evaluation metric used is the mean precision, which is related to precision and recall. The accuracy formula is as follows: The recall rate formula is as follows: Average precision refers to the area enclosed by the curve formed by recall on the horizontal axis and precision on the vertical axis, as shown in the formula below: The mean precision is the average precision of all classes in the dataset, and the formula is as follows: Where: TP represents true positives, FP represents false positives, TN represents true negatives, FN represents false negatives, AP value refers to the area under the PR curve, and N represents the total number of categories detected.