Target detection model training method
By constructing the improved YOLOv8 neural network structure and using sparse matrix module to fusion feature maps, combining CIoU and PIoU to calculate the loss value, the problem of difficulty in extracting small target features is solved, and the detection accuracy and efficiency are improved.
Patent Information
- Application Number
- CN202411826968.X
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-12-12
- Publication Date
- 2025-03-18
AI Technical Summary
In the prior art, it is difficult for the object detection model to effectively extract and recognize small target features during training, resulting in low recognition accuracy and accuracy.
By annotating and filtering images with small targets, and building an improved YOLOv8 neural network structure, the sparse matrix module is used to fuse the feature maps output from the backbone network, and the bounding box loss value is calculated by combining CIoU and PIoU to improve the model's ability to extract small target features.
The extraction ability and detection accuracy of the target detection model for small target features is improved, the model's detection ability of small target images is enhanced, and the prediction efficiency of the model is improved on the basis of ensuring the prediction effect.
Smart Images
Figure CN119649146A_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the technical field of image processing and recognition, and in particular to a training method for a target detection model. Background Art
[0002] In the field of computer vision and machine learning, the intersection over union (IoU) is widely used in object detection and image segmentation tasks as an important indicator for evaluating model performance. IoU measures the accuracy of the model by calculating the ratio of the intersection to the union of the predicted box and the true box.
[0003] With the in-depth study of image detection and recognition, it was found that using IoU to evaluate images containing small objects during model training will cause a large number of potential small objects to be misjudged or ignored, resulting in low recognition accuracy and precision of the trained model.
[0004] At the same time, since the loss function in the existing training method mainly uses IoU as the evaluation criterion, the trained model cannot effectively capture and represent the features of small targets with pixels within 32*32 when performing feature extraction, resulting in the accuracy of the trained target detection model always being unable to meet actual application needs.
[0005] Therefore, there is an urgent need for a technical solution for a model training method that can improve the target detection accuracy of the model. Summary of the invention
[0006] In view of the above analysis, an embodiment of the present invention aims to provide a method for training a target detection model, so as to solve the problem that the neural network model trained by the training method in the prior art is difficult to extract small target features and has low recognition accuracy.
[0007] An embodiment of the present invention provides a method for training a target detection model, and the training method includes:
[0008] Step S1: annotate a large number of historical images and select multiple images with small objects, the selected multiple annotated images form a sample set, and the sample set is divided into a training set, a verification set and a test set;
[0009] Step S2: construct an improved YOLOv8 neural network structure;
[0010] Step S3: using the training set to train the improved YOLOv8 neural network structure for multiple rounds, and at the end of each round, using the validation set to evaluate the improved YOLOv8 neural network structure after the training of the round, to obtain the evaluation result of each round; after the training is completed, comparing the evaluation results of each round, and obtaining the improved YOLOv8 neural network structure of the round corresponding to the optimal evaluation result as the image detection model;
[0011] Step S4: Use the test set to test the image detection model. If the test result is qualified, save the image detection model and end the training; if it is unqualified, return to step S3 to continue training.
[0012] Based on the further improvement of the above training method, a large number of historical images are annotated and multiple images with small objects are screened out, and the multiple annotated images screened out form a sample set, including:
[0013] Annotate the bounding box of the target in each historical image, determine the center point coordinates, height, width, area and target category of each bounding box in each historical image, and the total width, total height and total area of each historical image;
[0014] According to the height, width and area of each bounding box, combined with the total width, total height and total area of the corresponding historical image, the height ratio, width ratio and area ratio of each target in the corresponding historical image are calculated; if there is at least one target in the historical image whose height ratio is less than the height ratio threshold, or the width ratio is less than the width ratio threshold, or the area ratio is less than the area ratio threshold, then the historical image is taken as a sample in the sample set.
[0015] Based on the further improvement of the above training method, the improved YOLOv8 neural network structure is trained for multiple rounds using the training set, including:
[0016] In each round, multiple samples are extracted from the training set as a batch of samples, and each batch of samples is input into the improved YOLOv8 neural network structure;
[0017] According to the samples of each batch and the corresponding prediction results, the loss function value of the batch is calculated; according to the loss function value of the batch, the improved YOLOv8 neural network structure is updated through back propagation;
[0018] The updated improved YOLOv8 neural network structure is trained using the next batch of samples until the training of the preset batch is completed to obtain the improved YOLOv8 neural network structure of the current round;
[0019] The improved YOLOv8 neural network structure of the current round is trained for the next round until the preset round of training is completed.
[0020] Based on the further improvement of the above training method, the loss function value of each batch is calculated according to the samples of each batch and the corresponding prediction results, including:
[0021] Calculate the category loss value and bounding box loss value of all targets of all samples in the batch of samples; wherein each sample includes at least one target;
[0022] The average category loss value is calculated based on the category loss value of each target, and the average bounding box loss value is calculated based on the bounding box loss value of each target, thereby obtaining the loss function value of the batch.
[0023] Based on the further improvement of the above training method, the bounding box loss value of each target is calculated by the following formula:
[0024]
[0025]
[0026] in, represents the bounding box loss function, Sigmoid() represents the activation function, S t Represents the pixel area of the real border, S st represents the preset border pixel area threshold for distinguishing whether it is a small target, b represents the preset ratio parameter, (1-PIoU) represents the first border loss function, (1-CIoU) represents the second border loss function, and x t ,y t represents the true center coordinate, x p ,y p represents the predicted center coordinates, SD represents the average area of the bounding boxes of all objects in all training samples in each batch, A represents the true bounding box, B represents the predicted bounding box, l1 represents the distance between the center coordinates of the true bounding box and the predicted bounding box, l2 represents the length of the diagonal of the minimum circumscribed rectangular box of the true bounding box and the predicted bounding box, and w A 、h A Indicates the width and height of the actual border, w B 、h B represents the width and height of the predicted bounding box, and v represents the aspect ratio consistency factor function.
[0027] Based on the further improvement of the above training method, the category loss value of each target is calculated by the following formula:
[0028]
[0029] in, represents the category loss function, N represents the total number of categories, i represents the i-th category, and Y i represents the true category value of the i-th category, P i Represents the predicted category value of the i-th category.
[0030] Based on the further improvement of the above training method, the loss function value of each batch is calculated by the following formula:
[0031]
[0032] u1+u2=1;
[0033] in, represents the loss function of each batch, u1 represents the category weight coefficient, u2 represents the bounding box weight coefficient, represents the average loss value of the category, Represents the average loss value of the bounding box.
[0034] Based on the further improvement of the above training method, the neck network in the improved YOLOv8 neural network structure includes a sparse matrix module and an original neck network;
[0035] The original neck network is used to fuse other feature maps output by the backbone network except for the shallowest feature map;
[0036] The sparse matrix module is used to fuse the shallowest feature map and the deepest feature map output by the backbone network. The fused feature map and the fused feature map obtained by the original neck network are further fused in the neck network and then output.
[0037] Based on the further improvement of the above training method, the improved YOLOv8 neural network structure also includes:
[0038] The backbone network is used to extract features from the input samples, and input the obtained second feature map, third feature map, fourth feature map and fifth feature map from shallow to deep into the neck network;
[0039] The neck network also includes a dynamic weighted fusion module; the sparse matrix module fuses the second feature map and the fifth feature map to output a tenth feature map, and the original neck network fuses the third to fifth feature maps to obtain an eleventh feature map and a twelfth feature map;
[0040] The dynamic weighted fusion module is used to fuse the tenth feature map and the twelfth feature map to output a thirteenth feature map;
[0041] The original neck network is also used to further fuse the thirteenth feature map, the eleventh feature map and the fifth feature map to obtain the sixth feature map, the seventh feature map and the eighth feature map, and input the sixth feature map, the seventh feature map and the eighth feature map into the detection head network;
[0042] The detection head network is used to obtain a predicted target result of the sample according to the sixth feature map, the seventh feature map and the eighth feature map; wherein the predicted target result includes a predicted target category and a predicted target bounding box.
[0043] Based on the further improvement of the above training method, the sparse matrix module includes a double upsampling layer, a nonlinear activation layer, a first upsampling layer, a multiplication module and a first convolution module;
[0044] A double upsampling layer, used to upsample the fifth feature map of the input twice, and input the twice upsampled feature map into the nonlinear activation layer;
[0045] A non-linear activation layer is used to perform non-linear activation on the input feature map to obtain a sparse matrix and input it into the first upsampling layer;
[0046] A first upsampling layer, used for upsampling the sparse matrix to obtain a ninth feature map and inputting it into a multiplication module;
[0047] A multiplication module, used for multiplying the input second feature map and the ninth feature map, and inputting the multiplied feature map into the first convolution module;
[0048] The first convolution module is used to downsample and nonlinearly represent the input feature map to obtain a tenth feature map.
[0049] Compared with the prior art, the present invention can achieve at least one of the following beneficial effects:
[0050] 1. By marking and screening multiple images with small targets, training sets, validation sets and test sets are formed. The improved YOLOv8 neural network structure is trained with the training set. The improved YOLOv8 neural network structure of each round is evaluated with the validation set. The obtained image detection model is tested with the test set. In the improved YOLOv8 neural network structure, the sparse matrix module fuses the shallowest feature map and the deepest feature map output by the backbone network to improve the extraction of small target features. At the same time, only the feature map area where the target may exist is calculated during the fusion. The efficiency of the model prediction is improved as much as possible while ensuring the model prediction effect, so that the model can quickly detect and identify images.
[0051] 2. CIoU and PIoU are used to calculate the bounding box loss value of each target, so that the improved YOLOv8 neural network structure can learn more small target features and improve the detection ability of small target images;
[0052] 3. By training the improved YOLOv8 neural network structure, the learnable parameters corresponding to the tenth feature map and the twelfth feature map are obtained. The learnable parameters can be adaptively adjusted by fully considering the border sizes of different images. It is not only suitable for the detection of large target images, but also more conducive to improving the sensitivity of small target image detection.
[0053] In the present invention, the above-mentioned technical solutions can also be combined with each other to achieve more preferred combination solutions. Other features and advantages of the present invention will be described in the subsequent description, and some advantages can become obvious from the description, or can be understood by practicing the present invention. The purpose and other advantages of the present invention can be realized and obtained through the contents particularly pointed out in the description and the drawings. BRIEF DESCRIPTION OF THE DRAWINGS
[0054] The drawings are only for the purpose of illustrating particular embodiments and are not to be considered limiting of the present invention. Like reference symbols denote like components throughout the drawings.
[0055] Figure 1 A schematic diagram of a flow chart of a method for training a target detection model provided by an embodiment of the present invention;
[0056] Figure 2 A schematic diagram of the structure of an improved YOLOv8 neural network structure provided by an embodiment of the present invention;
[0057] Figure 3 A schematic diagram of the structure of a sparse matrix module provided in an embodiment of the present invention;
[0058] Figure 4 A schematic diagram of the structure of a detection head and a post-processing module provided in an embodiment of the present invention. DETAILED DESCRIPTION
[0059] The preferred embodiments of the present invention are described in detail below in conjunction with the accompanying drawings, wherein the accompanying drawings constitute a part of this application and are used together with the embodiments of the present invention to illustrate the principles of the present invention, but are not used to limit the scope of the present invention.
[0060] A specific embodiment of the present invention discloses a method for training a target detection model, such as Figure 1 As shown, the training method comprises:
[0061] Step S1: annotate a large number of historical images and select multiple images with small objects, the selected multiple annotated images form a sample set, and the sample set is divided into a training set, a verification set and a test set;
[0062] Step S2: construct an improved YOLOv8 neural network structure;
[0063] Step S3: using the training set to train the improved YOLOv8 neural network structure for multiple rounds, and at the end of each round, using the validation set to evaluate the improved YOLOv8 neural network structure after the training of the round, to obtain the evaluation result of each round; after the training is completed, comparing the evaluation results of each round, and obtaining the improved YOLOv8 neural network structure of the round corresponding to the optimal evaluation result as the image detection model;
[0064] Step S4: Use the test set to test the image detection model. If the test result is qualified, save the image detection model and end the training; if it is unqualified, return to step S3 to continue training.
[0065] Specifically, Figure 1 As shown, in step S1, a large number of historical images are obtained, and each historical image is labeled and screened to obtain multiple images with small targets. It can be understood that there may be targets or no targets in the historical images, so it is necessary to delete the historical images that do not include targets.
[0066] Preferably, the step of labeling a large number of historical images and selecting a plurality of images with small objects, wherein the selected plurality of labeled images form a sample set, comprises:
[0067] Annotate the bounding box of the target in each historical image, determine the center point coordinates, height, width, area and target category of each bounding box in each historical image, and the total width, total height and total area of each historical image;
[0068] According to the height, width and area of each bounding box, combined with the total width, total height and total area of the corresponding historical image, the height ratio, width ratio and area ratio of each target in the corresponding historical image are calculated; if there is at least one target in the historical image whose height ratio is less than the height ratio threshold, or the width ratio is less than the width ratio threshold, or the area ratio is less than the area ratio threshold, then the historical image is taken as a sample in the sample set.
[0069] Specifically, when annotating historical images, LabelImg and LabelMe can be used to annotate multiple bounding boxes of multiple targets included in each historical image, thereby determining the center point coordinates, height, width, area and target category of each bounding box in each historical image, and at the same time determining the total width, total height and total area of each historical image.
[0070] Specifically, the height ratio of each target in the corresponding historical image is calculated by the height of each bounding box and the total height of the corresponding historical image, the width ratio of each target in the corresponding historical image is calculated by the width of each bounding box and the total width of the corresponding historical image, and the area ratio of each target in the corresponding historical image is calculated by the area of each bounding box and the total area of the corresponding historical image.
[0071] Specifically, when screening historical images, determine whether there is at least one target in each historical image whose height ratio is less than a height ratio threshold, or whose width ratio is less than a width ratio threshold, or whose area ratio is less than an area ratio threshold; if so, take the historical image as a sample.
[0072] It is worth mentioning that before annotating historical images, they can be preprocessed. The preprocessing methods can be grayscale, geometric transformation, image enhancement and filtering, etc., which help to eliminate irrelevant information, restore real information, enhance the detectability of relevant information, and simplify the data to the maximum extent, thereby improving the reliability of feature extraction, image segmentation, matching and recognition, and obtaining historical images in a set format.
[0073] Specifically, Figure 1 As shown, in step S1, after obtaining a sufficient number of samples, the samples can be divided into a training set, a validation set, and a test set according to a preset allocation ratio, for example, according to a ratio of 7:1.5:1.5.
[0074] Specifically, Figure 1 As shown, in step S2, an improved YOLOv8 neural network structure is constructed.
[0075] Preferably, if Figure 2 As shown, the neck network in the improved YOLOv8 neural network structure includes a sparse matrix module and an original neck network;
[0076] The original neck network is used to fuse other feature maps output by the backbone network except for the shallowest feature map;
[0077] The sparse matrix module is used to fuse the shallowest feature map and the deepest feature map output by the backbone network. The fused feature map and the fused feature map obtained by the original neck network are further fused in the neck network and then output.
[0078] Specifically, Figure 2 As shown in the figure, the improved YOLOv8 neural network structure adds a sparse matrix module on the basis of the original YOLOv8 neural network structure. The added sparse matrix module fuses the shallowest feature map and the deepest feature map output by the backbone network, and the original neck network fuses other feature maps except the shallowest feature map. The fused feature map obtained by the original neck network and the fused feature map obtained by the sparse matrix module are further fused in the neck network and output.
[0079] Preferably, if Figure 2 As shown, the improved YOLOv8 neural network structure also includes:
[0080] The backbone network is used to extract features from the input samples, and input the obtained second feature map, third feature map, fourth feature map and fifth feature map from shallow to deep into the neck network;
[0081] The neck network also includes a dynamic weighted fusion module; the sparse matrix module fuses the second feature map and the fifth feature map to output a tenth feature map, and the original neck network fuses the third to fifth feature maps to obtain an eleventh feature map and a twelfth feature map;
[0082] The dynamic weighted fusion module is used to fuse the tenth feature map and the twelfth feature map to output a thirteenth feature map;
[0083] The original neck network is also used to further fuse the thirteenth feature map, the eleventh feature map and the fifth feature map to obtain the sixth feature map, the seventh feature map and the eighth feature map, and input the sixth feature map, the seventh feature map and the eighth feature map into the detection head network;
[0084] The detection head network is used to obtain a predicted target result of the sample according to the sixth feature map, the seventh feature map and the eighth feature map; wherein the predicted target result includes a predicted target category and a predicted target bounding box.
[0085] Specifically, Figure 2 As shown, the sample is input into the backbone network of the improved YOLOv8 neural network structure, and features of the sample are extracted in the backbone network to obtain the first feature map P1, the second feature map P2, the third feature map P3, the fourth feature map P4 and the fifth feature map P5. For example, Figure 2As shown, a 640*640*3 sample is input into the backbone network, and a 320*320*64 feature map is generated through the convolution module ConvModule as the first feature map P1; the first feature map P1 is passed through the convolution module ConvModule to generate a 160*160*128 feature map, and then passed through the residual convolution module CSPLayer to generate a 160*160*128 feature map as the second feature map P2; the second feature map P2 is passed through the convolution module ConvModule and the residual convolution module CSPLayer to generate a 80*80*256 feature map as the third feature map P3; the third feature map is passed through the convolution module ConvModule and the residual convolution module CSPLayer to generate a 40*40*512 feature map as the fourth feature map P4; the fourth feature map P4 is passed through the convolution module ConvModule, the residual convolution module CSPLayer and the spatial pyramid pooling module SPPF to generate a 20*20*512 feature map as the fifth feature map.
[0086] It can be understood that SPPF is fixed-space pyramid pooling, which further processes the deep feature map through multi-layer pooling and splicing operations without changing the shape of the input feature map.
[0087] Specifically, Figure 2 As shown, the backbone network inputs the obtained second feature map P2, third feature map P3, fourth feature map P4 and fifth feature map P5 into the neck network.
[0088] Specifically, the neck network includes an original neck network, a sparse matrix module and a dynamic weighted fusion module; the original neck network receives the third feature map P3, the fourth feature map P4 and the fifth feature map P5, and fuses them to obtain the eleventh feature map P11 and the twelfth feature map P12; the sparse matrix module receives the second feature map P2 and the fifth feature map P5 output by the backbone network and fuses them to obtain the tenth feature map P10; the dynamic weighted fusion module receives the tenth feature map P10 output by the sparse matrix module and the twelfth feature map P12 output by the original neck network, and fuses them to obtain the thirteenth feature map P13.
[0089] It is worth mentioning that the shallow high-resolution feature map P2 can effectively extract the features of small targets, and at the same time, the sparse matrix can be used to reduce the amount of calculation and achieve faster acquisition of the features of small targets.
[0090] Specifically, in the original neck network, the thirteenth feature map P13, the eleventh feature map P11 and the fifth feature map P5 are fused to obtain the sixth feature map P6, the seventh feature map P7 and the eighth feature map P8, and the sixth feature map P6, the seventh feature map P7 and the eighth feature map P8 are input into the detection head network.
[0091] Specifically, in the detection head network, the predicted target result corresponding to the sample is obtained according to the sixth feature map P6, the seventh feature map P7 and the eighth feature map P8; the predicted target result includes the predicted target category and the predicted target bounding box, that is, the category and bounding box of each target in the sample.
[0092] Preferably, if Figure 3 As shown, the sparse matrix module includes a double upsampling layer, a nonlinear activation layer, a first upsampling layer, a multiplication module and a first convolution module;
[0093] A double upsampling layer, used to upsample the fifth feature map of the input twice, and input the twice upsampled feature map into the nonlinear activation layer;
[0094] A non-linear activation layer is used to perform non-linear activation on the input feature map to obtain a sparse matrix and input it into the first upsampling layer;
[0095] A first upsampling layer, used for upsampling the sparse matrix to obtain a ninth feature map and inputting it into a multiplication module;
[0096] A multiplication module, used for multiplying the input second feature map and the ninth feature map, and inputting the multiplied feature map into the first convolution module;
[0097] The first convolution module is used to downsample and nonlinearly represent the input feature map to obtain a tenth feature map.
[0098] Specifically, Figure 3 As shown, the fifth feature map P5 is input into the sparse matrix module, and passes through the double upsampling layer, the nonlinear activation layer, and the first upsampling layer in sequence to obtain the ninth feature map P9; the second feature map P2 input into the sparse matrix module and the ninth feature map are multiplied point by point in the multiplication module, and the multiplied feature map is output to the first convolution module ConvModule1 to obtain the tenth feature map P10.
[0099] For example, Figure 3As shown, the fifth feature map P5 is 20*20*512, and is upsampled twice in the double upsampling layer to obtain a feature map 80*80*512, which only changes the size of the feature map, and does not change the number of channels of the feature map; the feature map 80*80*512 is nonlinearly activated in the nonlinear activation layer to obtain a feature map 80*80*512 as a sparse matrix; the sparse matrix 80*80*512 is input to the first upsampling layer for upsampling to obtain 160*160*128, as the ninth feature map P9; the second feature map P2 is 160*160*128, and the ninth feature map P9 and the second feature map P2 are multiplied point by point in the multiplication module to obtain a feature map 160*160*128; the feature map 160*160*128 after point-by-point multiplication is downsampled and nonlinearly represented in the first convolution module ConvModule1 to obtain 80*80*768, as the tenth feature map P10.
[0100] Preferably, if Figure 2 As shown, the original neck network is used to upsample the fifth feature map, and the upsampled fifth feature map is concatenated and fused with the fourth feature map. The concatenated and fused feature map is subjected to a residual convolution operation to obtain an eleventh feature map;
[0101] The eleventh feature map is upsampled, and the obtained feature map is concatenated and fused with the third feature map to obtain the twelfth feature map.
[0102] For example, Figure 2 As shown, the fifth feature map P5 is input to the original neck network, and is upsampled by the second upsampling layer in the original neck network to obtain a feature map of 40*40*512, and then passes through splicing 2, splicing 2 receives the fourth feature map P4 input by the backbone network, and in splicing 2, the fourth feature map P4 is spliced with the feature map 40*40*512 after the upsampling layer to obtain a feature map 40*40*1024, that is, the size remains unchanged in the splicing layer, and the channels are spliced. The feature map 40*40*1024 obtained by splicing 2 is input to the residual convolution layer CSPLayer for residual convolution operation to obtain a feature map 40*40*512 as the eleventh feature map P11. The eleventh feature map is upsampled by the second upsampling layer to obtain a feature map of 80*80*512, and then passes through splicing 1. Splicing 1 also receives the third feature map P3 output by the backbone network. In splicing 3, the feature map 80*80*512 and the third feature map P3 are spliced and fused to obtain a feature map of 80*80*768 as the twelfth feature map.
[0103] Preferably, if Figure 2 As shown, the dynamic weighted fusion module is used to fuse the tenth feature map and the twelfth feature map to output a thirteenth feature map, including:
[0104] Obtaining a first learnable parameter corresponding to the tenth feature map and a second learnable parameter corresponding to the twelfth feature map; wherein the sum of the first learnable parameter and the second learnable parameter is 1;
[0105] A thirteenth feature map is obtained by summing a product of the tenth feature map and the first learnable parameter and a product of the twelfth feature map and the second learnable parameter, wherein the first learnable parameter and the second learnable parameter are obtained through network training.
[0106] Specifically, Figure 2 As shown, the tenth feature map P10 and the twelfth feature map P12 are fused in the dynamic weighted fusion module to obtain the thirteenth feature map P13.
[0107] For example, Figure 2 As shown, the first learnable parameter α and the second learnable parameter (1-α) are used as training parameters of the improved YOLOv8 neural network structure, the tenth feature map P10 is point-by-point multiplied by the first learnable parameter α, the twelfth feature map P12 is point-by-point multiplied by the second learnable parameter (1-α), and the tenth feature map P10 and the twelfth feature map P12 after the point-by-point multiplication are added and summed to obtain a feature map 80*80*768, which is used as the thirteenth feature map P13.
[0108] Preferably, if Figure 2 As shown, in the original neck network,
[0109] The thirteenth feature map is subjected to a residual convolution operation to obtain the sixth feature map.
[0110] Preferably, if Figure 2 As shown, in the original neck network,
[0111] The sixth feature map is downsampled and nonlinearly represented by the second convolution module, and the obtained feature map is spliced with the eleventh feature map. The spliced feature map is subjected to a residual convolution operation to obtain the seventh feature map.
[0112] Preferably, if Figure 2 As shown, in the original neck network,
[0113] The seventh feature map is downsampled and nonlinearly represented by the second convolution module, and the obtained feature map is spliced with the fifth feature map. The spliced feature map is subjected to a residual convolution operation to obtain the eighth feature map.
[0114] Specifically, Figure 2As shown, in the original neck network, the thirteenth feature map P13 is received from the dynamic weighted fusion module, and the thirteenth feature map P13 is subjected to residual convolution operation by the residual convolution module CSPLayer to obtain a feature map 80*80*256 as the sixth feature map P6; the sixth feature map P6 is subjected to the second convolution module ConvModule2 to obtain a feature map 40*40*256, which is output to splicing 3; the eleventh feature map P11 is received in splicing 3, and is spliced and fused with the feature map 40*40*256 to obtain a feature map 40*40*768, which is then subjected to the residual convolution layer CS PLayer obtains a feature map of 40*40*512 as the seventh feature map P7; the seventh feature map P7 passes through the second convolution module ConvModule2 to obtain a feature map of 20*20*512, which is input to splicing 4. Splicing 4 receives the fourth feature map P4 output by the backbone network. In splicing 4, the fifth feature map P5 and the feature map 20*20*512 are spliced and fused to obtain a feature map of 20*20*1024. The feature map 20*20*1024 passes through the residual convolution module CSPLayer to obtain a feature map of 20*20*512 as the eighth feature map P8. The sixth feature map P6, the seventh feature map and the eighth feature map P8 are output to the detection head network and predicted in the detection head network.
[0115] It is worth noting that in the neck network, the second convolution module ConvModule2 downsamples and nonlinearly represents the input feature map, the residual convolution module performs residual convolution operation on the input feature map, the second upsampling layer upsamples the input feature map, splicing 1, splicing 2 and splicing 3 respectively splice and fuse the two input feature maps, and finally obtain the sixth feature map P6, the seventh feature map P7 and the eighth feature map P8.
[0116] Preferably, if Figure 2 As shown, the detection head network includes a first detection head, a second detection head, a third detection head and a post-processing module;
[0117] The first detection head, the second detection head, and the third detection head obtain corresponding prediction box results according to the sixth feature map, the seventh feature map, and the eighth feature map respectively;
[0118] Among them, each prediction box result includes the center point coordinates, width, height, confidence of each category and the probability of containing the target of the prediction box;
[0119] The post-processing module is used to filter the prediction box results to obtain the prediction target results.
[0120] Specifically, Figure 2As shown, the first detection head Head1 makes predictions based on the input sixth feature map P6, the second detection head Head2 makes predictions based on the input seventh feature map P7, and the third detection head Head3 makes predictions based on the input eighth feature map P8.
[0121] Specifically, Figure 4 As shown in the figure, the input feature map passes through the convolution module CBS and the convolution layer Conv2D in sequence. The convolution module CBS downsamples and nonlinearly represents the input feature map, and the convolution layer Conv2D performs a convolution operation on the input feature map to change the number of channels of the feature map.
[0122] Specifically, Head represents any one of the first detection head Head1, the second detection head Head2 or the third detection head Head3. Exemplarily, when the first detection head Head1 is represented, the frame prediction result 80*80*4 is predicted according to the sixth feature map P6, where 80*80 represents the predicted 6400 prediction boxes, 4 represents the center point coordinates, width and height of each prediction box, and the category prediction result 80*80*(cls_num+1) is predicted according to the sixth feature map P6, where 80*80 represents the 6400 prediction boxes, cls_num represents the confidence of each category, and +1 represents the probability of containing the target in each prediction box.
[0123] Specifically, the three detection heads, the first detection head, the second detection head and the third detection head, obtain a total of 6400+1600+400 prediction box results, each of which includes the center point coordinates, width, height, confidence of each category and the probability of containing the target. Figure 4 As shown, the prediction box result of 8400*(cls_num+5) is input into the post-processing module, and in the post-processing module, the prediction box result is screened to obtain the final prediction target result.
[0124] Preferably, the prediction box results are screened to obtain the predicted target results, including:
[0125] Delete the prediction boxes whose probability of containing the target in the prediction box results does not meet the target threshold;
[0126] For the remaining prediction boxes, traverse all combinations of two prediction boxes and calculate the intersection-over-union ratio of the two prediction boxes in each combination; when the intersection-over-union ratio is greater than the same target threshold, compare the highest category confidences of the two prediction boxes and delete the prediction box corresponding to the smaller highest category confidence; when the intersection-over-union ratio is less than or equal to the same target threshold, both prediction boxes are retained;
[0127] The category corresponding to the highest category confidence in the remaining prediction box is used as the predicted target category of the predicted target result, and the center point coordinates, width and height of the prediction box are used as the predicted target bounding box of the predicted target result.
[0128] Specifically, in the post-processing module, it is first determined whether the probability of each prediction box containing a target is greater than a preset target threshold. If it is greater than the target threshold, the prediction box is considered to contain a target. Otherwise, the prediction box is considered to not contain a target and the prediction box is deleted.
[0129] For the remaining prediction frames, the combination of any two prediction frames is judged to see whether the two prediction frames correspond to the same target. Specifically, the intersection-over-union ratio of the two prediction frames is calculated. If the intersection-over-union ratio is greater than the preset threshold for the same target, the two prediction frames are considered to correspond to the same target, otherwise they correspond to different targets. When the two prediction frames correspond to the same target, the category confidences corresponding to the two prediction frames are checked, and the prediction frame with the smaller highest category confidence is deleted, and the prediction frame with the larger highest category confidence is retained. That is, when the two prediction frames correspond to the same target, the prediction frame with the higher category confidence is retained.
[0130] Specifically, among the remaining prediction boxes after screening, each prediction box corresponds to the center point coordinates, width and height, which serves as the predicted target bounding box of the target; the category corresponding to the highest category confidence corresponding to each prediction box is used as the predicted target category of the target.
[0131] Specifically, Figure 1 As shown, in step S3, the improved YOLOv8 neural network structure is trained for multiple rounds using the training set, and the improved YOLOv8 neural network structure after the training is evaluated using the validation set at the end of each round to obtain the evaluation result of each round. For example, if 100 rounds of training are performed, 100 rounds of evaluation results are obtained, and 100 improved YOLOv8 neural network structures are obtained. The evaluation results of 100 rounds are compared, and the improved YOLOv8 neural network structure corresponding to the round with the best evaluation result is used as the image detection model.
[0132] It is understandable that the evaluation criteria are carried out according to the existing standards in the prior art and will not be described in detail here.
[0133] Preferably, the training of the improved YOLOv8 neural network structure using the training set for multiple rounds includes:
[0134] In each round, multiple samples are extracted from the training set as a batch of samples, and each batch of samples is input into the improved YOLOv8 neural network structure;
[0135] According to the samples of each batch and the corresponding prediction results, the loss function value of the batch is calculated; according to the loss function value of the batch, the improved YOLOv8 neural network structure is updated through back propagation;
[0136] The updated improved YOLOv8 neural network structure is trained using the next batch of samples until the training of the preset batch is completed to obtain the improved YOLOv8 neural network structure of the current round;
[0137] The improved YOLOv8 neural network structure of the current round is trained for the next round until the preset round of training is completed.
[0138] Specifically, in each round, multiple samples are randomly selected from the training set as a batch of samples, and then a batch of samples is input into the improved YOLOv8 neural network structure for training to obtain the prediction results of a batch of samples.
[0139] Specifically, after obtaining the prediction results of each batch, the loss function value of the batch is calculated according to the prediction results of each batch and the samples of each batch.
[0140] Preferably, the step of calculating the loss function value of each batch according to the samples of each batch and the corresponding prediction results includes:
[0141] Calculate the category loss value and bounding box loss value of all targets of all samples in the batch of samples; wherein each sample includes at least one target;
[0142] The average category loss value is calculated based on the category loss value of each target, and the average bounding box loss value is calculated based on the bounding box loss value of each target, thereby obtaining the loss function value of the batch.
[0143] Specifically, the bounding box loss value of each target is calculated by the following formula:
[0144]
[0145] in, represents the bounding box loss function, Sigmoid() represents the activation function, S t Represents the pixel area of the real border, S st represents the preset border pixel area threshold for distinguishing whether it is a small target, b represents the preset ratio parameter, (1-PIoU) represents the first border loss function, (1-CIoU) represents the second border loss function, and x t ,y t represents the true center coordinate, x p ,y prepresents the predicted center coordinates, SD represents the average area of the bounding boxes of all objects in all training samples in each batch, A represents the true bounding box, B represents the predicted bounding box, l1 represents the distance between the center coordinates of the true bounding box and the predicted bounding box, l2 represents the length of the diagonal of the minimum circumscribed rectangular box of the true bounding box and the predicted bounding box, and w A 、h A Indicates the width and height of the actual border, w B 、h B represents the width and height of the predicted bounding box, and v represents the aspect ratio consistency factor function.
[0146] Specifically, PIoU enables the improved YOLOv8 neural network structure to learn more small target features, and calculates the relative distance between the predicted target center point coordinates and the real target center point coordinates according to the size of the target. This allows the improved YOLOv8 neural network structure to calculate the loss of all small targets during the training phase, thereby optimizing the ability of the improved YOLOv8 neural network structure to learn small target features, and effectively solves the problem of being unable to effectively calculate the loss of small targets due to small targets, and ensures that each real bounding box can find a matching predicted bounding box during the training process.
[0147] Specifically, the proportion of the first border loss function and the second border loss function is adjusted by a preset ratio parameter b. The larger the b is, the greater the proportion of the first border loss function is.
[0148] Specifically, the category loss value of each target is calculated by the following formula:
[0149]
[0150] in, represents the category loss function, N represents the total number of categories, i represents the i-th category, and Y i represents the true category value of the i-th category, P i Represents the predicted category value of the i-th category.
[0151] Specifically, after calculating the category loss values and bounding box loss values of all targets, the average values are calculated respectively to obtain the category average loss value and the bounding box average loss value, and the loss function value of the batch is calculated based on the category average loss value and the bounding box average loss value.
[0152] Specifically, the loss function value of each batch is calculated by the following formula:
[0153]
[0154] u1+u2=1;
[0155] in, represents the loss function of each batch, u1 represents the category weight coefficient, u2 represents the bounding box weight coefficient, represents the average loss value of the category, Represents the average loss value of the bounding box.
[0156] Specifically, after calculating the loss function value of the batch, the improved YOLOv8 neural network structure is updated through back propagation according to the loss function value of the batch, and the weight parameters involved in the improved YOLOv8 neural network structure are updated, which will not be repeated here.
[0157] Specifically, the improved YOLOv8 neural network structure updated in the current batch is trained using the samples of the next batch until all the samples of the preset batches have completed the training, thereby obtaining the improved YOLOv8 neural network structure of the current round.
[0158] Specifically, after obtaining the improved YOLOv8 neural network structure of the current round, the improved YOLOv8 neural network structure after round training is evaluated through the validation set to obtain the evaluation result of the round.
[0159] The next round of training is started based on the improved YOLOv8 neural network structure of the current round, until the preset round of training is completed and the evaluation result of the preset round is obtained.
[0160] Specifically, Figure 1 As shown, in step S3, the evaluation results of the preset rounds are compared, and the improved YOLOv8 neural network structure corresponding to the round with the optimal evaluation result is determined as the image detection model.
[0161] Specifically, Figure 1 As shown, in step S4, the image detection model is tested using the test set, and it is determined whether the improved YOLOv8 neural network structure needs to be retrained based on the test results. If qualified, the image detection model is saved and can be applied; if not qualified, it needs to be retrained through step S3.
[0162] Compared with the prior art, a training method for a target detection model provided by an embodiment of the present invention forms a training set, a verification set and a test set by annotating and screening multiple images with small targets, uses the training set to train an improved YOLOv8 neural network structure, uses the verification set to evaluate the improved YOLOv8 neural network structure of each round, and uses the test set to test the obtained image detection model. In the improved YOLOv8 neural network structure, a sparse matrix module fuses the shallowest feature map and the deepest feature map output by the backbone network to improve the extraction of small target features. At the same time, only the feature map area where the target may exist is calculated during the fusion, so as to minimize the loss of small target features. It is possible to improve the efficiency of model prediction on the basis of ensuring the model prediction effect, so that the model can quickly detect and identify images; at the same time, CIoU and PIoU are used to calculate the bounding box loss value of each target, so that the improved YOLOv8 neural network structure can learn more small target features and improve the detection ability of small target images; and, by training the improved YOLOv8 neural network structure to obtain the learnable parameters corresponding to the tenth feature map and the twelfth feature map, the learnable parameters can be adaptively adjusted by fully considering the border size of different images, which is not only suitable for the detection of large target images, but also more conducive to improving the sensitivity of small target image detection.
[0163] Those skilled in the art will appreciate that all or part of the processes of the above-mentioned embodiments can be implemented by instructing related hardware through a computer program, and the program can be stored in a computer-readable storage medium, wherein the computer-readable storage medium is a disk, an optical disk, a read-only storage memory, or a random access memory, etc.
[0164] The above description is only a preferred specific implementation manner of the present invention, but the protection scope of the present invention is not limited thereto. Any changes or substitutions that can be easily conceived by any technician familiar with the technical field within the technical scope disclosed by the present invention should be covered within the protection scope of the present invention.
Claims
1. A method for training a target detection model, characterized in that: The training methods shown include: Step S1: annotate a large number of historical images and select multiple images with small objects, the selected multiple annotated images form a sample set, and the sample set is divided into a training set, a verification set and a test set; Step S2: construct an improved YOLOv8 neural network structure; Step S3: using the training set to train the improved YOLOv8 neural network structure for multiple rounds, and at the end of each round, using the validation set to evaluate the improved YOLOv8 neural network structure after the training of the round, to obtain the evaluation result of each round; after the training is completed, comparing the evaluation results of each round, and obtaining the improved YOLOv8 neural network structure of the round corresponding to the optimal evaluation result as the image detection model; Step S4: Use the test set to test the image detection model. If the test result is qualified, save the image detection model and end the training; if it is unqualified, return to step S3 to continue training.
2. The training method according to claim 1, characterized in that: The method of labeling a large number of historical images and selecting a plurality of images with small objects, wherein the selected plurality of labeled images form a sample set, includes: Annotate the bounding box of the target in each historical image, determine the center point coordinates, height, width, area and target category of each bounding box in each historical image, and the total width, total height and total area of each historical image; According to the height, width and area of each bounding box, combined with the total width, total height and total area of the corresponding historical image, the height ratio, width ratio and area ratio of each target in the corresponding historical image are calculated; if there is at least one target in the historical image whose height ratio is less than the height ratio threshold, or the width ratio is less than the width ratio threshold, or the area ratio is less than the area ratio threshold, then the historical image is taken as a sample in the sample set.
3. The training method according to claim 1, characterized in that: The improved YOLOv8 neural network structure is trained for multiple rounds using the training set, including: In each round, multiple samples are extracted from the training set as a batch of samples, and each batch of samples is input into the improved YOLOv8 neural network structure; According to the samples of each batch and the corresponding prediction results, the loss function value of the batch is calculated; according to the loss function value of the batch, the improved YOLOv8 neural network structure is updated through back propagation; The updated improved YOLOv8 neural network structure is trained using the next batch of samples until the training of the preset batch is completed to obtain the improved YOLOv8 neural network structure of the current round; The improved YOLOv8 neural network structure of the current round is trained for the next round until the preset round of training is completed.
4. The training method according to claim 3, characterized in that: The loss function value of each batch is calculated based on the samples of each batch and the corresponding prediction results, including: Calculate the category loss value and bounding box loss value of all targets of all samples in the batch of samples; wherein each sample includes at least one target; The average category loss value is calculated based on the category loss value of each target, and the average bounding box loss value is calculated based on the bounding box loss value of each target, thereby obtaining the loss function value of the batch.
5. The training method according to claim 4, characterized in that: The bounding box loss value of each target is calculated by the following formula: in, represents the bounding box loss function, Sigmoid() represents the activation function, S t Represents the pixel area of the real border, S st represents the preset border pixel area threshold for distinguishing whether it is a small target, b represents the preset ratio parameter, (1-PIoU) represents the first border loss function, (1-CIoU) represents the second border loss function, and x t ,y t represents the true center coordinate, x p ,y p represents the predicted center coordinates, SD represents the average area of the bounding boxes of all objects in all training samples in each batch, A represents the true bounding box, B represents the predicted bounding box, l1 represents the distance between the center coordinates of the true bounding box and the predicted bounding box, l2 represents the length of the diagonal of the minimum circumscribed rectangular box of the true bounding box and the predicted bounding box, and w A 、h A Indicates the width and height of the actual border, w B 、h B represents the width and height of the predicted bounding box, and v represents the aspect ratio consistency factor function.
6. The training method according to claim 5, characterized in that: The category loss value of each target is calculated by the following formula: in, represents the category loss function, N represents the total number of categories, i represents the i-th category, and Y i represents the true category value of the i-th category, P i Represents the predicted category value of the i-th category.
7. The training method according to claim 6, characterized in that: The loss function value of each batch is calculated by the following formula: u1+u2=1; in, represents the loss function of each batch, u1 represents the category weight coefficient, u2 represents the bounding box weight coefficient, represents the average loss value of the category, Represents the average loss value of the bounding box.
8. The training method according to claim 1, characterized in that: The neck network in the improved YOLOv8 neural network structure includes a sparse matrix module and an original neck network; The original neck network is used to fuse other feature maps output by the backbone network except for the shallowest feature map; The sparse matrix module is used to fuse the shallowest feature map and the deepest feature map output by the backbone network. The fused feature map and the fused feature map obtained by the original neck network are further fused in the neck network and then output.
9. The training method according to claim 8, characterized in that: The improved YOLOv8 neural network structure also includes: The backbone network is used to extract features from the input samples, and input the obtained second feature map, third feature map, fourth feature map and fifth feature map from shallow to deep into the neck network; The neck network also includes a dynamic weighted fusion module; the sparse matrix module fuses the second feature map and the fifth feature map to output a tenth feature map, and the original neck network fuses the third to fifth feature maps to obtain an eleventh feature map and a twelfth feature map; The dynamic weighted fusion module is used to fuse the tenth feature map and the twelfth feature map to output a thirteenth feature map; The original neck network is also used to further fuse the thirteenth feature map, the eleventh feature map and the fifth feature map to obtain the sixth feature map, the seventh feature map and the eighth feature map, and input the sixth feature map, the seventh feature map and the eighth feature map into the detection head network; The detection head network is used to obtain a predicted target result of the sample according to the sixth feature map, the seventh feature map and the eighth feature map; wherein the predicted target result includes a predicted target category and a predicted target bounding box.
10. The training method according to claim 9, characterized in that: The sparse matrix module includes a double upsampling layer, a nonlinear activation layer, a first upsampling layer, a multiplication module and a first convolution module; A double upsampling layer, used to upsample the fifth feature map of the input twice, and input the twice upsampled feature map into the nonlinear activation layer; A non-linear activation layer is used to perform non-linear activation on the input feature map to obtain a sparse matrix and input it into the first upsampling layer; A first upsampling layer, used for upsampling the sparse matrix to obtain a ninth feature map and inputting it into a multiplication module; A multiplication module, used for multiplying the input second feature map and the ninth feature map, and inputting the multiplied feature map into the first convolution module; The first convolution module is used to downsample and nonlinearly represent the input feature map to obtain a tenth feature map.