Target detection method and device based on improved YOLOv5 model
By improving the YOLOv5 model and combining the weight sampling method, the traditional object detection method solves the problems of false detection and missed detection when dealing with category sample imbalance, extreme small targets and extreme aspect ratio targets, achieving higher detection accuracy and robustness.
Patent Information
- Application Number
- CN202411716879.X
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-11-27
- Publication Date
- 2025-03-18
AI Technical Summary
When traditional object detection methods deal with data set category sample imbalance, extreme small targets and extreme aspect ratio targets, the error detection rate and missed detection rate are high, making it difficult to adapt to complex image environments.
The improved YOLOv5 model is adopted to pre-train the convolutional layer and specific modules of the YOLOv5 model, and initialize the model with the weight sampling method, solving the problem of sample imbalance in the data set category, and improving the model's detection ability of small targets and extreme aspect ratio targets.
It effectively reduces the false detection rate and missed detection rate, improves the accuracy and robustness of target detection, and is suitable for target detection tasks in complex scenarios.
Smart Images
Figure CN119648997A_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the field of computer vision technology, and in particular to a target detection method and device based on an improved YOLOv5 model. Background Art
[0002] Object detection is a key task in computer vision. Its goal is to accurately detect and locate objects of interest in images or videos. This technology is widely used in various fields, such as autonomous driving, video surveillance, face recognition, industrial quality inspection, and medical image analysis.
[0003] Traditional object detection methods primarily rely on hand-crafted feature extractors and classifiers, such as Haar features, histogram of oriented gradient features, and scale-invariant feature transforms. These methods manually select and design image features to identify and locate objects. However, the limitations of such methods are also very obvious: first, feature design requires a wealth of domain knowledge and experience and often struggles to adapt to complex image environments; second, when processing complex scenes, hand-crafted features perform poorly in some special cases, such as small object detection and recognition of objects with extreme aspect ratios. Furthermore, the imbalance of class samples in datasets (i.e., the difference in the number of samples across different categories) also contributes to the inadequacy of traditional methods in detection accuracy. This is especially true when faced with images with high background noise, which often results in high false detection and missed detection rates.
[0004] Therefore, effective methods are urgently needed to solve the problems of high false detection and missed detection rates caused by imbalanced dataset category samples, extremely small objects, and objects with extreme aspect ratios. Summary of the Invention
[0005] In view of this, an embodiment of the present invention provides a target detection method and device based on an improved YOLOv5 model to solve the problems of high false detection rate and missed detection rate.
[0006] To achieve the above objectives, the embodiments of the present invention provide the following technical solutions:
[0007] A first aspect of the present invention discloses a target detection method based on an improved YOLOv5 model, the method comprising:
[0008] Acquire a data set to be detected, wherein the data set to be detected includes image data, and the image data includes a target;
[0009] Inputting the data set to be detected into a pre-trained target detection model, and outputting target location and category information corresponding to the data set to be detected; wherein the target detection model is pre-trained using a sample image to train an improved YOLOv5 model; the improved YOLOv5 model is obtained by pre-training the convolutional layer and specific modules of the YOLOv5 model;
[0010] The target position and category information corresponding to the data set to be detected are formatted to obtain an output result.
[0011] Preferably, the process of pre-training the improved YOLOv5 model using sample images to obtain the target detection model includes:
[0012] Collecting an original data set and annotating images in the original data set to obtain sample images;
[0013] Dividing the sample images into a training set and a validation set;
[0014] Initialize the improved YOLOv5 model according to the weight sampling method to obtain the initial target detection model;
[0015] Inputting the training set into the initial target detection model to obtain target position and category information corresponding to the training set;
[0016] Calculating the loss function of the initial target detection model based on the target position and category information corresponding to the training set;
[0017] Determine whether the loss function meets a preset value;
[0018] If the loss function satisfies a preset value, the initial target detection model is optimized using the validation set to obtain a final target detection model;
[0019] If the loss function does not meet the preset value, the parameters of the initial target detection model are updated according to the gradient of the loss function, and the step of calculating the loss function of the initial target detection model according to the target position and category information corresponding to the training set is returned to execute.
[0020] Preferably, the step of labeling the images in the original data set to obtain sample images includes:
[0021] For each image in the original data set, draw a bounding box of the image using a preset annotation tool;
[0022] Adding a category label to the image after drawing the bounding box to obtain a labeled image;
[0023] Mark all annotated images as sample images.
[0024] Preferably, the improved YOLOv5 model is initialized according to the weight sampling method to obtain an initial target detection model, including:
[0025] Read the true target box in the initialization sample through the improved YOLOv5 model, and write the true target box into a list;
[0026] Counting the number of bounding boxes corresponding to each category label in the list through the improved YOLOv5 model, and determining the weight of each category label based on the inverse of the number of bounding boxes;
[0027] The improved YOLOv5 model is initialized based on the weighted sampling of the category labels to obtain an initial object detection model.
[0028] Preferably, if the loss function satisfies a preset value, the initial target detection model is optimized using the validation set to obtain a final target detection model, including:
[0029] If the loss function satisfies a preset value, the initial target detection model is verified using the verification set;
[0030] If the verification result of the initial target detection model does not meet the verification conditions, adjust the hyperparameters of the initial target detection model, and return to the step of verifying the initial target detection model using the verification set if the loss function meets the preset value;
[0031] If the verification result of the initial target detection model meets the verification conditions, it is determined that the final target detection model is obtained.
[0032] Preferably, the process of pre-training the convolutional layer and specific modules of the YOLOv5 model to obtain the improved YOLOv5 model includes:
[0033] Get the YOLOv5 model;
[0034] Modify the convolution layer of the YOLOv5 model, modify the standard convolution to the inner convolution, and obtain a first YOLOv5 model;
[0035] Modify the Concat structure of the first YOLOv5 model to a Concat_bifpn structure to obtain a second YOLOv5 model;
[0036] The second YOLOv5 model is pre-trained according to the pre-training samples to obtain an improved YOLOv5 model.
[0037] A second aspect of the present invention discloses a target detection device based on an improved YOLOv5 model, the device comprising:
[0038] An acquisition unit, configured to acquire a data set to be detected, wherein the data set to be detected includes image data, and the image data includes a target;
[0039] A first input unit is configured to input the dataset to be detected into a pre-trained target detection model and output target location and category information corresponding to the dataset to be detected; wherein the target detection model is pre-trained using a sample image to train an improved YOLOv5 model; and the improved YOLOv5 model is pre-trained on a convolutional layer and a specific module of the YOLOv5 model;
[0040] The conversion unit is used to convert the target position and category information corresponding to the data set to be detected into a format to obtain an output result.
[0041] Preferably, the device further comprises:
[0042] A collecting unit, configured to collect an original data set and annotate images in the original data set to obtain sample images;
[0043] A division unit, configured to divide the sample image into a training set and a validation set;
[0044] An initialization unit is used to initialize the improved YOLOv5 model according to the weight sampling method to obtain an initial target detection model;
[0045] A second input unit is used to input the training set into the initial target detection model to obtain target position and category information corresponding to the training set;
[0046] A calculation unit, configured to calculate a loss function of the initial target detection model based on target positions and category information corresponding to the training set;
[0047] A judgment unit, configured to judge whether the loss function satisfies a preset value;
[0048] an optimization unit, configured to optimize the initial target detection model using the validation set to obtain a final target detection model if the loss function satisfies a preset value;
[0049] An updating unit is used to update the parameters of the initial target detection model according to the gradient of the loss function if the loss function does not meet the preset value, and return to the execution calculation unit.
[0050] Preferably, the collecting unit is specifically used for:
[0051] For each image in the original data set, a bounding box of the image is drawn using a preset annotation tool; a category label is added to the image after the bounding box is drawn to obtain an annotated image; and all the annotated images are marked as sample images.
[0052] Preferably, the initialization unit includes:
[0053] A reading module is used to read the true target bounding box in the initialization sample through the improved YOLOv5 model and write the true target bounding box into a list;
[0054] A statistics module, configured to count the number of bounding boxes corresponding to each category label in the list using the improved YOLOv5 model, and determine the weight of each category label based on the inverse of the number of bounding boxes;
[0055] An initialization module is used to initialize the improved YOLOv5 model based on the weight sampling of the category labels to obtain an initial target detection model.
[0056] Based on the above-mentioned embodiment of the present invention, a target detection method and device based on an improved YOLOv5 model are provided, which obtains a data set to be detected; inputs the data set to be detected into a pre-trained target detection model, and outputs the target position and category information corresponding to the data set to be detected; wherein the target detection model is obtained by pre-training the improved YOLOv5 model using sample images; the improved YOLOv5 model is obtained by pre-training the convolutional layer and specific modules of the YOLOv5 model; the target position and category information corresponding to the data set to be detected are formatted to obtain an output result. The target detection model is used to perform target detection on images or videos in the data set to be detected, and the output results of the model are applied to actual scenarios. The trained target detection model effectively solves the problems caused by unbalanced class samples in the data set, extremely small targets, and targets with extreme aspect ratios, thereby improving the detection accuracy. BRIEF DESCRIPTION OF THE DRAWINGS
[0057] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the following briefly introduces the drawings required for use in the embodiments or the description of the prior art. Obviously, the drawings described below are merely embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on the provided drawings without paying any creative work.
[0058] Figure 1 A flowchart of a target detection method based on an improved YOLOv5 model provided by an embodiment of the present invention;
[0059] Figure 2A network structure diagram of YOLOv5 provided in an embodiment of the present invention;
[0060] Figure 3 A flowchart of training an improved YOLOv5 model to obtain a target detection model provided by an embodiment of the present invention;
[0061] Figure 4 Flowchart of the improved YOLOv5 model obtained through pre-training provided by an embodiment of the present invention;
[0062] Figure 5 A schematic diagram of an inward convolution provided by an embodiment of the present invention;
[0063] Figure 6 This is a structural block diagram of a target detection device based on an improved YOLOv5 model provided by an embodiment of the present invention. DETAILED DESCRIPTION
[0064] The following will clearly and completely describe the technical solutions in the embodiments of the present invention in conjunction with the accompanying drawings. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative efforts are within the scope of protection of the present invention.
[0065] In this application, the terms "comprises," "comprising," or any other variations thereof are intended to encompass non-exclusive inclusion, such that a process, method, article, or apparatus that includes a list of elements includes not only those elements but also other elements not explicitly listed, or elements inherent to such process, method, article, or apparatus. In the absence of further limitations, an element defined by the phrase "comprising a ..." does not preclude the presence of additional identical elements in the process, method, article, or apparatus that includes the element.
[0066] As we can see from the background, traditional object detection methods rely on manually selecting and designing image features to identify and locate targets. However, feature design requires extensive domain knowledge and experience and often struggles to adapt to complex image environments. Furthermore, when dealing with complex scenes, manually designed features often perform poorly in some specific situations. Furthermore, the imbalance of class samples in datasets (i.e., the disparity in the number of samples across different categories) also contributes to the inadequacy of traditional methods in detection accuracy. This is particularly true when dealing with images with high background noise, often leading to high false positive and missed detection rates.
[0067] Therefore, an embodiment of the present invention provides a target detection method and device based on an improved YOLOv5 model, which obtains a data set to be detected; inputs the data set to be detected into a pre-trained target detection model, and outputs the target position and category information corresponding to the data set to be detected; wherein the target detection model is obtained by pre-training the improved YOLOv5 model using sample images; the improved YOLOv5 model is obtained by pre-training the convolutional layer and specific modules of the YOLOv5 model; the target position and category information corresponding to the data set to be detected are format-converted to obtain an output result. The target detection model is used to perform target detection on images or videos in the data set to be detected, and the output results of the model are applied to actual scenarios. The trained target detection model effectively solves the problems caused by unbalanced class samples in the data set, extremely small targets, and targets with extreme aspect ratios, thereby improving the detection accuracy.
[0068] See also Figure 1 , shows a flowchart of a target detection method based on an improved YOLOv5 model provided by an embodiment of the present invention.
[0069] It should be noted that the YOLOv5 network is a target detection network proposed by Ultralytics, a machine learning training and deployment platform. Figure 2 The network structure diagram of YOLOv5 is shown. YOLOv5 adds Mosaic data enhancement, anchor box adaptive calculation, and image adaptive scaling processing before input; adds Focus structure, CSP structure, and SPP structure to the backbone part; uses PAN structure and FPN structure in the neck part, and uses GIOU_Loss in the head part.
[0070] The method includes:
[0071] Step S101: Obtain a data set to be detected.
[0072] It should be noted that the data set to be detected includes image data, and the image data contains targets.
[0073] In some embodiments, the data set to be detected also includes video data, frames are extracted from the video data, and each frame includes a target.
[0074] In the specific implementation of step S101, the data set to be detected is obtained and preprocessed, including but not limited to image normalization, size adjustment, color space conversion, etc., to ensure that the preprocessed data set to be detected matches the input requirements of the target detection model.
[0075] Step S102: Input the data set to be detected into a pre-trained target detection model, and output the target position and category information corresponding to the data set to be detected.
[0076] It is understood that the target detection model is obtained by pre-training the improved YOLOv5 model using sample images. The specific training process is detailed in the following embodiments of the present invention. Figure 3 The content in.
[0077] It should be noted that the target detection model can also be obtained by pre-training models such as Faster R-CNN, YOLO, and SSD using sample images.
[0078] It is understood that the improved YOLOv5 model is obtained by pre-training the convolutional layer and specific modules of the YOLOv5 model. The specific training process is detailed in the following embodiments of the present invention. Figure 4 The content in.
[0079] In the specific implementation of step S102, the data set to be detected is input into a pre-trained target detection model, and a series of sample balancing, convolution, pooling and activation function operations are performed on the data set to be detected through various layers in the target detection model. Then, feature extraction and detection are performed in the data set to be detected to identify the target position and category information, thereby outputting the target position and category information corresponding to the data set to be detected.
[0080] Step S103: convert the target position and category information corresponding to the dataset to be detected into a format to obtain an output result.
[0081] It's understandable that after obtaining the object location and category information corresponding to the target dataset output by the object detection model, some post-processing operations are usually required to improve the accuracy and stability of the results. Post-processing operations include, but are not limited to, non-maximum suppression (NMS) and bounding box regression. The processed results are then parsed into a format that is easy to understand and use, such as bounding box coordinates and category labels.
[0082] In some embodiments, the output results are displayed on an interactive interface, or used for subsequent processing and decision-making in actual applications.
[0083] In practical applications, the performance of the target detection model is monitored and optimized to ensure its stability and accuracy in different scenarios. The performance indicators of the target detection model are monitored and the target detection model is updated and optimized as needed to adapt to new data and requirements.
[0084] In an embodiment of the present invention, a trained object detection model is deployed in a practical application and used to detect objects in images or videos in a dataset. Through preprocessing, forward propagation, feature extraction and detection, postprocessing, and result analysis, the model's output is applied to real-world scenarios, and the model's performance is continuously monitored and optimized. This trained object detection model effectively addresses issues caused by dataset class imbalance, extremely small objects, and objects with extreme aspect ratios, thereby improving detection accuracy.
[0085] The above embodiments of the present invention Figure 1 For the specific implementation of the target detection model obtained by pre-training the improved YOLOv5 model with sample images, see Figure 3 , which shows a flow chart of training an improved YOLOv5 model to obtain a target detection model provided by an embodiment of the present invention, including:
[0086] Step S301: collect original data sets and annotate images in the original data sets to obtain sample images.
[0087] In the specific implementation of step S301, image data or video data containing the target is collected, and for each image in the image data and / or video data, the bounding box and category of the target in the image are annotated to obtain a sample image. The specific process is as follows (process A1 to process A3):
[0088] Process A1: For each image in the original dataset, use the preset annotation tool to draw the bounding box of the image.
[0089] It should be noted that the preset labeling tools include but are not limited to LabelImg, Labelbox, and CVAT.
[0090] When specifically implementing process A1, for each image in the original data set, a bounding box is drawn using a rectangular box tool provided by a preset annotation tool.
[0091] It should be noted that this bounding box should surround the target object as accurately as possible and frame the outer contour of the object.
[0092] Process A2: Add category labels to the image after drawing the bounding box to obtain the labeled image.
[0093] In the specific implementation of process A2, a category label is added to the image after the bounding box is drawn. The category label is usually set according to the task requirements, such as "person", "car", "animal", etc. The label can be a numerical code (such as 0 for person and 1 for car) or text (such as "person", "car"). The labeled image is obtained.
[0094] It should be noted that for all annotated images, each bounding box corresponds to a category label.
[0095] Process A3: Label all annotated images as sample images.
[0096] That is, all annotated images are used as sample images to train the improved YOLOv5 model.
[0097] Step S302: Divide the sample images into a training set and a validation set.
[0098] In the specific implementation of step S302, the sample images are divided into a training set and a validation set, the improved YOLOv5 model is trained with the training set, and the training effect is verified with the validation set.
[0099] In some embodiments, the sample images are further divided into a training set, a validation set, and a test set.
[0100] The performance of the final object detection model is evaluated through the test set. Common evaluation indicators include precision, recall rate, F1 score, etc.
[0101] Step S303: Initialize the improved YOLOv5 model according to the weight sampling method to obtain an initial target detection model.
[0102] It should be noted that the weighted sampling method is also the problem of adaptive category sample imbalance, that is, high weights are given to targets that appear less frequently, and low weights are given to samples that appear more frequently, so as to achieve oversampling of a small number of samples.
[0103] In the specific implementation of step S303, the improved YOLOv5 model is initialized according to the weight sampling method to obtain the initial object detection model. This solves the problem of imbalanced class samples in the dataset and achieves a balanced class sample effect. This helps improve the model's ability to detect rare objects. The specific initialization process is as follows (process B1 to process B3):
[0104] Process B1: Read the true target box in the initialization sample through the improved YOLOv5 model and write the true target box into a list.
[0105] In the specific implementation process B1, the ground truth of the initialization sample is read through the improved YOLOv5 model, and the ground truth of the initialization sample is written into a list.
[0106] Process B2: The improved YOLOv5 model is used to count the number of bounding boxes corresponding to each category label in the list, and the weight of each category label is determined based on the inverse of the number of bounding boxes.
[0107] In the specific implementation process B2, the number of bounding boxes corresponding to each category label in the statistical list is counted through the improved YOLOv5 model, and then the weight of each category label is determined based on the inverse of the number of bounding boxes. That is, the category label with more bounding boxes has a smaller weight.
[0108] During each training epoch, we balance the samples by selecting the image samples for each class label according to the class label weights. We use the random.choices function in Python to sample according to the weights and change the training image index to achieve a balanced sample.
[0109] Process B3: Initialize the improved YOLOv5 model based on the weight sampling of the category label to obtain the initial object detection model.
[0110] In the specific implementation process B3, the improved YOLOv5 model is initialized based on the weight sampling of the category label and trained using a large-scale image dataset to enhance the randomness and robustness of the model, improve the generalization ability of the model, achieve better training results and performance, and finally obtain the initial object detection model.
[0111] Step S304: Input the training set into the initial target detection model to obtain the target position and category information corresponding to the training set.
[0112] In the specific implementation of step S304, the training set in the sample image is input into the initial target detection model, and the initial target detection model is used to perform detection to obtain target position and category information corresponding to the training set.
[0113] Step S305: Calculate the loss function of the initial target detection model based on the target position and category information corresponding to the training set.
[0114] It is understandable that the loss function used to optimize the model parameters is predefined. Commonly used loss functions include position loss (such as Smooth L1 Loss) and classification loss (such as cross entropy loss) in target detection.
[0115] In the specific implementation of step S305 , the loss function of the initial target detection model is calculated based on the target position and category information corresponding to the training set.
[0116] Step S306: Determine whether the loss function meets a preset value.
[0117] In the specific implementation of step S306, it is determined whether the loss function meets the preset value; if the loss function meets the preset value, step S307 is executed; if the loss function does not meet the preset value, step S308 is executed.
[0118] Step S307: If the loss function meets the preset value, the initial target detection model is optimized using the validation set to obtain the final target detection model.
[0119] It is understandable that using the validation set in the sample images to evaluate and tune the initial object detection model can help avoid overfitting and improve the generalization ability of the model.
[0120] In the specific implementation of step S307, when the loss function meets the preset value, the initial target detection model is optimized according to the validation set in the sample image to obtain the final target detection model. The specific optimization process is as follows (process C1 to process C3):
[0121] Process C1: If the loss function meets the preset value, the initial target detection model is verified using the validation set.
[0122] In the specific implementation of process C1, when the loss function meets the preset value, the initial object detection model is verified using the validation set to obtain the validation result. It is then determined whether the validation result meets the validation conditions. If the validation result does not meet the validation conditions, process C2 is executed; if the validation result does meet the validation conditions, process C3 is executed.
[0123] Process C2: If the verification result of the initial target detection model does not meet the verification conditions, the hyperparameters of the initial target detection model are adjusted and the process returns to process C1.
[0124] When implementing process C2, if the verification result of the initial target detection model does not meet the verification conditions, the hyperparameters of the initial target detection model are adjusted, specifically involving techniques such as learning rate adjustment and regularization, and then the process returns to execution process C1.
[0125] Process C3: If the verification result of the initial target detection model meets the verification conditions, the final target detection model is determined.
[0126] In the specific implementation process C3, when the verification result of the initial target detection model meets the verification condition, it is determined that the final target detection model is obtained.
[0127] Step S308: If the loss function does not meet the preset value, the parameters of the initial target detection model are updated according to the gradient of the loss function, and the process returns to step S305.
[0128] In the specific implementation of step S308, when the loss function does not meet the preset value, the back propagation algorithm and optimizer (such as stochastic gradient descent) are used to update the parameters of the initial target detection model, and the process returns to step S305.
[0129] In this embodiment of the present invention, the object detection model is trained by comprehensively considering multiple steps, including data preparation, model selection and definition, training and parameter adjustment, evaluation, and deployment. Machine learning, deep learning, and computer vision techniques are integrated to train the object detection model. A weighted sampling method is introduced to oversample samples that occur less frequently. By assigning higher weights to samples in minority categories, the model pays more attention to minority categories during training, achieving a balanced sample size across categories. This effectively improves the model's ability to detect rare objects.
[0130] The above embodiments of the present invention Figure 1 For the specific implementation of the improved YOLOv5 model obtained by pre-training the convolutional layers and specific modules of the YOLOv5 model, see Figure 4 , which shows a flow chart of the improved YOLOv5 model obtained through pre-training provided by an embodiment of the present invention, including:
[0131] Step S401: Obtain a YOLOv5 model.
[0132] In the specific implementation of step S401, a YOLOv5 model is obtained and a pre-trained YOLOv5 model is loaded.
[0133] Step S402: Modify the convolution layer of the YOLOv5 model, modify the standard convolution to the inner convolution, and obtain the first YOLOv5 model.
[0134] It's important to note that standard convolution kernels possess two fundamental properties: spatial invariance and channel specificity. Inversely, inward convolution exhibits both channel invariance and spatial specificity. Channel invariance reduces redundant information between channels, improving computational efficiency. Spatial specificity enables convolution kernels to capture diverse feature representations across spatial locations and mitigates long-range pixel dependencies.
[0135] As you can understand, assuming standard convolution ignores batch size (batch_size), the input feature map size is [c,w,h], and each group of feature maps corresponds to n convolution kernels of size [c,k,k]. Within the same group, each convolution kernel slides across the feature map, sharing the same weights in the convolution kernel at different spatial locations, and different weights for different channels. The number of convolution kernels does not depend on the input feature map.
[0136] It should be noted that, combined with Figure 5The following diagram shows an inward convolution. Assuming that inward convolution does not consider batch size (batch_size), the input feature map size is [c,w,h]. Each group of feature maps corresponds to h×w convolution kernels of size [1,k,k]. Within the same group, the convolution kernels at the same spatial location are expanded to [c,k,k]. Weights are shared between channels, and a weighted summation is performed to obtain a feature of size [c,1,1] in the output feature map. Different spatial locations correspond to different convolution kernels. The number of convolution kernels depends on the size of the input feature map, and the output feature map has the same number of channels as the input feature map.
[0137] Step S403: Modify the Concat structure of the first YOLOv5 model to a Concat_bifpn structure to obtain a second YOLOv5 model.
[0138] It should be noted that in the first YOLOv5 model, the Concat structure is replaced with the Concat_bifpn structure. The original Concat structure directly inputs feature maps of different scales for splicing, while Concat_bifpn encapsulates two weight parameters W1 = (w1, w2) and W2 = (w1, w2, w3), which are two-dimensional and three-dimensional respectively, indicating that two feature maps can be fused or three feature maps can be fused. Then, a 1*1 convolution kernel and a Relu activation function are declared. During forward propagation, the number of fused feature maps is determined. If there are two, the weight parameter W1 is used to calculate the output value, which is expressed as follows (Formula (1)):
[0139]
[0140] Step S404: Pre-train the second YOLOv5 model according to the pre-training sample to obtain an improved YOLOv5 model.
[0141] In the specific implementation of step S404, the second YOLOv5 model is pre-trained using the pre-training samples to obtain an improved YOLOv5 model.
[0142] In an embodiment of the present invention, a feature fusion adaptive mechanism is introduced. By assigning different weights to the feature maps of each scale during the feature fusion process, the model can flexibly adjust the contribution of features of different scales, which helps to enhance the sensitivity of the model to small targets, thereby improving the detection precision and accuracy. In addition, the present invention also introduces inward convolution to further optimize the performance of the model. Inward convolution can solve the limitations of traditional convolution in dealing with long-distance pixel dependence by introducing spatial specificity. For targets with extreme aspect ratios, inward convolution can effectively improve the detection effect and improve the model's recognition accuracy for such targets. In summary, by combining the adaptive weight adjustment of feature fusion and inward convolution, the model can more efficiently detect targets of different scales and shapes, and improve the detection capability of small targets and targets with extreme aspect ratios.
[0143] Corresponding to the target detection method based on the improved YOLOv5 model provided in the above embodiment of the present invention, see Figure 6 , shows a structural block diagram of a target detection device based on an improved YOLOv5 model provided by an embodiment of the present invention.
[0144] The device includes: an acquisition unit 601 , a first input unit 602 , and a conversion unit 603 .
[0145] The acquisition unit 601 is configured to acquire a data set to be detected, where the data set to be detected includes image data, and the image data includes a target.
[0146] The first input unit 602 is used to input the data set to be detected into a pre-trained target detection model, and output the target position and category information corresponding to the data set to be detected; wherein, the target detection model is pre-trained using the sample image to train the improved YOLOv5 model; the improved YOLOv5 model is obtained by pre-training the convolution layer and specific modules of the YOLOv5 model.
[0147] The conversion unit 603 is used to convert the target position and category information corresponding to the data set to be detected into a format to obtain an output result.
[0148] In an embodiment of the present invention, a trained object detection model is deployed in a practical application and used to detect objects in images or videos in a dataset. Through preprocessing, forward propagation, feature extraction and detection, postprocessing, and result analysis, the model's output is applied to real-world scenarios, and the model's performance is continuously monitored and optimized. This trained object detection model effectively addresses issues caused by dataset class imbalance, extremely small objects, and objects with extreme aspect ratios, thereby improving detection accuracy.
[0149] Combine Figure 6The content shown in the figure further includes: a collecting unit, a dividing unit, an initializing unit, a second input unit, a calculating unit, a judging unit, an optimizing unit and an updating unit.
[0150] The collection unit is used to collect the original data set and annotate the images in the original data set to obtain sample images.
[0151] The collection unit is specifically used to: for each image in the original data set, use a preset annotation tool to draw the bounding box of the image; add a category label to the image after the bounding box is drawn to obtain an annotated image; and mark all the annotated images as sample images.
[0152] The division unit is used to divide the sample images into a training set and a validation set.
[0153] The initialization unit is used to initialize the improved YOLOv5 model according to the weight sampling method to obtain the initial target detection model.
[0154] The second input unit is used to input the training set into the initial target detection model to obtain the target position and category information corresponding to the training set.
[0155] The calculation unit is used to calculate the loss function of the initial target detection model according to the target position and category information corresponding to the training set.
[0156] The judgment unit is used to judge whether the loss function meets the preset value.
[0157] The optimization unit is used to optimize the initial target detection model using the validation set if the loss function meets the preset value to obtain the final target detection model.
[0158] The update unit is used to update the parameters of the initial target detection model according to the gradient of the loss function if the loss function does not meet the preset value, and return to the execution calculation unit.
[0159] Combine Figure 6 The content shown, the initialization unit, includes: a reading module, a statistics module and an initialization module.
[0160] The reading module is used to read the true target box in the initialization sample through the improved YOLOv5 model and write the true target box into a list.
[0161] The statistics module is used to count the number of bounding boxes corresponding to each category label in the statistics list through the improved YOLOv5 model, and determine the weight of each category label based on the inverse of the number of bounding boxes.
[0162] The initialization module is used to initialize the improved YOLOv5 model based on the weight sampling of the category label to obtain the initial target detection model.
[0163] Combine Figure 6 The content shown, the optimization unit, includes: a verification module, a hyperparameter adjustment module and a determination module.
[0164] The verification module is used to verify the initial target detection model using the verification set if the loss function meets the preset value.
[0165] The hyperparameter adjustment module is used to adjust the hyperparameters of the initial target detection model if the verification result of the initial target detection model does not meet the verification conditions, and return to the execution verification module.
[0166] The determination module is used to determine the final target detection model if the verification result of the initial target detection model meets the verification conditions.
[0167] Combine Figure 6 The device further comprises:
[0168] The second acquisition unit is used to acquire the YOLOv5 model.
[0169] The first modification unit is used to modify the convolution layer of the YOLOv5 model, modifying the standard convolution to the inner convolution, to obtain the first YOLOv5 model.
[0170] The second modification unit is used to modify the Concat structure of the first YOLOv5 model into a Concat_bifpn structure to obtain a second YOLOv5 model.
[0171] The pre-training unit is used to pre-train the second YOLOv5 model according to the pre-training sample to obtain an improved YOLOv5 model.
[0172] In summary, the present invention proposes an object detection method and device based on an improved YOLOv5 model, which achieves improved performance and robustness when handling unbalanced datasets, extremely small objects, and objects with extreme aspect ratios. It can more accurately identify the location, size, and category of an object, making it suitable for object detection tasks in a variety of scenarios, including security, industry, agriculture, and military.
[0173] Each embodiment in this specification is described in a progressive manner. The same or similar parts between the embodiments can be referred to each other. Each embodiment focuses on the differences from other embodiments. In particular, for system or system embodiments, since they are basically similar to method embodiments, the description is relatively simple. For relevant parts, refer to the partial description of the method embodiment. The system and system embodiments described above are merely schematic, wherein the units described as separate components may or may not be physically separated, and the components displayed as units may or may not be physical units, that is, they may be located in one place, or they may be distributed on multiple network units. Some or all of the modules can be selected according to actual needs to achieve the purpose of the solution of this embodiment. A person of ordinary skill in the art can understand and implement it without expending creative work.
[0174] Professionals may further appreciate that the units and algorithm steps of each example described in conjunction with the embodiments disclosed herein can be implemented in electronic hardware, computer software, or a combination of the two. In order to clearly illustrate the interchangeability of hardware and software, the above description has generally described the components and steps of each example according to their functions. Whether these functions are performed in hardware or software depends on the specific application and design constraints of the technical solution. Professionals and technicians may use different methods to implement the described functions for each specific application, but such implementation should not be considered beyond the scope of the present invention.
[0175] The above description of the disclosed embodiments is intended to enable one skilled in the art to implement or use the present invention. Various modifications to these embodiments will be readily apparent to one skilled in the art, and the general principles defined herein may be implemented in other embodiments without departing from the scope of the present invention. Therefore, the present invention is not limited to the embodiments shown herein, but is intended to conform to the widest scope consistent with the principles and novel features disclosed herein.
Claims
1. A target detection method based on an improved YOLOv5 model, characterized in that: The method comprises: Acquire a data set to be detected, wherein the data set to be detected includes image data, and the image data includes a target; Input the data set to be detected into a pre-trained target detection model, and output the target position and category information corresponding to the data set to be detected; wherein the target detection model is obtained by pre-training an improved YOLOv5 model using sample images; the improved YOLOv5 model is obtained by pre-training the convolutional layer and specific modules of the YOLOv5 model; The target position and category information corresponding to the data set to be detected are formatted and outputted.
2. The method according to claim 1, characterized in that: The process of pre-training the improved YOLOv5 model using sample images to obtain the target detection model includes: Collecting original data sets, and annotating images in the original data sets to obtain sample images; Dividing the sample images into a training set and a validation set; Initialize the improved YOLOv5 model according to the weight sampling method to obtain the initial target detection model; Inputting the training set into the initial target detection model to obtain target position and category information corresponding to the training set; Calculating the loss function of the initial target detection model according to the target position and category information corresponding to the training set; Determine whether the loss function meets a preset value; If the loss function satisfies a preset value, the initial target detection model is optimized using the validation set to obtain a final target detection model; If the loss function does not meet the preset value, the parameters of the initial target detection model are updated according to the gradient of the loss function, and the step of calculating the loss function of the initial target detection model according to the target position and category information corresponding to the training set is returned to execute.
3. The method according to claim 2, characterized in that The step of labeling the images in the original data set to obtain sample images includes: For each image in the original data set, draw a bounding box of the image using a preset annotation tool; Adding a category label to the image after drawing the bounding box to obtain a labeled image; Mark all annotated images as sample images.
4. The method according to claim 2, characterized in that: Initializing the improved YOLOv5 model according to the weight sampling method to obtain an initial target detection model includes: Read the true target frame in the initialization sample through the improved YOLOv5 model, and write the true target frame into a list; Counting the number of bounding boxes corresponding to each category label in the list through the improved YOLOv5 model, and determining the weight of each category label based on the inverse of the number of bounding boxes; The improved YOLOv5 model is initialized based on the weighted sampling of the category labels to obtain an initial target detection model.
5. The method according to claim 2, characterized in that: If the loss function satisfies a preset value, the initial target detection model is optimized using the validation set to obtain a final target detection model, including: If the loss function satisfies a preset value, the initial target detection model is verified using the verification set; If the verification result of the initial target detection model does not meet the verification condition, the hyperparameters of the initial target detection model are adjusted, and the step of verifying the initial target detection model using the verification set if the loss function meets the preset value is returned to be executed; If the verification result of the initial target detection model meets the verification condition, it is determined that the final target detection model is obtained.
6. The method according to claim 1, characterized in that The process of pre-training the convolutional layers and specific modules of the YOLOv5 model to obtain an improved YOLOv5 model includes: Get the YOLOv5 model; Modify the convolution layer of the YOLOv5 model, modify the standard convolution to an inner convolution, and obtain a first YOLOv5 model; Modify the Concat structure of the first YOLOv5 model to a Concat_bifpn structure to obtain a second YOLOv5 model; The second YOLOv5 model is pre-trained according to the pre-training samples to obtain an improved YOLOv5 model.
7. A target detection device based on an improved YOLOv5 model, characterized in that: The device comprises: A first acquisition unit is used to acquire a data set to be detected, wherein the data set to be detected includes image data, and the image data includes a target; A first input unit is used to input the data set to be detected into a pre-trained target detection model, and output the target position and category information corresponding to the data set to be detected; wherein the target detection model is obtained by pre-training an improved YOLOv5 model using sample images; the improved YOLOv5 model is obtained by pre-training the convolutional layer and specific modules of the YOLOv5 model; The conversion unit is used to convert the target position and category information corresponding to the data set to be detected into a format to obtain an output result.
8. The device according to claim 7, characterized in that The device also includes: A collecting unit, used to collect an original data set and annotate images in the original data set to obtain sample images; A division unit, used for dividing the sample image into a training set and a verification set; An initialization unit, used to initialize the improved YOLOv5 model according to the weight sampling method to obtain an initial target detection model; A second input unit, used to input the training set into the initial target detection model to obtain target position and category information corresponding to the training set; A calculation unit, used to calculate the loss function of the initial target detection model according to the target position and category information corresponding to the training set; A judging unit, used to judge whether the loss function satisfies a preset value; An optimization unit, configured to optimize the initial target detection model using the validation set to obtain a final target detection model if the loss function satisfies a preset value; An updating unit is used to update the parameters of the initial target detection model according to the gradient of the loss function if the loss function does not meet the preset value, and return to the execution calculation unit.
9. The device according to claim 8, characterized in that The collecting unit is specifically used for: For each image in the original data set, a bounding box of the image is drawn using a preset annotation tool; a category label is added to the image after the bounding box is drawn to obtain an annotated image; and all the annotated images are marked as sample images.
10. The device according to claim 8, characterized in that The initialization unit comprises: A reading module, used for reading the real target frame in the initialization sample through the improved YOLOv5 model, and writing the real target frame into a list; A statistical module, used to count the number of bounding boxes corresponding to each category label in the list through the improved YOLOv5 model, and determine the weight of each category label based on the inverse of the number of bounding boxes; An initialization module is used to initialize the improved YOLOv5 model based on the weight sampling of the category label to obtain an initial target detection model.