A small target detection method based on shallow feature enhancement for complex scenes

Abstract

The present application relates to the technical field of image processing, in particular to a small target detection method based on shallow feature enhancement for complex scenes, which mainly solves the problems of difficulty in identifying small targets and high model complexity of existing detection methods. The present application first adds an attention module directly behind the output layer of the backbone network to filter redundant information and select representative original shallow features. Next, the idea of cross-stage partial connection is used to reduce the computational burden of the model. Then, a shallow feature detection head is used instead of a deep feature detection head in the feature fusion part to extract more shallow features while reducing the parameter quantity of the model. Finally, the SIoU loss function is used in the training stage to speed up the convergence speed of the model. The present application has the advantage of improving the detection accuracy of small targets without increasing the parameter quantity of the original model.

Description

Technical Field

[0001] This invention relates to the field of image processing technology, and specifically to a method for small target detection based on shallow feature enhancement for complex scenes. Background Technology

[0002] In image processing, it is often necessary to identify targets composed of very few pixels; these targets are typically defined as small targets. Small target detection is crucial in many practical applications: for example, detecting early small fire sources in aerial images helps prevent severe forest fires, and locating small key landmarks in drone images plays a vital role in disaster emergency response. Therefore, accurate and rapid detection of small targets has significant practical value and has attracted widespread attention from numerous researchers.

[0003] Currently, the latest small object detection models are improvements based on deep learning methods. Since small object detection is closely related to the shallow features of the target, many small object detection models attempt to extract more shallow features with localization information from the image. One effective approach is to increase the number of detection heads, as this reduces feature loss during the network downsampling stage. Another common approach is to introduce attention mechanisms, such as using an attention module in the feature fusion part when fusing information from different backbone network output layers to highlight shallow features ignored by traditional object detection models.

[0004] In small object detection, shallow feature detectors are often more effective than deep feature detectors. Therefore, by adding additional shallow feature detectors while discarding deep feature detectors, the model complexity can remain unchanged, and more shallow features can be extracted. Adding an attention module after each backbone network output layer can directly filter out unnecessary raw information and highlight shallow features closely related to small objects. Furthermore, since the raw features passed from the backbone network are connected in the feature fusion section, the number of parameters in the feature fusion section is much larger than the number of parameters in the backbone network. Moving the attention module from the feature fusion section to the backbone network can reduce the number of model parameters.

[0005] To further improve the detection accuracy of small targets in complex scenes, this invention proposes a small target detection method based on shallow feature enhancement for complex scenes. Summary of the Invention

[0006] The purpose of this invention is to address the difficulties in identifying small targets in complex scenes and the high model complexity of existing detection methods. Based on the YOLOv5 model, this invention proposes a shallow feature enhancement-based small target detection method for complex scenes: 1) To select representative original features and reduce redundant model parameters, an attention module is directly added after the output layer of each backbone network; 2) To extract more shallow features, a shallow feature detection head is used instead of a deep feature detection head in the feature fusion part; 3) To reduce the computational burden of the model, the idea of ​​cross-level partial connections is used; 4) To accelerate the convergence speed of the model, the CIoU loss function is replaced with the SIoU loss function.

[0007] The technical solution adopted in this invention includes the following steps: a small target detection method based on shallow feature enhancement for complex scenes, the method including the following steps:

[0008] Step A: Obtain images of small targets in complex scenes, divide them into training and testing sets, load the model, and initialize the relevant parameters;

[0009] Step B: Preprocess each image in the training set samples and load the preprocessed image into the model;

[0010] Step C: Extract image features in the model backbone network and feed the extracted features into the feature fusion part for multi-scale fusion;

[0011] Step D: Calculate the loss function of the model based on the fused features, then minimize the loss function and update the model parameters using the gradient descent algorithm;

[0012] Step E: After training, save the optimal model parameter weights for subsequent model testing;

[0013] Step F: Predict the image in the test set. Based on the model parameter weights trained in Step E, predict the location and category of the target. Then, obtain the final prediction result through non-maximum suppression.

[0014] Step B, which involves preprocessing each image in the training set, specifically includes:

[0015] Step B1: The four images are stitched together into one image by random scaling, random cropping, and random arrangement. This not only enriches the background of the image, but also greatly increases the number of images that the model can read in.

[0016] Step B2: When the defined anchor box size is not suitable for the training set, the k-means algorithm is used to cluster the bounding box labels of the training set so that the generated anchor box size is more suitable for the training set image.

[0017] Step B3: Scale the read-in image to a fixed size of 640×640 pixels. By unifying the image to a fixed size, the model can process more images at once, thus improving the detection speed.

[0018] Step C, which involves extracting image features from the model backbone network and then feeding the extracted features into the feature fusion part for multi-scale fusion, specifically involves:

[0019] Step C1: Use the preprocessed sample image described in step B as the input to the backbone network, and add an attention module after each output layer of the backbone network. This attention module consists of an input, a channel attention module, a spatial attention module, and an output. Assume the input feature map is F∈R. C×H×W Then, after passing through the channel attention module, feature map M is generated. c ∈R C ×1×1 , the result M c The output feature map is formed by multiplying it with the input feature map F. Then, F′ is used as input and passed through the spatial attention module to generate feature map M. s ∈R 1×H×W The final result M s The final output feature map is formed by multiplying it with the input feature map F′. in This represents element-wise multiplication. The channel attention module focuses more on category information in the input image, while the spatial attention module focuses more on location information. Combining these two modules can alter the weights of the original feature channels, thereby enhancing effective features and suppressing ineffective ones. Furthermore, adding an attention module to the backbone network instead of the feature fusion part can further reduce the model's parameters, as the feature fusion part contains more parameters than the backbone network.

[0020] Step C2: Divide the original SPPF module into two parts to reduce the computational burden: The first part is processed using the original SPPF module to extract more hidden features at multiple scales, and the second part is processed using only one convolution module. Finally, the outputs of the two parts are concatenated and passed to the feature fusion part.

[0021] Step C3: Multi-scale fusion of features extracted from the backbone network: In the feature fusion part, this invention removes the largest-scale detector head in PANet and adds a small-scale detector head. This effectively extracts richer shallow information from the image and reduces the additional parameters added by the additional detector head. The fusion process is as follows: First, the fine features output from layers 2 to 5 of the backbone network are passed to the SPAN feature pyramid. Next, in this feature pyramid, transposed convolution is used to upsample the large-scale feature map and connect it with the adjacent small-scale feature map. Then, GSConv is used to downsample the small-scale feature map and connect it with the adjacent large-scale feature map. Finally, the fused features from the downsampling process are passed to the detector head for prediction.

[0022] Step D, which involves calculating the model's loss function, specifically involves:

[0023] Step D1: The loss function consists of three parts. The first part is the regression loss, calculated using the SIoU loss function instead of the more common CIoU loss function. SIoU loss takes into account angular loss and can add more penalty to predicted bounding boxes that are far from the ground truth bounding boxes. The SIoU loss function is as follows:

[0024]

[0025] Among them, L IoU Ω is the intersection-over-union ratio of the predicted bounding box and the ground truth bounding box, Δ is the angle and distance loss, and Ω is the shape loss.

[0026] Step D2, the second part of the loss function is the classification loss L. cls It represents the classification error of the target, L. cls The smaller the value, the closer the predicted label is to the true label;

[0027] Step D3, the third part of the loss function is the confidence loss L. obj It is used to exclude empty prediction boxes. The total loss function is a linear weighted combination of these three losses:

[0028] L=W SIoU L SIoU +W cls L cls +W obj L obj

[0029] Among them, W SIoU W cls W obj These represent the weights of regression loss, classification loss, and confidence loss, respectively. Attached Figure Description

[0030] Figure 1 This is a flowchart illustrating the implementation of an embodiment of the present invention.

[0031] Figure 2 This is a schematic diagram of a complete model according to an embodiment of the present invention.

[0032] Figure 3 This is a schematic diagram of the attention module according to an embodiment of the present invention.

[0033] Figure 4 This is a schematic diagram of the cross-level connection and integration with the SPPF module in an embodiment of the present invention.

[0034] Figure 5 This is a schematic diagram of the feature fusion portion of an embodiment of the present invention.

[0035] Figure 6 This is a visualization comparison of the results between embodiments of the present invention and other methods. Detailed Implementation

[0036] To make the objectives, technical solutions, and advantages of the embodiments of the present invention clearer, the technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings.

[0037] like Figure 1 , Figure 2 As shown, to address the difficulties in identifying small targets and the high model complexity of existing detection methods, this embodiment proposes a small target detection method based on shallow feature enhancement for complex scenes, building upon the YOLOv5 model: 1) To select representative original features and reduce redundant model parameters, an attention module is directly added after the output layer of each backbone network; 2) To extract more shallow features, a shallow feature detection head is used instead of a deep feature detection head in the feature fusion part; 3) To reduce the computational burden of the model, the idea of ​​cross-level partial connections is used; 4) To accelerate the convergence speed of the model, the CIoU loss function is replaced with the SIoU loss function.

[0038] The method specifically includes the following steps:

[0039] Step A: Obtain images of small targets in complex scenes, divide them into training and testing sets, load the model, and initialize the relevant parameters;

[0040] Step B: Preprocess each image in the training set samples and load the preprocessed image into the model;

[0041] Step C: Extract image features in the model backbone network and feed the extracted features into the feature fusion part for multi-scale fusion;

[0042] Step D: Calculate the loss function of the model based on the fused features, then minimize the loss function and update the model parameters using the gradient descent algorithm;

[0043] Step E: After training, save the optimal model parameter weights for subsequent model testing;

[0044] Step F: Predict the image in the test set. Based on the model parameter weights trained in Step E, predict the location and category of the target. Then, obtain the final prediction result through non-maximum suppression.

[0045] Step B, which involves preprocessing each image in the training set, specifically includes:

[0046] Step B1: The four images are stitched together into one image by random scaling, random cropping, and random arrangement. This not only enriches the background of the image, but also greatly increases the number of images that the model can read in.

[0047] Step B2: When the defined anchor box size is not suitable for the training set, the k-means algorithm is used to cluster the bounding box labels of the training set so that the generated anchor box size is more suitable for the training set image.

[0048] Step B3: Scale the read-in image to a fixed size of 640×640 pixels. By unifying the image to a fixed size, the model can process more images at once, thus improving the detection speed.

[0049] Step C, which involves extracting image features from the model backbone network and then feeding the extracted features into the feature fusion part for multi-scale fusion, specifically involves:

[0050] Step C1: Use the preprocessed sample image described in step B as the input to the backbone network, and add an attention module after each output layer of the backbone network: (e.g.) Figure 3 As shown, this attention module consists of an input, a channel attention module, a spatial attention module, and an output. Assume the input feature map is F∈R. C×H×W Then, after passing through the channel attention module, feature map M is generated. c ∈R C×1×1 , the result M c The output feature map is formed by multiplying it with the input feature map F. Then, F′ is used as input and passed through the spatial attention module to generate feature map M. s ∈R 1×H×W The final result M S The final output feature map is formed by multiplying it with the input feature map F′. in This represents element-wise multiplication. The channel attention module focuses more on category information in the input image, while the spatial attention module focuses more on location information. Combining these two modules can alter the weights of the original feature channels, thereby enhancing effective features and suppressing ineffective ones. Furthermore, adding an attention module to the backbone network instead of the feature fusion part can further reduce the model's parameters, as the feature fusion part contains more parameters than the backbone network.

[0051] Step C2, as follows Figure 4 As shown, the original SPPF module is divided into two parts to reduce the computational burden: the first part is processed using the original SPPF module to extract more hidden features at multiple scales, and the second part is processed using only one convolution module. Finally, the outputs of the two parts are concatenated and passed to the feature fusion part.

[0052] Step C3, as follows Figure 5 As shown, the features extracted by the backbone network are fused at multiple scales: In the feature fusion part, this invention removes the largest-scale detector head in PANet and adds a small-scale detector head, which effectively extracts richer shallow information from the image and reduces the additional parameters added by the additional detector head. The fusion process is as follows: First, the fine features output from the second to fifth layers of the backbone network are passed to the SPAN feature pyramid. Next, in this feature pyramid, the large-scale feature map is upsampled using transposed convolution and connected to the adjacent small-scale feature map. Then, the small-scale feature map is downsampled using GSConv and connected to the adjacent large-scale feature map. Finally, the fused features from the downsampling process are passed to the detector head for prediction.

[0053] Step D, which involves calculating the model's loss function, specifically involves:

[0054] Step D1: The loss function consists of three parts. The first part is the regression loss, calculated using the SIoU loss function instead of the more common CIoU loss function. SIoU loss takes into account angular loss and can add more penalty to predicted bounding boxes that are far from the ground truth bounding boxes. The SIoU loss function is as follows:

[0055]

[0056] Among them, L IoU Ω is the intersection-over-union ratio of the predicted bounding box and the ground truth bounding box, Δ is the angle and distance loss, and Ω is the shape loss.

[0057] Step D2, the second part of the loss function is the classification loss L. cls It represents the classification error of the target, L. clsThe smaller the value, the closer the predicted label is to the true label;

[0058] Step D3, the third part of the loss function is the confidence loss L. obj It is used to exclude empty prediction boxes. The total loss function is a linear weighted combination of these three losses:

[0059] L=W SIoU L SIoU +W cls L cls +W obj L obj

[0060] Among them, W SIoU W cls W obj These represent the weights of regression loss, classification loss, and confidence loss, respectively.

[0061] like Figure 6 As shown, this invention selects some representative images to visually demonstrate its advantages. The first and second columns of the image respectively show the recognition results of YOLOv5 and this invention. Taking the first row as an example, YOLOv5 only identifies a few vehicles in the distance, while the method of this invention successfully labels almost all vehicles. Furthermore, in the second row, this invention not only successfully identifies vehicles at the end of the road that YOLOv5 ignores, but also corrects some objects incorrectly labeled as "cars" by YOLOv5. For example, an object next to the roof of a building in the upper right corner of the image was incorrectly labeled as a "car" by YOLOv5; this incorrect label is removed from the recognition results of this invention.

[0062] The above description is merely a preferred embodiment of the present invention and is not intended to limit the invention. Various modifications and variations can be made to the present invention by those skilled in the art. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the scope of the claims of the present invention.

Claims

1. A small target detection method based on shallow feature enhancement for complex scenes, characterized in that, The method comprises the following steps: Step A: Obtain small target images in a complex scene, and divide them into a training set and a test set, then load the model and initialize the related parameters; Step B: Preprocess each image in the training set sample, and load the preprocessed image into the model; Step C: Extract image features in the model backbone network, and send the extracted features to the feature fusion part for multi-scale fusion; Step D: Calculate the loss function of the model according to the fused features, then minimize the loss function, and update the parameters of the model through the gradient descent algorithm; Step E: Save the optimal model parameter weight after training, which is used for subsequent model testing; Step F: Predict the image of the test set, predict the position and category of the target according to the model parameter weight trained in step E, and then obtain the final prediction result through non-maximum suppression; In step C, the image features are extracted in the model backbone network, and the extracted features are sent to the feature fusion part for multi-scale fusion, which is specifically: Step C1: Use the preprocessed sample image described in Step B as the input to the backbone network, and add an attention module after each output layer of the backbone network. This attention module consists of an input, a channel attention module, a spatial attention module, and an output. Assume the input feature map is... Then, after passing through the channel attention module, feature maps are generated. , the results With input feature map Multiplication forms the output feature map Then The input is processed by the spatial attention module to generate a feature map. The final result With input feature map Multiplication forms the final output feature map. ,in This represents element-wise multiplication; the channel attention module focuses more on category information in the input image, while the spatial attention module focuses more on location information in the input image. Combining the two modules can change the weight information of the original feature channels, thereby enhancing effective features and suppressing ineffective features; in addition, adding the attention module to the backbone network instead of the feature fusion part can further reduce the model's parameters, because the feature fusion part contains more parameters than the backbone network. Step C2: Divide the original SPPF module into two parts to reduce the computational burden: the first part uses the original SPPF module for processing to extract more hidden features at multiple scales, and the second part only uses a convolution module for processing, and finally the outputs of the two parts are connected together and transmitted to the feature fusion part; Step C3: Multi-scale fusion of the features extracted by the backbone network: in the feature fusion part, the detection head of the largest scale in PANet is removed, and a small-scale detection head is added, which effectively extracts more rich shallow information from the image and reduces the additional parameters increased by the additional detection head; The fusion process is as follows: first, the fine features output by the second to fifth layers of the backbone network are transmitted to the SPAN feature pyramid; Next, in this feature pyramid, the large-scale feature map is upsampled using transposed convolution and connected with the small-scale feature map next to it; Then, the small-scale feature map is downsampled using GSConv and connected with the large-scale feature map next to it; Finally, the fused features in the downsampling process are transmitted to the detection head for prediction.

2. The method of claim 1, wherein, Step B: Preprocess each image in the training set sample, which is specifically: Step B1: Four images are spliced into one image by random scaling, random cropping and random arrangement, which not only enriches the image background, but also greatly increases the number of model image inputs; Step B2: When the defined anchor box size is not suitable for the training set, use the k-means algorithm to cluster the boundary box labels of the training set, so that the generated anchor box size is more suitable for the training set image; Step B3: Scale the read-in image to 640x640 pixels to a fixed size, and unify the size to a fixed size, so that the model can process more images at a time, and the detection speed will be improved.

3. The method of claim 1, wherein, Step D: The calculation of the loss function of the model is specifically: Step D1, the loss function is composed of three parts; the first part is the regression loss, here the SIoU loss function is used instead of the commonly used CIoU loss function; the SIoU loss considers the angle loss, and can add more punishment to the predicted bounding box far away from the real box; the SIoU loss function is as follows: wherein, is the intersection over union of the predicted box and the true box, is the angle and distance loss, is the shape loss; Step D2, the second part of the loss function is the classification loss which represents the classification error of the target, The smaller the value of the smaller the predicted label is closer to the true label; Step D3, the third part of the loss function is the confidence loss which is used to exclude empty predicted boxes; the total loss function is a linear weighted combination of the three losses: where, , , respectively represent the weights of regression loss, classification loss and confidence loss.

Images