Small object detection method based on improved yolov8

Abstract

Disclosed in the present invention is a small object detection method based on improved YOLOv8. The small object detection method comprises: inputting, into a pre-trained small object detection model based on improved YOLOv8, a small object image to be subjected to detection for identification, so as to obtain a detection result. A training method for a small object detection model based on improved YOLOv8 comprises: acquiring a small object image dataset, and dividing same into a training set and a validation set; using a backbone network ATDeNet to replace a YOLOv8 backbone network, so as to construct a small object detection model based on improved YOLOv8; and using the training set and the validation set to train the constructed small object detection model, so as to obtain a trained small object detection model based on improved YOLOv8. The accuracy and efficiency of small object detection can be significantly improved.

Description

A small target detection method based on improved YOLOv8 TECHNICAL FIELD

[0001] The application belongs to the field of target detection of computer vision, and particularly relates to a small target detection method based on improved YOLOv8. BACKGROUND

[0002] With the rapid development of the unmanned aerial vehicle industry, unmanned aerial vehicles equipped with cameras have been deployed in various industries. The pictures transmitted by unmanned aerial vehicles generally have the characteristics of large field of view and rich information, and can provide information that ground shooting cannot provide. However, as the angle of view is pulled up, the size of the objects in the image is reduced, and the processing difficulty of the image information of the unmanned aerial vehicle is also increased compared with traditional images. The characteristics of the image of the unmanned aerial vehicle are small target size, less appearance and geometric information, complex background, and the like, which bring great challenges to target detection. Therefore, how to improve small target detection is an important topic in the field of computer vision.

[0003] Since the first launch of YOLO (You Only Look Once) algorithm in 2015, it has played an important role in the field of real-time target detection technology. After several iterations and upgrades, each version is committed to improving the detection efficiency, accuracy and practicality. YOLOv1 first adopted a whole prediction framework, abandoned the candidate region generation, and realized ultra-fast real-time detection speed, although the positioning accuracy was relatively insufficient, but with its unique advantages, it led a new direction of real-time target detection. Subsequently, in the evolution process of YOLOv2 to YOLOv8, the latest research results of deep learning are continuously absorbed and integrated, through multi-scale training, network structure adjustment, data enhancement technology and various innovative technologies such as weighted residual connection, cross-stage partial connection, etc., the detection accuracy is gradually improved and the high-speed running characteristics are maintained. Especially in the YOLOv8 stage, not only major breakthroughs have been made in speed and accuracy, but also more attention is paid to meeting the requirements of real-time and high classification performance in industrial scenarios, so that the YOLO series has been widely used in automated quality inspection and other related fields, and has become a model of progress in real-time target detection technology in the field of computer vision.

[0004] However, YOLOv8 still faces the problems of easy loss of feature information in deep network and insufficient feature fusion capability in the field of small target detection, which makes it difficult to guarantee the accuracy and efficiency of small target detection. SUMMARY

[0005] The purpose of the present application is to overcome the deficiencies in the prior art, provide a small target detection method based on improved YOLOv8, provide a new backbone network ATDeNet which fuses ATDE module and mixed spatial pyramid structure MixSPP, used to replace the original backbone network of YOLOv8, can alleviate the problem that small targets are easy to lose features in deep network, improve the accuracy and efficiency of small target detection.

[0006] To achieve the above purpose, the present application is realized by adopting the following technical scheme:

[0007] In a first aspect, the present application 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 a detection result;

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

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

[0011] Replace the backbone network of YOLOv8 with the backbone network ATDeNet to construct a small target detection model based on improved YOLOv8;

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

[0013] Further, the step of replacing the backbone network of YOLOv8 with the backbone network ATDeNet to construct a small target detection model based on improved YOLOv8 comprises:

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

[0015] The backbone network ATDeNet adopts a design of gradually reducing the size of a convolution kernel, specifically (7, 5, 3, 3), and a design of a convolution depth of (1, 1, 3, 1), which can increase the receptive field when starting detection while reducing the receptive field of the subsequent prediction feature map, thereby avoiding excessive redundant background information.

[0016] The first convolution module ATDE can splice the image in multiple levels, and extract spatial detail information of the image by using 1x1 convolution dimensionality reduction.

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

[0018] Further, the first convolution module ATDE includes a depth separable convolution and a convolution kernel of 1x1, so that the network can well capture the model input data 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, so that even if the response of the pyramid pool to the target is weak, the risk of overfitting can be reduced after the multiple 1x1 convolution is stacked after extracting the features, and the robustness of the model is improved.

[0020] Further, the convolution layer includes batch normalization BN and a SiLU activation function, and different convolution parameters are set.

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

[0022] Further, the small target image dataset includes small target images under different scenes, weather and light conditions.

[0023] Further, the method further includes: dividing a test set from the small target image dataset, and using the test set to analyze the detection precision, detection speed and parameter amount of the small target detection model based on the improved YOLOv8 after training.

[0024] In a second aspect, the present application provides a small target detection system based on an improved YOLOv8, which comprises:

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

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

[0027] A small target image dataset is acquired and divided into a training set and a verification set;

[0028] The backbone network ATDeNet is used to replace the backbone network of YOLOv8 to construct a small target detection model based on the improved YOLOv8.

[0029] The training set and the verification set are used to train the constructed small target detection model to obtain a trained small target detection model based on the improved YOLOv8.

[0030] In a third aspect, the present application provides a computer storage medium having a computer program stored thereon, wherein the computer program, when executed by a processor, implements the small target detection method based on the improved YOLOv8.

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

[0032] (1) The detection method proposed in the present application is based on the YOLOv8 target detection algorithm, which has high detection accuracy in normal environment. The backbone extraction network of the YOLOv8 is replaced by the backbone network ATDeNet proposed in the present application for small target detection, so that the model can effectively extract the spatial detail information of the image and improve the detection accuracy of the small target.

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

[0034] (3) The present application proposes a backbone network-ATDeNet for small targets, which 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) The present application proposes a new hybrid spatial pyramid structure-MixSPP, which parallelly connects the spatial pyramid pooling and the convolution module and adopts a layer-by-layer fusion mode to reduce the loss of multi-scale information in the fusion process. BRIEF DESCRIPTION OF DRAWINGS

[0036] Fig. 1 shows a structural schematic diagram of the backbone network ATDeNet;

[0037] Figs. 2 and 3 show structural schematic diagrams of the ATDE convolution module;

[0038] Figs. 4 and 5 show structural schematic diagrams of the hybrid pyramid MixSPP module;

[0039] Fig. 6 shows the overall network structure of the small target detection method based on improved YOLOv8 of the present application. DETAILED DESCRIPTION

[0040] The technical concept of the present application is to replace the backbone network of YOLOv8 with 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. Meanwhile, the hybrid spatial pyramid structure MixSPP is adopted, which parallelly connects the spatial pyramid pooling and convolution modules and adopts the layer-by-layer fusion mode to reduce the loss of multi-scale information in the fusion process.

[0041] The technical solutions in the embodiments of the present application will be described clearly and completely below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are only part of the embodiments of the present application, not all the embodiments. The description of the at least one exemplary embodiment is actually only illustrative, not as any limitation on the present application and its application or use.

[0042] Embodiment 1

[0043] The present embodiment discloses a small target detection method based on improved YOLOv8, comprising:

[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 a detection result.

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

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

[0047] Replace the backbone network of YOLOv8 with a backbone network ATDeNet (attention to detail network) to construct a small target detection model based on improved YOLOv8.

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

[0049] The small target detection method based on improved YOLOv8 provided in the present embodiment will be described below in combination with a preferred embodiment to explain the content involved in the above embodiments:

[0050] The software configuration environment of the 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 unmanned aerial vehicle dataset selected in the embodiment is the VisDrone2019 target detection dataset opened by a team of Tianjin University and others in 2019, which officially provides 6471 training sets and 548 validation sets, containing 10 types of targets. The dataset is collected using different models of unmanned aerial vehicles under different scenes, weather and lighting conditions, and can objectively reflect the performance of the detector.

[0052] S2: Construct the backbone network ATDeNet. Referring to the structural diagram of the backbone network ATDeNet shown in FIG. 1, it includes an input layer Stem and four backbone parts Stage1, Stage2, Stage3 and Stage4.

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

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

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

[0056] Referring to FIG. 1, the Stage2 and Stage3 are similar in structure to the Stage1 described above, and the only difference is the different convolution parameters, which will not be described again here.

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

[0058] Referring to FIGS. 2 and 3, the ATDE module is composed of n depth separable convolutions and three convolution kernels of 1x1, each of which is followed by a batch normalization layer and an activation function. The ATDE module splices the input image and the image after a depth separable convolution and n depth separable convolutions to form a new feature map to increase the expression ability of the feature map. Then, the spliced feature map is mapped to a high dimension by 1x1 convolution, where the dimension is (1+d)xn, and then restored to the original dimension by 1x1 convolution.

[0059] Referring to FIGS. 4 and 5, the MixSPP module is composed of 5 convolution layers, 3 max pooling layers, depth separable convolutions, etc. The MixSPP module uses a convolution layer with a convolution kernel size of 1x1, a stride of 1, no padding, and an output channel number of 512 to perform convolution operation on the image, and then uses batch normalization and SiLU activation function to preliminarily process the image; then, the image is divided into three regions of different sizes for max pooling operation. The convolution kernel sizes of the three regions are k13x13, k9x9, and k5x5, respectively, the strides are all 1, and there is no padding. The max pooling results of each region will be used as a separate feature map; then, these feature maps pass through a depth separable convolution layer (DSConv) composed of a depth convolution layer and a pointwise convolution layer. The parameters of this layer are k=3, s=1, p=1, c=512, indicating that a depth separable convolution layer with a convolution kernel size of 3x3, a stride of 1, padding of 1, and an output channel number of 512 is used. Similarly, this layer also includes batch normalization and SiLU activation function; finally, the feature map of the original input image is spliced with the feature map after the DSConv layer. The spliced feature map passes through three consecutive convolution layers, each of which has a convolution kernel size of 1x1, a stride of 1, no padding, and an output channel number of 1024, and is followed by batch normalization and SiLU activation function; finally, the spliced feature map passes through a last convolution layer with parameters k=1, s=1, p=0, c=512 to generate an output feature map. Finally, the spliced feature map passes through a last convolution layer with parameters k=1, s=1, p=0, c=512 to generate an output feature map. The MixSPP module can reduce the loss of multi-scale information in the fusion process by combining different size pooling operations and depth separable convolutions, even if the pyramid pool has a weak response to the target, and the stacking of multiple 1x1 convolutions after feature extraction can reduce the risk of overfitting and improve the robustness of the model.

[0060] S3: Replace the backbone network of YOLOv8 with the backbone network ATDeNet in step S2 to construct a small target detection model based on the improved YOLOv8.

[0061] S4: training the small target detection model built in step S3. The training set and the validation set in the data set in step S1 are input into the built small target detection model based on the improved YOLOv8 for training, and a trained model is obtained.

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

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

[0064] where λ cls is the classification loss weight, indicating the classification loss weight; λ box is the box loss weight, indicating the frame loss weight; λ dfl is the distribution focal loss weight, indicating the distribution focal loss weight;

[0065] q x,y , ν x,y , α x,y , and The specific formulae are as follows:

[0066] where N pos is the total number of grid points containing target objects; is an indicator function indicating which grid points contain objects; β x,y and represent the coordinates, width and height information of the real boundary box and the predicted boundary box, respectively, so the value range of q x,y is from 0 to 1, and the value closer to 1 indicates that the matching degree of the predicted frame and the real frame is higher, and the overlapping area is larger; b x,y is the coordinate representation of the center point of the real boundary box; y c is the true value label of the input image for each grid point (x, y) belonging to the category c, regardless of whether there is an object in the grid point; w x,y and h x,y represent the width and height of the corresponding boundary box, respectively; ν x,y and α x,y are auxiliary variables for adjusting the loss related to the aspect ratio; The predicted probability of the class, generated by the sigmoid activation function; The objectness score output by the softmax function; And Respectively represent the width and height of the predicted bounding box, when they are used as a ratio, this value reflects the shape information of the predicted box, that is, its aspect ratio.

[0067] Post-processing is performed using the Soft-NMS method, which is an improved version of the traditional NMS method. Instead of simply discarding overlapping bounding boxes, Soft-NMS applies a soft threshold penalty to reduce the confidence of overlapping regions, which can better preserve bounding boxes that may still be valid detection results despite some degree of overlap, thereby improving detection accuracy and reducing redundant detection;

[0068] According to the set training parameters and processing methods, train the small target detection model based on the algorithm;

[0069] The trained small target detection model based on the improved YOLOv8 is evaluated for performance, and a well-converged network model is obtained.

[0070] S5: Evaluate the model. According to the small target detection model based on the improved YOLOv8 obtained after training, evaluate the average detection accuracy and detection speed of the model.

[0071] By setting the corresponding parameters, input the sample images in the training set into the small target detection model based on the improved YOLOv8, set the number of iterations, and train the best weight file for small target detection. The loss function is used to continuously adjust the direction of network training, the map value of the validation set is calculated to verify whether the training achieves the expected effect, the best weight file is selected by screening the trained weight file, and it is loaded into the small target detection model based on the improved YOLOv8 as the model for small target detection in the unmanned aerial vehicle scene.

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

[0073] Referring to FIG. 6, in this example, under the software configuration environment of 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, ATDeNet has higher average precision under the YOLOv8 framework compared to CSPDarkNet, especially when processing small and medium-sized targets. At the same time, ATDeNet has less parameter quantity, which means it may have better computational efficiency. In addition, although the FPS of ATDeNet is slightly higher, the gap is not large, so it may not have a significant impact in actual application. Overall, ATDeNet is a better choice under the YOLOv8 framework, especially when it is necessary to improve the accuracy of small target detection.

[0074] Embodiment 2

[0075] This embodiment discloses a small target detection system based on improved YOLOv8 on the basis of embodiment 1, comprising:

[0076] a detection module for inputting a small target image to be detected into a pre-trained small target detection model based on improved YOLOv8 for recognition to obtain a detection result;

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

[0078] obtain a small target image dataset and divide it into a training set and a validation set;

[0079] use the backbone network ATDeNet to replace the backbone network of YOLOv8 to construct a small target detection model based on improved YOLOv8;

[0080] use the training set and the validation set to train the constructed small target detection model to obtain a trained small target detection model based on improved YOLOv8.

[0081] Embodiment 3

[0082] This embodiment introduces a computer readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the small target detection method based on improved YOLOv8 described in embodiment 1.

[0083] Those skilled in the art will appreciate that embodiments of the present application can be readily used as a method, a system or a computer program product. Accordingly, the present application can take the form of an entirely hardware embodiment, an entirely software embodiment or an embodiment combining software and hardware aspects. Furthermore, the present application can take the form of a computer program product on one or more computer-usable storage media (including, but not limited to, disk storage, CD-ROMs, optical storage devices, etc.) embodying computer-readable program code.

[0084] The present application is described in reference to the flowchart illustrations and / or block diagrams of the methods, apparatus (systems) and computer program products according to embodiments of the application. It will be understood that each block of the flowchart illustrations and / or block diagrams, and combinations of blocks in the flowchart illustrations and / or block diagrams, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general purpose computer, special purpose computer, embedded processing machine, or other programmable data processing apparatus to produce a machine, such that the instructions, which execute via the processor of the computer or other programmable data processing apparatus, create means for implementing the functions specified in the flowchart illustrations and / or block diagrams block or blocks.

[0085] These computer program instructions can also be stored in a computer- readable memory that can direct a computer or other programmable data processing apparatus to function in a particular manner, such that the instructions stored in the computer-readable memory produce an article of manufacture including instructions which implement the function specified in the flowchart illustrations and / or block diagrams block or blocks.

[0086] The above description is only preferred embodiments of the present application. It is understood that, for those skilled in the art, some improvements and modifications can be made without departing from the technical principles of the present application, and these improvements and modifications should also be considered as within the scope of the present application.

Claims

1. A small target detection method based on improved YOLOv8, characterized in that, The method comprises the following steps: inputting a small target image to be detected into a pre-trained small target detection model based on improved YOLOv8 for recognition to obtain a detection result; The training method of the small target detection model based on improved YOLOv8 comprises the following steps: obtain a small target image dataset and divide it into a training set and a validation set; replace the backbone network of YOLOv8 with a backbone network ATDeNet to construct a small target detection model based on improved YOLOv8; train the constructed small target detection model using the training set and the validation set to obtain a trained small target detection model based on improved YOLOv8.

2. The small target detection method based on improved YOLOv8 according to claim 1, wherein the step of replacing the backbone network of YOLOv8 with a backbone network ATDeNet to construct a small target detection model based on improved YOLOv8 comprises the following steps: construct a backbone network ATDeNet; the backbone network ATDeNet comprises 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 comprises a convolution layer; the first backbone part Stage1 comprises a convolution layer and an ATDE module; the second backbone part Stage2 comprises a convolution layer and an ATDE module; the third backbone part Stage3 comprises a convolution layer and a first convolution module ATDE; the fourth backbone part Stage4 comprises a convolution layer and a pyramid structure module MixSPP; wherein the first convolution module ATDE can perform multi-level splicing on the image and utilize 1x1 convolution to increase dimension and then reduce dimension to extract spatial detail information of the image; the pyramid structure module MixSPP parallelly connects a spatial pyramid pooling and a convolution module and performs layer-by-layer fusion; replace the backbone network of YOLOv8 with the backbone network ATDeNet to construct a small target detection model based on improved YOLOv8.

3. The small target detection method based on improved YOLOv8 according to claim 2, wherein the first convolution module ATDE comprises a depth separable convolution and a convolution kernel with a size of 1x1.

4. The small target detection method based on improved YOLOv8 according to claim 2, wherein the pyramid structure module MixSPP comprises a convolution layer, a max pooling layer and a depth separable convolution.

5. The small target detection method based on improved YOLOv8 according to claim 2, wherein the convolution layers of the input layer and the backbone part contain batch normalization BN and SiLU activation functions, and different convolution parameters are set.

6. The small target detection method based on improved YOLOv8 according to claim 1, wherein the small target image dataset comprises a VisDrone2019 target detection dataset.

7. The small target detection method based on improved YOLOv8 according to claim 1, wherein ​ ​ ​ ​ ​ The small target image dataset contains small target images under different scenes, weather and lighting conditions.

8. The small target detection method based on improved YOLOv8 according to claim 1, wherein, The method further comprises: dividing a test set from the small target image dataset, and using the test set to analyze the detection precision, detection speed and parameter quantity of the trained small target detection model based on the improved YOLOv8.

9. A small target detection system based on improved YOLOv8, characterized in that, Comprise: The detection module is configured to input the small target image to be detected into the pre-trained small target detection model based on the improved YOLOv8 for identification to obtain a detection result. The training method of the small target detection model based on the improved YOLOv8 comprises: Obtain a small target image dataset and divide it into a training set and a validation set; Replace the backbone network of YOLOv8 with the backbone network ATDeNet to construct a small target detection model based on the improved YOLOv8; Train the constructed small target detection model using the training set and the validation set to obtain a trained small target detection model based on the improved YOLOv8.

10. A computer storage medium having stored thereon a computer program, characterized in that The computer program is executed by the processor to implement the small target detection method based on the improved YOLOv8 as claimed in any one of claims 1-8.

Images