Sample imbalance target detection method based on improved YOLOX model

Abstract

The present invention relates to a sample imbalance target detection method based on an improved YOLOX model. The improved YOLOX model is used to perform feature extraction on an image to be detected collected by a sweeping robot to obtain a plurality of first features. The plurality of first features are input into a feature fusion layer for cross-layer fusion to obtain a plurality of prediction features. Finally, the plurality of prediction features are processed by a prediction layer to obtain a target detection result of the image to be detected. The feature extraction layer includes a plurality of CSPDarkNets with an attention mechanism, and the cross-layer fusion processing includes fusing part or all of the intermediate features obtained after the fusion processing of the plurality of first features, so that shallow features and deep features can be further fused, thereby improving the effect of target detection for different types of targets. At the same time, the improved YOLOX model uses a Varifocal loss function and a balanced loss function, so that targets with low frequency of occurrence can be effectively identified when the samples are imbalanced.

Description

Technical Field

[0001] The present invention relates to the field of target detection, and in particular to a sample imbalance target detection method based on an improved YOLOX model. Background Art

[0002] With the development of target detection technology, the target detection algorithm based on deep learning has been proposed, and its accuracy far exceeds that of traditional target detection algorithms. However, existing algorithms are generally developed based on public data sets and cannot be directly applied to sweeping robots. How to accurately detect specific targets requires specific designs. There are a large number of small target objects in the home scene environment, such as paper balls, table and chair legs, pet feces, etc. At the same time, the difficulty of sampling different categories of data is different. For example, it is easy to collect a large number of samples for table and chair legs and paper balls, while it is expensive to collect pet feces and it is difficult to obtain a large number of samples. Existing target detection methods are generally targeted at public data sets with relatively balanced samples, which leads to inductive bias in specific sample imbalanced data sets, and ultimately affects the performance of the method. Summary of the Invention

[0003] Based on this, it is necessary to provide a sample imbalance target detection method based on the improved YOLOX model to address the problem of low accuracy in imbalanced sample target detection.

[0004] A sample imbalance target detection method based on an improved YOLOX model, wherein the improved YOLOX model includes a feature extraction layer, a feature fusion layer, and a prediction layer; the feature extraction layer includes multiple CSPDarkNets with an attention mechanism, and the loss functions used by the improved YOLOX model are a Varifocal loss function and a balanced loss function. The method includes:

[0005] Obtain the image to be detected collected by the sweeping robot;

[0006] Inputting the image to be detected into the feature extraction layer for processing, and obtaining a plurality of first features according to output results of the feature extraction layer;

[0007] Inputting the multiple first features into the feature fusion layer for cross-layer fusion processing to obtain multiple prediction features; wherein the cross-layer fusion processing includes fusing part or all of the intermediate features obtained after the fusion processing of the multiple first features;

[0008] The plurality of prediction features are processed by the prediction layer, and an object detection result of the image to be detected is obtained according to the processing result.

[0009] In one embodiment, the method further comprises:

[0010] Obtain a data set of images from the perspective of a sweeping robot, wherein the data set includes a training set and a validation set;

[0011] Inputting the training set into the improved YOLOX model to be trained for training, and verifying the improved YOLOX model using the verification set;

[0012] Obtain the Varifocal loss and balanced loss of the improved YOLOX model during training, and update the model parameters based on the reverse gradient propagation of the loss function until the model converges to obtain the improved YOLOX model after training.

[0013] In one embodiment, the method includes:

[0014] The balance loss Loss is obtained by the following formula:

[0015]

[0016] Where N is the number of samples, I i is the i-th picture, W C is the designed classification weight of the Cth category, L eiou is the effective intersection-to-union loss, L cls is the classification loss, L obj Varifocal loss.

[0017] In one embodiment, the method includes:

[0018] The Varifocal loss L is obtained by the following formula obj :

[0019]

[0020] Among them, p is the IoU-aware classification score predicted by the Varifocal loss function, q is the target score, q is the intersection over union (IoU) for positive samples and 0 for negative samples, g is the modulation factor, and a is the scaling factor.

[0021] In one embodiment, the method includes:

[0022] The classification weight W of the Cth class is obtained by the following formula C :

[0023]

[0024] Among them, n i is the number of labels of category i, n C is the number of labels in category C, n class is the number of categories in the dataset;

[0025] The effective intersection-over-union loss L is obtained by the following formula eiou :

[0026]

[0027] Among them, c represents the minimum closure diagonal distance that can contain both the predicted box and the real box, c w 、c h are the closure width and closure height of the two rectangles, r is the Euclidean distance between the two center points, b, w, and h are the center point, width, and height of the prediction box respectively, b gt 、w gt 、h gt They represent the center point, width and height of the real box respectively, and IoU is the intersection over union ratio.

[0028] In one embodiment, updating the model parameters based on the loss function reverse gradient propagation includes:

[0029] The weights of the improved YOLOX model during training are obtained through the following algorithm:

[0030]

[0031] Among them, W t and W t+1 is the weight of the tth iteration and the t+1th iteration, m W is the learning rate of the back-propagation algorithm, is the partial derivative of the loss function with respect to the weight W;

[0032] And, get the offset of the improved YOLOX model in training through the following algorithm:

[0033]

[0034] Among them, b t and b t+1 is the offset between the tth iteration and the t+1th iteration, m b is the learning rate of the back-propagation algorithm, is the partial derivative of the loss function with respect to the offset b;

[0035] In one embodiment, the method includes:

[0036] The model parameters are updated by an exponential moving average method, wherein the exponential moving average method is implemented by the following formula:

[0037]

[0038]

[0039] Among them, W EMA,k、b EMA,k They represent the parameter values ​​after exponential sliding average update of the weights and offsets after the kth round of training, a is the attenuation factor, are the weights and offsets after the k-1th round of training.

[0040] In one embodiment, the CSPDarkNet includes a Focus module and a plurality of Dark modules connected in series with the Focus module;

[0041] Each Dark module is equipped with an attention mechanism module. The input depth of the attention mechanism module should be the output depth of the preceding Dark module, and the input depth and output depth of the attention mechanism module are consistent with the input depth.

[0042] In one of the embodiments, the embedded position after the attention mechanism module is equipped with a detection head; the number of the detection heads is determined according to the size of the target to be detected in the image to be detected.

[0043] In one embodiment, the prediction layer includes a Head layer, which is used to calculate the loss in the model training stage; and to process the multiple prediction features in the model application stage to obtain the target detection result.

[0044] The above-mentioned sample imbalance target detection method based on the improved YOLOX model extracts features from the image to be detected collected by the sweeping robot through the improved YOLOX model to obtain multiple first features. The multiple first features are input into the feature fusion layer for cross-layer fusion to obtain multiple prediction features. Finally, the multiple prediction features are processed by the prediction layer to obtain the target detection result of the image to be detected. Among them, the feature extraction layer includes multiple CSPDarkNets with attention mechanisms. The cross-layer fusion processing includes fusing part or all of the intermediate features obtained after the fusion processing of the multiple first features, so that the shallow features and deep features can be further fused, thereby improving the effect of target detection for different types of targets. At the same time, the improved YOLOX model uses the Varifocal loss function and the balanced loss function, so that in the case of sample imbalance, targets with low frequency of occurrence can be effectively identified. BRIEF DESCRIPTION OF THE DRAWINGS

[0045] Figure 1 A structural diagram of an improved YOLOX model in one embodiment;

[0046] Figure 2 A flowchart for improving YOLOX model training in one embodiment;

[0047] Figure 3 Flowchart of a sample imbalance target detection method based on an improved YOLOX model in one embodiment;

[0048] Figure 4 1 is a schematic diagram of detection results of a sample imbalance target detection method based on an improved YOLOX model in one embodiment;

[0049] Figure 5 Schematic diagram of detection results of a sample imbalance target detection method based on an improved YOLOX model in one embodiment. DETAILED DESCRIPTION

[0050] In order to more clearly understand the above-mentioned objects, features and advantages of the present invention, the present invention is further described in detail below with reference to the accompanying drawings and specific embodiments. It should be noted that the embodiments of the present application and the features therein can be combined with each other unless there is a conflict.

[0051] In the following description, many specific details are set forth to fully understand the present invention. However, the present invention may also be implemented in other ways different from those described herein. Therefore, the scope of protection of the present invention is not limited to the specific embodiments disclosed below.

[0052] The following describes, with reference to the accompanying drawings, a small target detection method based on an improved YOLOX model according to some embodiments of the present invention. The method of the present invention can be applied to the cloud or server, as well as to terminals capable of image detection, such as mobile robots, wearable devices, and mobile phones. The following describes the application of this method to a robot vacuum cleaner as an example.

[0053] like Figure 1 As shown, this embodiment discloses a structure of an improved YOLOX model, which includes a feature extraction layer, a feature fusion layer and a prediction layer.

[0054] like Figure 2 As shown in Figure 2, the training process of the model includes:

[0055] During the training data preparation phase, a dataset of images from the sweeper's perspective is obtained.

[0056] Among them, the image dataset includes a training set and a validation set. The training set is obtained by collecting images from the perspective of the sweeper, and the parameters of the image can be determined based on the camera equipment of the sweeper. For example, the image width is 640 pixels and the height is 640 pixels. The commonly used image annotation tool LabelImg can also be used to target the images in the dataset to obtain correct label data. In some cases, the dataset can also be enhanced, such as translation, symmetry, rotation, color enhancement or other strong data enhancement methods.

[0057] For example, a sample dataset might include easily captured images of objects like table and chair legs and paper balls, as well as images of objects like rings and pet feces that are difficult to obtain and appear less frequently. This means that training is performed with an unbalanced sample set.

[0058] During the model training phase, the training set is input into the improved YOLOX model to be trained for training, and the improved YOLOX model is verified through the validation set to execute the model training process.

[0059] Among them, the dataset can be divided into 10:1, 9:1, and 5:1 according to the size of the dataset, and the divided dataset is input into the improved YOLOX model to be trained.

[0060] The improved YOLOX model to be trained can be a pre-built model, and its model structure is as follows: Figure 1 As shown, it includes feature extraction layer, feature fusion layer and prediction layer.

[0061] Among them, the feature extraction layer is CSPDarkNet with attention mechanism, which includes a Focus module and multiple Dark modules connected in series with the Focus module. Figure 1 In D2, D3, D4, and D5, each Dark module contains a basic convolution block and a CSPLayer. An attention mechanism module is added after CSPLayer. The attention mechanism module includes a channel attention module and a spatial attention module. The input depth of the attention mechanism module should be the output depth of the preceding Dark module, and the input depth and output depth of the attention mechanism module are consistent with the input depth. Adding an attention mechanism after each CSPDarkNet module is conducive to strengthening the extraction of shallow features. The attention mechanism module can be a CBAM (Convolutional Block Attention Module, Convolutional Block Attention Module) module, or other attention mechanism modules can be used, such as a Squeeze-and-Excitation (SE) module, an AccuracyBooster (AB) module, and an Accuracy Booster Plus (ABPlus) module.

[0062] The embedding position after the attention mechanism module is equipped with detection heads. The number of detection heads is adjusted according to the size of the target in the image to be detected. For example, when a suspected small target is detected in the image to be detected, the number of detection heads can be increased to improve the extraction of shallow and deep features.

[0063] Shallow features are extracted from layers closer to the input of the object detection model and generally provide a more abstract interpretation of the image. Deep features are extracted from layers closer to the output of the object detection model and generally provide a more specific interpretation of the image. For example, D2 is a shallow layer relative to D3, and so on.

[0064] The feature fusion layer is connected in series with each CSPDarkNet, including the P layer and the N layer, which is used to fuse the features output by the feature extraction layer. The P layer realizes the fusion of shallow features to deep features, while the N layer realizes the fusion of deep features to shallow features.

[0065] In some cases, such as Figure 1 As shown in the figure, in order to achieve cross-layer fusion, an additional fourth feature fusion layer can be introduced after the D2 layer, and convolution and upsampling can be used to perform feature fusion with other feature layers; a cross-layer connection position is added after the D2 layer to enable direct cross-layer feature fusion with the P4 layer features; a cross-layer connection is added after the D3 layer to enable direct cross-layer feature fusion with the P5 layer features; a cross-layer connection is added after the P4 layer to enable direct cross-layer feature fusion with the N2 layer features; a cross-layer connection is added after the P5 layer to enable direct cross-layer feature fusion with the N3 layer features; a convolution operation is performed after each feature fusion to achieve channel compression.

[0066] The prediction layer includes multiple Head layers, which can be used to calculate the loss in the model training phase and process the multiple prediction features output by the feature fusion layer in the model application phase to obtain the detection results.

[0067] In addition, before model training, the above model can be initialized, the network offset is initialized to 0, and the network weight parameters are initialized using the Kaiming Gaussian initialization method. The Kaiming Gaussian initialization method obeys the following distribution:

[0068]

[0069] Among them, W l is the weight of the first layer, N is the Gaussian distribution, a is the ReLU activation function, n l is the data dimension of layer l. The number of iterations can be set to max_epoch = 100 according to the size of the dataset.

[0070] During the model training process, the Varifocal loss and balanced loss of the improved YOLOX model during training are obtained, and the model parameters are updated based on the reverse gradient propagation of the loss function until the model converges to obtain the improved YOLOX model after training.

[0071] During training, you can configure the varifocal loss function and the balanced loss function to improve learning for imbalanced samples. Because sample imbalance leads to insufficient learning for few-shot classes, the balanced loss function can enhance the learning strength and improve convergence for few-shot classes.

[0072] Specifically, the balance loss Loss is obtained by the following formula:

[0073]

[0074] Where N is the number of samples, I i is the i-th picture, W C is the designed classification weight of the Cth category, L eiou is the effective intersection-to-union loss, L cls is the classification loss, L obj Varifocal loss.

[0075] The classification weight W of the Cth class is obtained by the following formula C :

[0076]

[0077] Among them, n i is the number of labels of category i, n C is the number of labels in category C, n class is the number of categories in the dataset; through W C It can amplify the loss of categories with a small number of samples and reduce the loss of categories with a large number of samples.

[0078] The effective intersection-over-union loss L is obtained by the following formula eiou :

[0079]

[0080] Among them, c represents the minimum closure diagonal distance that can contain both the predicted box and the real box, c w 、c h are the closure width and closure height of the two rectangles, r is the Euclidean distance between the two center points, b, w, and h are the center point, width, and height of the prediction box respectively, b gt 、w gt 、h gt They represent the center point, width and height of the real box respectively, and IoU is the intersection over union ratio.

[0081] The Varifocal loss L is obtained by the following formula obj :

[0082]

[0083] Among them, p is the IoU-aware classification score predicted by the Varifocal loss function, q is the target score, q is the intersection over union (IoU) for positive samples and 0 for negative samples, g is the modulation factor, and a is the scaling factor.

[0084] Specifically, the model parameters are updated based on the loss function through reverse gradient propagation until the model converges. First, the loss function is calculated and reverse gradient propagation is performed to obtain the weights and offsets that need to be updated after each iteration. This reverse gradient propagation is performed once per iteration. The model is then updated based on the parameters calculated using the exponential moving average. This process is repeated once per epoch until the model converges. One epoch means that all the data in the dataset is fed into the network, completing one forward calculation and backward propagation process.

[0085] Among them, after obtaining the above loss function, the weight W that needs to be updated after iteration can be calculated according to the following formula t+1 :

[0086]

[0087] Among them, W t and W t+1 is the weight of the tth iteration and the t+1th iteration, m W is the learning rate of the back-propagation algorithm, is the partial derivative of the loss function with respect to the weight W.

[0088] And, get the offset of the improved YOLOX model in training through the following algorithm:

[0089]

[0090] The offset is updated using the backpropagation algorithm, and its formula is as follows:

[0091]

[0092] Among them, b t and b t+1 is the offset between the tth iteration and the t+1th iteration, m b is the learning rate of the back-propagation algorithm, is the partial derivative of the loss function with respect to the offset b.

[0093] The exponential moving average (EMA) is performed on the model to update the model parameters. The formula is as follows:

[0094]

[0095]

[0096] Among them, WEMA,k 、b EMA,k They represent the parameter values ​​after exponential sliding average update of the weights and offsets after the kth round of training, a is the attenuation factor, are the weights and offsets after the k-1th round of training.

[0097] The improved YOLOX model uses the Varifocal loss function and the Balanced loss function to improve learning for samples with imbalanced categories, and the Effective Intersection-over-Union (EIOU) loss function to enhance network convergence, enabling the model to achieve better training results even with imbalanced samples. Furthermore, by configuring an attention mechanism within the CSPDarkNet module and overlaying a feature fusion layer with a cross-layer fusion mechanism, target detection is enhanced without losing auxiliary feature information. This results in a model suitable for target recognition by sweeping robots, improving recognition of small targets, especially those with low probability of appearance.

[0098] In one embodiment, Figure 3 As shown, a sample imbalance target detection method based on an improved YOLOX model is provided. In target detection, the improved YOLOX model trained by the above embodiment is used. The method includes:

[0099] Step S310: obtaining an image to be detected captured by an image acquisition device of the sweeping robot.

[0100] The images to be detected can include small objects such as paper balls, table and chair legs, coins, and pet feces. They can also include objects that appear less frequently, also known as low-frequency objects. These objects are difficult to obtain or appear infrequently, so they account for a relatively low proportion of the training sample set for the improved YOLOX model.

[0101] Before detecting the image to be detected, image preprocessing may be performed, and the preprocessing may include translation, symmetry, rotation, color enhancement, mosic data enhancement, and mixup data enhancement of the image.

[0102] Step S320: input the image to be detected into the feature extraction layer for processing, and obtain a plurality of first features according to the output result of the feature extraction layer.

[0103] The first features refer to the features extracted by each CSPDarkNet with attention mechanism in the feature extraction layer. The first features can include deep features and shallow features.

[0104] Specifically, the algorithm can input the image to be detected into the focus of the feature extraction layer, and then pass through layers D2 to D5 in sequence, and output the corresponding first feature in each layer.

[0105] In step S330 , the multiple first features are input into a feature fusion layer for cross-layer fusion processing to obtain multiple prediction features.

[0106] Among them, cross-layer fusion processing includes fusing part or all of the intermediate features obtained after the fusion of multiple first features. The feature fusion layer may include multiple feature fusion branches, or it may fuse the first features once and then perform a secondary fusion. In CSPDarkNet and each feature fusion branch, a detection head can be configured to determine the features to be fused at the position of the detection head. For example, if the target of a specific sweeper data is small, a detection head is added in the shallow layer. If the target is large, a detection head can be added in the deep layer.

[0107] Specifically, the algorithm can input the features to be fused into the corresponding feature fusion branch according to the position of the cross-layer fusion detection head to obtain the feature fusion result.

[0108] Step S340: Process the plurality of prediction features through the prediction layer, and obtain the target detection result of the image to be detected according to the processing result.

[0109] The target detection result refers to the detection of small targets with low frequency in samples from the image to be detected.

[0110] Specifically, the algorithm can input multiple prediction features into the Head layer and obtain the target detection result based on the output of the Head layer.

[0111] The method of the above embodiment performs feature extraction on the image to be detected collected by the sweeping robot through the improved YOLOX model to obtain multiple first features. The multiple first features are input into the feature fusion layer for cross-layer fusion to obtain multiple prediction features. Finally, the multiple prediction features are processed by the prediction layer to obtain the target detection result of the image to be detected. Among them, the feature extraction layer includes multiple CSPDarkNets with attention mechanisms, and the cross-layer fusion processing includes fusing part or all of the intermediate features obtained after the fusion processing of the multiple first features, so that the shallow features and deep features can be further fused, thereby improving the effect of target detection for different types of targets. At the same time, the improved YOLOX model uses the Varifocal loss function and the balanced loss function, so that when the samples are unbalanced, targets with low frequency of occurrence can be effectively identified. Figure 4 and Figure 5 The figure is a schematic diagram of the recognition effect when using the above model for detection. The technical features of the above embodiments can be combined in any way. To keep the description concise, not all possible combinations of the technical features in the above embodiments are described. However, as long as there is no contradiction in the combination of these technical features, they should be considered to be within the scope of this specification.

[0112] The above-described embodiments merely represent several implementation methods of the present invention. While the descriptions are relatively specific and detailed, they should not be construed as limiting the scope of the patent. It should be noted that a person of ordinary skill in the art may make various modifications and improvements without departing from the spirit of the present invention, all of which fall within the scope of protection of the present invention. Therefore, the scope of protection of the patent for this invention shall be determined by the appended claims.

Claims

1. The sample imbalance target detection method based on the improved YOLOX model is characterized by: The improved YOLOX model includes a feature extraction layer, a feature fusion layer, and a prediction layer; the feature extraction layer includes multiple CSPDarkNets with an attention mechanism, and the loss function used by the improved YOLOX model is a Varifocal loss function and a balanced loss function. The method includes: Obtain a data set of images from the perspective of a sweeping robot, wherein the data set includes a training set and a validation set; Inputting the training set into the improved YOLOX model to be trained for training, and verifying the improved YOLOX model using the verification set; Obtain the Varifocal loss and balanced loss of the improved YOLOX model during training, and update the model parameters based on the reverse gradient propagation of the loss function until the model converges to obtain the improved YOLOX model after training; Obtain the image to be detected collected by the sweeping robot; Inputting the image to be detected into the feature extraction layer for processing, and obtaining a plurality of first features according to output results of the feature extraction layer; Inputting the multiple first features into the feature fusion layer for cross-layer fusion processing to obtain multiple prediction features; wherein the cross-layer fusion processing includes fusing part or all of the intermediate features obtained after the fusion processing of the multiple first features; Processing the plurality of prediction features by the prediction layer, and obtaining an object detection result of the image to be detected according to the processing result; The balance loss Loss is obtained by the following formula: Where N is the number of samples, I i is the i-th picture, W C is the designed classification weight of the Cth category, L eiou is the effective intersection-to-union loss, L cls is the classification loss, L obj is Varifocal loss; Among them, the classification weight W of the Cth category is obtained by the following formula C : Among them, n i is the number of labels of category i, n C is the number of labels in category C, n class is the number of categories in the dataset; The effective intersection-over-union loss L is obtained by the following formula eiou : Among them, c represents the minimum closure diagonal distance that can contain both the predicted box and the real box, c w 、c h are the closure width and closure height of the two rectangles, r is the Euclidean distance between the two center points, b, w, and h are the center point, width, and height of the prediction box respectively, b gt 、w gt 、h gt They represent the center point, width and height of the real box respectively, and IoU is the intersection over union ratio.

2. The method according to claim 1, characterized in that include: The Varifocal loss L is obtained by the following formula obj : Among them, p is the IoU-aware classification score predicted by the Varifocal loss function, q is the target score, q is the intersection over union (IoU) for positive samples and 0 for negative samples, g is the modulation factor, and a is the scaling factor.

3. The method according to claim 1, characterized in that The updating of model parameters based on the reverse gradient propagation of the loss function includes: The weights of the improved YOLOX model during training are obtained through the following algorithm: Among them, W t and W t+1 is the weight of the tth iteration and the t+1th iteration, m W is the learning rate of the back-propagation algorithm, is the partial derivative of the loss function with respect to the weight W; And, get the offset of the improved YOLOX model in training through the following algorithm: Among them, b t and b t+1 is the offset between the tth iteration and the t+1th iteration, m b is the learning rate of the back-propagation algorithm, is the partial derivative of the loss function with respect to the offset b.

4. The method according to claim 1, wherein include: The model parameters are updated by an exponential moving average method, wherein the exponential moving average method is implemented by the following formula: Among them, W EMA，k 、b EMA，k They represent the parameter values ​​after exponential sliding average update of the weights and offsets after the kth round of training, a is the attenuation factor, are the weights and offsets after the k-1th round of training.

5. The method according to claim 1, wherein The CSPDarkNet includes a Focus module and a plurality of Dark modules connected in series with the Focus module; Each Dark module is equipped with an attention mechanism module. The input depth of the attention mechanism module should be the output depth of the preceding Dark module, and the input depth and output depth of the attention mechanism module are consistent with the input depth.

6. The method according to claim 5, characterized in that include: The embedding position after the attention mechanism module is equipped with a detection head; the number of the detection heads is determined according to the size of the target to be detected in the image to be detected.

7. The method according to claim 1, characterized in that The prediction layer includes a Head layer, which is used to calculate the loss in the model training stage and process the multiple prediction features in the model application stage to obtain the target detection result.

Images