A small target detection method based on improved YOLOv8

Abstract

The invention discloses a small target detection method based on improved YOLOv8, comprising: inputting a small target image to be detected into a pre-trained small target detection model based on improved YOLOv8 for identification, and obtaining a detection result; wherein the training method of the small target detection model based on improved YOLOv8 comprises: obtaining a small target image data set, and dividing it into a training set and a verification set; using a backbone network ATDeNet to replace the backbone network of YOLOv8, and constructing a small target detection model based on improved YOLOv8; using the training set and the verification set to train the constructed small target detection model, and obtaining a trained small target detection model based on improved YOLOv8. The accuracy and efficiency of small target detection can be significantly improved.

Description

Technical Field

[0001] The present application belongs to the field of target detection in computer vision, and in particular to a small target detection method based on improved YOLOv8. Background Art

[0002] With the rapid development of the drone industry, drones equipped with cameras have been deployed in all walks of life. The images transmitted by drones generally have a wide field of view and rich information, and can provide information that ground photography cannot provide. However, as the viewing angle increases and the size of objects in the image decreases, the difficulty of processing drone image information increases compared to traditional images. Drone images are characterized by small target size, little appearance and geometric information, and complex backgrounds, all of which pose great challenges to target detection. Therefore, how to improve small target detection is an important topic in the field of computer vision.

[0003] Since its first release in 2015, the YOLO (You Only Look Once) algorithm has played an important role in the field of real-time target detection technology. After multiple iterations and upgrades, each generation of versions is committed to improving detection efficiency, accuracy and practicality. YOLOv1 was the first to adopt an overall prediction framework, abandoning candidate region generation, and achieved ultra-fast real-time detection speed. Although the positioning accuracy was relatively insufficient, it led a new direction for real-time target detection with its unique advantages. Subsequently, during the evolution from YOLOv2 to YOLOv8, the latest research results of deep learning were continuously absorbed and integrated. Through the use of multi-scale training, network structure adjustment, data enhancement technology, and various innovative technologies such as weighted residual connection and cross-stage partial connection, the detection accuracy was gradually improved and the high-speed operation characteristics were maintained. In particular, in the YOLOv8 stage, not only did it make major breakthroughs in speed and accuracy, but it also paid more attention to meeting the requirements of real-time and high classification performance in industrial scenarios, making the YOLO series widely used in automated quality inspection and other related fields, becoming a model for the advancement of real-time target detection technology in the field of computer vision.

[0004] However, YOLOv8 still faces the problems of feature information being easily lost in deep networks and insufficient feature fusion capabilities in the field of small target detection, making it difficult to ensure the accuracy and efficiency of small target detection. Summary of the invention

[0005] The purpose of the present invention is to overcome the shortcomings in the prior art, provide a small target detection method based on improved YOLOv8, and provide a new backbone network ATDeNet that integrates the ATDE module and the hybrid spatial pyramid structure MixSPP to replace the original backbone network of YOLOv8, which can alleviate the problem that the features of small targets are easily lost in deep networks and improve the accuracy and efficiency of small target detection.

[0006] To achieve the above object, the present invention is implemented by adopting the following technical solutions:

[0007] In a first aspect, the present invention provides a small target detection method based on improved YOLOv8, comprising:

[0008] The small target image to be detected is input into the pre-trained small target detection model based on improved YOLOv8 for recognition to obtain the detection result;

[0009] The training method of the small target detection model based on the improved YOLOv8 includes:

[0010] Obtain a small object image dataset and divide it into a training set and a validation set;

[0011] Use the backbone network ATDeNet to replace the backbone network of YOLOv8 and build a small target detection model based on the improved YOLOv8;

[0012] The constructed small target detection model is trained using the training set and the validation set to obtain a trained small target detection model based on the improved YOLOv8.

[0013] Furthermore, the backbone network ATDeNet is used to replace the backbone network of YOLOv8 to build a small target detection model based on the improved YOLOv8, including:

[0014] Constructing a backbone network ATDeNet; the backbone network ATDeNet includes an input layer Stem and a first backbone part Stage1, a second backbone part Stage2, a third backbone part Stage3 and a fourth backbone part Stage4; the input layer Stem includes a convolutional layer; the first backbone part Stage1 includes a convolutional layer and an ATDE module; the second backbone part Stage2 includes a convolutional layer and an ATDE module; the third backbone part Stage3 includes a convolutional layer and a first convolutional module ATDE; the fourth backbone part Stage4 includes a convolutional layer and a pyramid structure module MixSPP;

[0015] Among them, the backbone network ATDeNet adopts a design of gradually reducing the convolution kernel size, specifically (7, 5, 3, 3), and adopts a convolution depth design of (1, 1, 3, 1), which can increase the receptive field at the beginning of detection while reducing the receptive field of subsequent predicted feature maps, avoiding excessive redundant background information.

[0016] The first convolution module ATDE can perform multi-level splicing of images and use 1x1 convolution to increase and then reduce the dimension to extract the spatial detail information of the image; the pyramid structure module MixSPP parallelizes the spatial pyramid pooling and convolution modules, and adopts a layer-by-layer fusion method;

[0017] The backbone network of YOLOv8 is replaced with the backbone network ATDeNet, and a small target detection model based on the improved YOLOv8 is constructed.

[0018] Furthermore, the first convolution module ATDE includes a depthwise separable convolution and a convolution with a convolution kernel of 1x1, so that the network can capture the model input data well and improve the expression ability and information processing ability of the model.

[0019] The pyramid structure module MixSPP includes a convolution layer, a maximum pooling layer and a depth-separable convolution, which can reduce the loss of multi-scale information in the fusion process. Even if the pyramid pool has a weak response to the target, multiple 1x1 convolutions are stacked after extracting features, which can reduce the risk of overfitting and improve the robustness of the model.

[0020] Furthermore, the convolution layer includes batch normalization BN and SiLU activation functions, and different convolution parameters are set.

[0021] Furthermore, the small target image dataset is derived from the VisDrone2019 target detection dataset.

[0022] Furthermore, the small object image dataset contains small object images under different scenes, weather and lighting conditions.

[0023] Furthermore, the method also includes: dividing a test set from the small target image data set, and using the test set to analyze the detection accuracy, detection speed and parameter quantity of the trained small target detection model based on the improved YOLOv8.

[0024] In a second aspect, the present invention provides a small target detection system based on improved YOLOv8, comprising:

[0025] The detection module is used to input the small target image to be detected into the pre-trained small target detection model based on improved YOLOv8 for recognition and obtain the detection result;

[0026] The training method of the small target detection model based on the improved YOLOv8 includes:

[0027] Obtain a small object image dataset and divide it into a training set and a validation set;

[0028] Use the backbone network ATDeNet to replace the backbone network of YOLOv8 and build a small target detection model based on the improved YOLOv8;

[0029] The constructed small target detection model is trained using the training set and the validation set to obtain a trained small target detection model based on the improved YOLOv8.

[0030] In a third aspect, the present invention provides a computer storage medium having a computer program stored thereon, characterized in that when the computer program is executed by a processor, the small target detection method based on the improved YOLOv8 as described above is implemented.

[0031] Compared with the prior art, the present invention has the following beneficial effects:

[0032] (1) The detection method proposed in the present invention is based on the YOLOv8 target detection algorithm, which enables it to have high detection accuracy under normal conditions. The present invention replaces the backbone extraction network of YOLOv8 with the backbone network ATDeNet proposed for small target detection, so that the model can effectively extract the spatial detail information of the image and improve the detection accuracy of small targets.

[0033] (2) The present invention proposes a convolutional module with stronger feature extraction capability, ATDE, which can effectively extract the spatial detail information of the image and improve the detection accuracy of small target detection.

[0034] (3) The present invention proposes a backbone network for small targets - ATDeNet. The network redesigns the receptive field so that the receptive field of the bottom feature map is small and the receptive field of the top feature map is large, which is suitable for the characteristics of small target detection.

[0035] (4) This paper proposes a new hybrid spatial pyramid structure - MixSPP, which combines spatial pyramid pooling and convolution modules in parallel and adopts a layer-by-layer fusion approach to reduce the loss of multi-scale information in the fusion process. BRIEF DESCRIPTION OF THE DRAWINGS

[0036] Figure 1 Shown is a schematic diagram of the structure of the backbone network ATDeNet;

[0037] Figure 2 and Figure 3 Shown is a schematic diagram of the structure of the ATDE convolution module;

[0038] Figure 4 and Figure 5 Shown is a schematic diagram of the structure of the mixed pyramid MixSPP module;

[0039] Figure 6The figure shows the overall network structure diagram of the small target detection method based on the improved YOLOv8 of the present invention. DETAILED DESCRIPTION

[0040] The technical concept of the present invention is: replacing the backbone network of YOLOv8 with the ATDeNet network for small target detection, so that the model can effectively extract the spatial detail information of the image and improve the detection accuracy of small targets. The ATDE module can effectively extract the spatial detail information of the image and improve the detection accuracy of small target detection. At the same time, a hybrid spatial pyramid structure MixSPP is adopted, which parallelizes the spatial pyramid pooling and convolution modules, and adopts a layer-by-layer fusion method to reduce the loss of multi-scale information in the fusion process.

[0041] The following will be combined with the drawings in the embodiments of the present application to clearly and completely describe the technical solutions in the embodiments of the present application. Obviously, the described embodiments are only part of the embodiments of the present application, not all of the embodiments. The following description of at least one exemplary embodiment is actually only illustrative and is by no means intended to limit the present application and its application or use.

[0042] Example 1

[0043] This embodiment discloses a small target detection method based on improved YOLOv8, including:

[0044] The small target image to be detected is input into the pre-trained small target detection model based on improved YOLOv8 for recognition to obtain the detection result;

[0045] The training method of the small target detection model based on the improved YOLOv8 includes:

[0046] Obtain a small object image dataset and divide it into a training set and a validation set;

[0047] Use the backbone network ATDeNet (attention to detail network) to replace the backbone network of YOLOv8 and build a small target detection model based on the improved YOLOv8;

[0048] The constructed small target detection model is trained using the training set and the validation set to obtain a trained small target detection model based on the improved YOLOv8.

[0049] The small target detection method based on improved YOLOv8 provided in this embodiment is described below in conjunction with a preferred embodiment to illustrate the contents involved in the above embodiment:

[0050] The software configuration environment of this embodiment is: CentOS 7 operating system, Python version 3.8.18, GPU version NVIDIA TITAN Xp 12G, CUDA version 11.3, Pytorch version 1.12.1 and MMDetection version 3.2.0.

[0051] S1: Obtain a small target image dataset and divide it into a training set, a validation set, and a test set. The drone dataset used in this embodiment is the VisDrone2019 target detection dataset, which was open-sourced by the Tianjin University team in 2019. The official training set includes 6471 images and the validation set includes 548 images, covering 10 categories of targets. This dataset was collected using different models of drones in different scenes, weather, and lighting conditions, and can objectively reflect the performance of the detector.

[0052] S2: Construct the backbone network ATDeNet. Figure 1 The structural diagram of the backbone network ATDeNet shown includes: an input layer Stem and four backbone parts Stage1, Stage2, Stage3 and Stage4;

[0053] The input layer inputs an image with a resolution of 640x640.

[0054] The input layer Stem includes a convolution layer with a convolution kernel of 3x3, a stride of 2, a padding of 1, and an output channel number of 32. The convolution layer includes batch normalization BN and SiLU activation functions.

[0055] The Stage 1 includes a convolution layer with a convolution kernel of 3x3, a stride of 2, a padding of 1, an output channel number of 64, a batch normalization BN and a SiLU activation function, and an ATDE module with parameter settings n of 1 and d of 4.

[0056] refer to Figure 1 The structures of Stage 2 and Stage 3 are similar to those of Stage 1, and the only difference is the convolution parameters, which will not be described here.

[0057] The Stage 4 includes a convolution layer with a convolution kernel of 3x3, a stride of 2, a padding of 1, an output channel number of 512, a batch normalization BN and a SiLU activation function, and a MixSPP module.

[0058] refer to Figure 2 and Figure 3The ATDE module consists of n depth-wise separable convolutions and three convolutions with 1x1 kernels, each of which is followed by a batch normalization layer and an activation function. The ATDE module concatenates the input image with the image that has undergone a depth-wise separable convolution and the image that has undergone n depth-wise separable convolutions to form a new feature map to increase the expressiveness of the feature map. The concatenated feature map is then mapped to a high dimension using a 1x1 convolution, where the dimension is (1+d)xn, and then restored to the original dimension using a 1x1 convolution.

[0059] refer to Figure 4 and Figure 5 , the MixSPP module consists of 5 convolutional layers, 3 maximum pooling layers, depthwise separable convolution, etc. The MixSPP module uses a convolutional layer with a convolution kernel size of 1x1, a stride of 1, no padding, and an output channel number of 512 to perform a convolution operation on the image, and then uses batch normalization and SiLU activation function to perform preliminary processing on the image; then, the image is divided into three regions of different sizes for maximum pooling operations. The convolution kernel sizes of these three regions are k13 x 13, k9 x 9, and k5 x 5, respectively, with a stride of 1 and no padding. The maximum pooling results of each region will be used as separate feature maps; then, these feature maps pass through a depthwise separable convolution layer (DSConv), which consists of a depthwise convolution layer and a pointwise convolution layer. The parameters of this layer are k=3, s=1, p=1, c=512, which means that a depthwise separable convolutional layer with a kernel size of 3x3, stride of 1, padding of 1, and output channels of 512 is used. Again, this layer also includes batch normalization and SiLU activation function; finally, the feature map of the original input image is concatenated with the feature map after the DSConv layer. The concatenated feature map passes through three consecutive convolutional layers, each with a kernel size of 1x1, stride of 1, no padding, and output channels of 1024, followed by batch normalization and SiLU activation function; finally, the concatenated feature map passes through the last convolutional layer with parameters of k=1, s=1, p=0, c=512 to generate the output feature map. Finally, the concatenated feature map passes through the last convolutional layer with parameters of k=1, s=1, p=0, c=512 to generate the output feature map. The MixSPP module can reduce the loss of multi-scale information in the fusion process by combining pooling operations of different sizes and depth-wise separable convolutions. Even if the pyramid pool has a weak response to the target, multiple 1x1 convolutions stacked after feature extraction can reduce the risk of overfitting and improve the robustness of the model.

[0060] S3: Use the backbone network ATDeNet in step S2 to replace the backbone network of YOLOv8 and build a small target detection model based on the improved YOLOv8.

[0061] S4: training the small target detection model constructed in step S3. The training set and the validation set in the data set in step S1 are sent to the constructed small target detection model based on improved YOLOv8 for training to obtain a trained model.

[0062] Set the training parameters, batch size is set to 8, stochastic gradient descent (SGD) is used as the optimizer, Kaiming is used as the weight initialization method, the number of training rounds is set to 150, the number of rounds of data enhancement is turned off to 130, the learning rate lr is 0.01, and the weight decay is 0.0005;

[0063] A comprehensive loss function is used in training, which combines classification loss, bounding box regression loss, and distribution focus loss, combines various losses and weights, and adds a regularization term to prevent overfitting. The specific form of the loss function is as follows:

[0064] &ℒ= λ box N POS ∑ x,y 1 c x,y ∗ 1− q x,y + ∥ b x,y − b x,y ∥ 2 2 ρ 2 + α x,y ν x,y &+ λ cls N POS ∑ x,y ∑ c∈classes y c log( y ̂ c )+(1− y c )log(1− y ̂ c ) &+ λ dfl N POS ∑ x,y 1 c x,y ∗ [−( q (x,y)+1 − q x,y )log( q ̂ x,y ) &+( q x,y − q (x,y)−1 )log( q ̂ (x,y)+1 )]

[0065] Among them, λ cls is the classification loss weight, which means the classification loss weight; box is the box lossweight, meaning the box loss weight; dfl It is the distribution focal loss weight, which means the distribution focal loss weight;

[0066] , , , and The specific formula is as follows:

[0067]

[0068]

[0069] in, is the total number of grid points containing the target object; is an indicator function that shows which grid points contain objects; and Represent the coordinates, width, and height information of the real bounding box and the predicted bounding box respectively. Therefore, The value range is from 0 to 1. The closer the value is to 1, the higher the match between the predicted box and the true box, and the larger the overlapping area. is the coordinate representation of the center point of the true bounding box; Each grid point of the input image The ground truth label belongs to category c, regardless of whether there is an object in the grid point; and Respectively represent the width and height of the corresponding bounding box; and is an auxiliary variable used to adjust the aspect ratio related losses; Represents the predicted probability of the category, Activation function generation; yes The objectness score of the function output; and Represent the width and height of the predicted bounding box respectively. When they are used as a ratio, this value reflects the shape information of the predicted box, that is, its aspect ratio.

[0070] The Soft-NMS method is used for post-processing, which is an improved version of the traditional NMS method. Soft-NMS does not simply discard the overlapping bounding boxes, but applies a soft threshold penalty to them to reduce the confidence of the overlapping area. This can better retain those bounding boxes that may still be valid detection results despite a certain degree of overlap, thereby improving the accuracy of detection and reducing redundant detection;

[0071] Train the algorithm-based small target detection model according to the set training parameters and processing methods;

[0072] The trained small target detection model based on improved YOLOv8 is used for performance evaluation and assessment to obtain a well-converged network model.

[0073] S5: Evaluate the model. Based on the small target detection model obtained after training based on the improved YOLOv8, the average detection accuracy and detection speed of the model are evaluated.

[0074] By setting the corresponding parameters, the sample images in the training set are input into the small target detection model based on the improved YOLOv8, and the number of iterations is set to train the best weight file for small target detection. The direction of network training is continuously adjusted through the loss function, and the map value of the validation set is calculated to verify whether the training has achieved the expected effect. The best weight file is selected by screening the trained weight file and loaded into the small target detection model based on the improved YOLOv8 as the model for small target detection in drone scenarios.

[0075] S6: Input the road surface disease image to be detected into the trained small target detection model based on improved YOLOv8 for recognition to obtain the detection result.

[0076] refer to Figure 6 In this example, under the software configuration environment of CentOS 7 operating system, Python version 3.8.18, GPU version NVIDIATITAN Xp 12G, CUDA version 11.3, Pytorch version 1.12.1, and MMDetection version 3.2.0, ATDeNet has a higher average accuracy than CSPDarkNet under the YOLOv8 framework, especially when dealing with small and medium-sized targets. At the same time, ATDeNet has fewer parameters, which means it may have better computational efficiency. In addition, although ATDeNet has a slightly higher FPS, the difference is not large, so it may not have a significant impact in practical applications. Overall, ATDeNet is a better choice under the YOLOv8 framework, especially when the accuracy of small target detection needs to be improved.

[0077] Example 2

[0078] Based on Example 1, this embodiment discloses a small target detection system based on improved YOLOv8, including:

[0079] The detection module is used to input the small target image to be detected into the pre-trained small target detection model based on improved YOLOv8 for recognition and obtain the detection result;

[0080] The training method of the small target detection model based on the improved YOLOv8 includes:

[0081] Obtain a small object image dataset and divide it into a training set and a validation set;

[0082] Use the backbone network ATDeNet to replace the backbone network of YOLOv8 and build a small target detection model based on the improved YOLOv8;

[0083] The constructed small target detection model is trained using the training set and the validation set to obtain a trained small target detection model based on the improved YOLOv8.

[0084] Example 3

[0085] This embodiment introduces a computer-readable storage medium on which a computer program is stored. When the computer program is executed by a processor, the small target detection method based on improved YOLOv8 described in Embodiment 1 is implemented.

[0086] It will be appreciated by those skilled in the art that embodiments of the present invention may be provided as methods, systems, or computer program products. Therefore, the present invention may take the form of a complete hardware embodiment, a complete software embodiment, or an embodiment combining software and hardware. Furthermore, the present invention may take the form of a computer program product implemented on one or more computer-usable storage media (including but not limited to disk storage, CD-ROM, optical storage, etc.) containing computer-usable program code.

[0087] The present invention is described with reference to flowcharts and / or block diagrams of methods, devices (systems), and computer program products according to embodiments of the present invention. It should be understood that each process and / or block in the flowchart and / or block diagram, as well as the combination of processes and / or blocks in the flowchart and / or block diagram, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, a special-purpose computer, an embedded processor, or other programmable data processing device to produce a machine, so that the instructions executed by the processor of the computer or other programmable data processing device generate instructions for implementing the processes in the flowchart and / or block diagram. Figure 1 A process or multiple processes and / or boxes Figure 1 A device that provides the functions specified in a block or multiple blocks.

[0088] These computer program instructions may also be stored in a computer-readable memory capable of directing a computer or other programmable data processing device to operate in a specific manner, so that the instructions stored in the computer-readable memory produce an article of manufacture comprising an instruction device, which implements the process Figure 1 A process or multiple processes and / or boxes Figure 1 A function specified in one or more boxes.

[0089] The above is only a preferred embodiment of the present invention. It should be pointed out that for ordinary technicians in this technical field, several improvements and modifications can be made without departing from the technical principles of the present invention. These improvements and modifications should also be regarded as the scope of protection of the present invention.

Claims

1. A small target detection method based on improved YOLOv8, characterized in that: include: The small target image to be detected is input into the pre-trained small target detection model based on improved YOLOv8 for recognition to obtain the detection result; The training method of the small target detection model based on the improved YOLOv8 includes: Obtain a small object image dataset and divide it into a training set and a validation set; Use the backbone network ATDeNet to replace the backbone network of YOLOv8 and build a small target detection model based on the improved YOLOv8; The backbone network ATDeNet includes an input layer Stem and four main parts Stage1, Stage2, Stage3 and Stage4; the input layer Stem includes a convolutional layer; the Stage1 includes a convolutional layer and an ATDE module; the Stage2 includes a convolutional layer and an ATDE module; the Stage3 includes a convolutional layer and an ATDE module; the Stage4 includes a convolutional layer and a pyramid structure module MixSPP; The convolution layers of the input layer and the backbone contain batch normalization BN and SiLU activation functions, and different convolution parameters are set; the ATDE module includes depthwise separable convolution and convolution with a convolution kernel of 1x1; the ATDE module can perform multi-level splicing of images and use 1x1 convolution to increase and then reduce the dimension to extract the spatial detail information of the image; the MixSPP module includes a convolution layer, a maximum pooling layer and a depthwise separable convolution; the MixSPP module parallelizes the spatial pyramid pooling and convolution modules, and performs layer-by-layer fusion; The constructed small target detection model is trained using the training set and the validation set to obtain a trained small target detection model based on the improved YOLOv8.

2. The small target detection method based on improved YOLOv8 according to claim 1, characterized in that: The small target image dataset includes the VisDrone2019 target detection dataset.

3. The small target detection method based on improved YOLOv8 according to claim 1, characterized in that: The small object image dataset contains small object images under different scenes, weather and lighting conditions.

4. The small target detection method based on improved YOLOv8 according to claim 1, characterized in that: The method also includes: dividing a test set from the small target image data set, and using the test set to analyze the detection accuracy, detection speed and parameter quantity of the trained small target detection model based on improved YOLOv8.

5. A small target detection system based on improved YOLOv8, characterized in that: include: The detection module is used to input the small target image to be detected into the pre-trained small target detection model based on improved YOLOv8 for recognition and obtain the detection result; The training method of the small target detection model based on the improved YOLOv8 includes: Obtain a small object image dataset and divide it into a training set and a validation set; Use the backbone network ATDeNet to replace the backbone network of YOLOv8 and build a small target detection model based on the improved YOLOv8; The backbone network ATDeNet includes an input layer Stem and four main parts Stage1, Stage2, Stage3 and Stage4; the input layer Stem includes a convolutional layer; the Stage1 includes a convolutional layer and an ATDE module; the Stage2 includes a convolutional layer and an ATDE module; the Stage3 includes a convolutional layer and an ATDE module; the Stage4 includes a convolutional layer and a pyramid structure module MixSPP; The convolution layers of the input layer and the backbone contain batch normalization BN and SiLU activation functions, and different convolution parameters are set; the ATDE module includes depthwise separable convolution and convolution with a convolution kernel of 1x1; the ATDE module can perform multi-level splicing of images and use 1x1 convolution to increase and then reduce the dimension to extract the spatial detail information of the image; the MixSPP module includes a convolution layer, a maximum pooling layer and a depthwise separable convolution; the MixSPP module parallelizes the spatial pyramid pooling and convolution modules, and performs layer-by-layer fusion; The constructed small target detection model is trained using the training set and the validation set to obtain a trained small target detection model based on the improved YOLOv8.

6. A computer storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the small target detection method based on improved YOLOv8 as described in any one of claims 1 to 5 is implemented.

Images