A lightweight garbage detection method based on improved YOLOv3

Abstract

The present application relates to a kind of lightweight garbage detection method based on improved YOLOv3, belong to computer vision field.The detection method includes: obtaining data set and processing;MobileNetv2 lightweight backbone network is constructed to replace Darknet53 original backbone network in YOLOv3;new backbone network is fused with improved ECA attention mechanism to carry out feature extraction;the feature map extracted is fused with multi-scale feature;The feature map after fusion is sent into detection head and is predicted to output, and target information is sent into loss function iterative training model;The optimal weight of model training is loaded in model, and the image to be detected is verified and output.The present application realizes real-time garbage detection, model parameter quantity is less, detection speed is fast, and the performance requirement needed for deployment embedded device is reduced.

Description

Technical Field

[0001] This invention relates to the field of computer vision, and more specifically to a lightweight garbage detection method based on an improved YOLOv3. Background Technology

[0002] my country is a major producer of waste, generating tens of millions of tons of garbage every day. However, problems such as low sorting and recycling efficiency still exist in waste management. Using computer technology in conjunction with automated equipment to replace manual labor would greatly improve the efficiency of each stage of waste management. The key to this is a robust waste identification and detection algorithm capable of sensing and processing waste.

[0003] In recent years, with the rapid development of artificial intelligence technology, object detection, as a branch of deep learning computer vision, has also achieved some breakthroughs. With these technological advancements, object detection has been widely applied in various fields such as facial recognition, healthcare, and autonomous driving. For waste sorting, using object detection technology to determine the location and type of waste can help people quickly classify it. Currently, object detection technologies mainly fall into two categories: one is two-stage detection algorithms, represented by Faster-RCNN, which have high detection accuracy but slow speed; the other is one-stage detection algorithms, represented by SSD and YOLO, which have fast detection speed but lower accuracy than two-stage detection algorithms.

[0004] Although waste detection technology has made significant progress, some problems still exist. For waste detection scenarios, target detection algorithms often need to be deployed in embedded devices with limited computing power, which poses a challenge to the device's computing power and memory. Furthermore, many existing methods are large in size and slow in detection speed, failing to meet the requirements of accuracy and real-time performance. Specific related patent documents, such as "A Method for Detecting Kitchen Waste Based on Neural Networks" (Application Publication No. CN115205521A) and "A Waste Detection System and Method Based on Deep Learning" (Application Publication No. CN115240134A), exhibit problems such as large network parameter counts and slow detection speeds. Summary of the Invention

[0005] To overcome the shortcomings of the above-mentioned technologies, this invention provides a lightweight garbage detection method based on improved YOLOv3, which effectively reduces the number of model parameters and improves the target detection speed of the model while ensuring detection accuracy.

[0006] The technical solution adopted in this invention is:

[0007] A lightweight garbage detection method based on improved YOLOv3 includes the following steps:

[0008] Step 1: Obtain junk images and perform data augmentation on them. Then, combine the augmented junk images into a training set and a validation set in a 7:3 ratio, which will be used for training and validation of the initial improved YOLOv3 model, respectively.

[0009] Step 2: Construct the MobileNetv2 lightweight network to replace the original Darknet53 backbone network in the traditional YOLOv3 model, obtaining the initial improved YOLOv3 model backbone network. The specific construction steps of the MobileNetv2 lightweight network are as follows:

[0010] Step 2.1: Construct the inverse residual structure A using channel-wise convolution and pointwise convolution. Specifically:

[0011] The channel-wise convolution is a convolution with a kernel size of 3×3, a stride of 1 or 2, and a number of groups equal to the number of output channels. The pointwise convolution is a convolution with a kernel size of 1×1, a stride of 1, and a number of groups of 1. The inverse residual structure A consists of a first pointwise convolution, a channel-wise convolution, a first batch normalization layer, a ReLU activation function layer, a second pointwise convolution, and a second batch normalization layer arranged in sequence. When the stride of the channel-wise convolution is 2, the output of the second batch normalization layer is the output of the inverse residual structure A. When the stride of the channel-wise convolution is 1, the sum of the output of the second batch normalization layer and the input of the first pointwise convolution is the output of the inverse residual structure A.

[0012] Step 2.2: Insert the improved ECA attention mechanism into the inverted residual structure A built in step 2.1. The insertion position is between the ReLU activation function layer and the second pointwise convolution.

[0013] The inverse residual structure containing a channel-wise convolution with a stride of 2 and an improved ECA attention mechanism is denoted as inverse residual structure B, and the inverse residual structure containing a channel-wise convolution with a stride of 1 and an improved ECA attention mechanism is denoted as inverse residual structure C.

[0014] Step 2.3: The MobileNetv2 lightweight network is constructed using inverted residual structure B, inverted residual structure C, and intermediate convolution. The intermediate convolution is a convolution with a kernel size of 3×3, a stride of 2, and an output channel of 32.

[0015] Step 3: Use the MobileNetv2 lightweight network to extract features from the garbage images in the training set to obtain feature map M4 and feature map M5;

[0016] Step 4: Perform multi-scale feature fusion on feature map M4 and feature map M5 obtained in step 3 to obtain feature map P5 and feature map P4. Specifically: convolve feature map M5 twice with a kernel size of 3×3 and a stride of 1 to obtain feature map P5; upsample feature map M5 and then add it to feature map M4 on the channel to obtain feature map P4.

[0017] Step 5: The feature maps P4 and P5 generated in Step 4 are fed into the detection head to extract the predicted bounding box information. Then, the predicted bounding box information and the target label information are fed into the loss function to iteratively train the initial improved YOLOv3 model and save the optimal weights.

[0018] The predicted bounding box information includes the predicted bounding box location coordinates of the feature map, the width and height of the bounding box, and the category information of the bounding box;

[0019] The target label information includes the true bounding box location coordinates of the feature map, the width and height of the bounding box, and the category information of the bounding box;

[0020] Step 6: Load the optimal weights saved in Step 5 into the initial improved YOLOv3 model to obtain the improved YOLOv3 model with the best performance, and define the improved YOLOv3 model with the best performance as the final improved YOLOv3 model.

[0021] Step 7: Feed the garbage images in the validation set into the final improved YOLOv3 model for validation and obtain the final detection results.

[0022] Preferably, the garbage images in step 1 include garbage images from the Huawei competition's publicly available garbage dataset and garbage images taken with a mobile phone; the data augmentation refers to adjusting the angle, saturation, and exposure of the garbage images, and the data augmentation methods include cropping, translation, brightness enhancement, adding Gaussian noise, rotation, and mirroring. After adjustment, the garbage image has a pixel size of 416×416.

[0023] Preferably, the structure of the improved ECA attention mechanism described in step 2 includes a globally average pooling layer and a one-dimensional convolutional layer connected in sequence. After the one-dimensional convolutional layer, it is separated into two channels. The feature map input to the globally average pooling layer is denoted as feature map γ, and the feature maps output by the two channels are denoted as feature map α and feature map β, respectively. The height, width, and number of channels of feature map α, feature map β, and feature map γ are multiplied together to output a feature map F of [H×W×C]. This feature map F is the output of the improved ECA attention mechanism, where H, W, and C are the height, width, and number of channels of feature map F, respectively.

[0024] Preferably, the specific method for feeding the garbage images in the validation set into the final improved YOLOv3 model for validation in step 7 is as follows: the garbage images in the validation set are fed into the trained final improved YOLOv3 model, and two sets of feature images are output with sizes of [13, 13, 256] and [16, 16, 512], respectively. Then, the predicted detection box information is output through the detection head, and redundant detection boxes are removed by using non-maximum suppression operation to generate the final detection result.

[0025] The three values ​​in [] represent the height, width, and number of channels of the feature map, respectively.

[0026] Preferably, the specific steps for constructing the MobileNetv2 lightweight network using inverted residual structure B, inverted residual structure C, and in-middle convolution are as follows:

[0027] S1, using in-convolution for downsampling;

[0028] S2, downsampling is performed using inverted residual structure B, and then passed through one set of inverted residual structure C;

[0029] S3, use the inverted residual structure B for downsampling, and then pass through 2 sets of inverted residual structures C;

[0030] S4, downsampling is performed using inverted residual structure B, and then the data is processed through 3 sets of inverted residual structures C;

[0031] S5, after passing through 3 sets of inverted residual structures C;

[0032] S6 uses inverted residual structure B for downsampling, and then passes through 3 sets of inverted residual structures C.

[0033] Compared with the prior art, the beneficial effects of the present invention are mainly reflected in:

[0034] (1) This invention uses the lightweight network MobileNetv2 instead of the original Darknet53 as the backbone network, which significantly reduces the number of model parameters and improves the detection speed of the algorithm.

[0035] (2) This invention integrates the improved ECA attention mechanism into the backbone network. The improved ECA attention mechanism assigns two weights to each channel of the input feature map, representing the horizontal and vertical weights of the feature map respectively, thereby enhancing the feature extraction capability of the MobileNetv2 backbone network and significantly improving the model accuracy with a small increase in the number of parameters.

[0036] (3) In view of the limited computing power and memory of actual embedded devices, this invention proposes a lightweight garbage detection method based on improved YOLOv3. The model has a small number of parameters and fast detection speed while ensuring accuracy, which can be better deployed on embedded devices. Attached Figure Description

[0037] Figure 1 This is a flowchart of the lightweight garbage detection method based on the improved YOLOv3 of the present invention;

[0038] Figure 2 This is a structural diagram of the inverted residual structure in an embodiment of the present invention;

[0039] Figure 3 This is a structural diagram of the improved ECA attention mechanism of the present invention. Detailed Implementation

[0040] The implementation of the present invention will now be further described with reference to the accompanying drawings.

[0041] Figure 1 This is a flowchart of a lightweight garbage detection method based on improved YOLOv3 in an embodiment of the present invention. Figure 1 As can be seen, the detection method of the present invention includes the following steps:

[0042] Step 1: Obtain junk images and perform data augmentation on them. Then, combine the augmented junk images into a training set and a validation set in a 7:3 ratio, which will be used for training and validation of the initial improved YOLOv3 model, respectively.

[0043] In this embodiment, the garbage images include garbage images from the Huawei competition's publicly available garbage dataset and garbage images taken with a mobile phone; the data augmentation refers to adjusting the angle, saturation, and exposure of the garbage images. The data augmentation methods include cropping, translation, brightness enhancement, adding Gaussian noise, rotation, and mirroring. After adjustment, the garbage image has a pixel size of 416×416.

[0044] Step 2: Construct the MobileNetv2 lightweight network to replace the original Darknet53 backbone network in the traditional YOLOv3 model, obtaining the initial improved YOLOv3 model backbone network. The specific construction steps of the MobileNetv2 lightweight network are as follows:

[0045] Step 2.1: Construct the inverse residual structure A using channel-wise convolution and pointwise convolution. Specifically:

[0046] The channel-wise convolution is a convolution with a kernel size of 3×3, a stride of 1 or 2, and a group number equal to the number of output channels. The pointwise convolution is a convolution with a kernel size of 1×1, a stride of 1, and a group number of 1. The inverse residual structure A consists of a first pointwise convolution, a channel-wise convolution, a first batch normalization layer, a ReLU activation function layer, a second pointwise convolution, and a second batch normalization layer arranged in sequence. When the stride of the channel-wise convolution is 2, the output of the second batch normalization layer is the output of the inverse residual structure A. When the stride of the channel-wise convolution is 1, the sum of the output of the second batch normalization layer and the input of the first pointwise convolution is the output of the inverse residual structure A.

[0047] Figure 2 This is a structural diagram of the inverted residual structure in an embodiment of the present invention.

[0048] Step 2.2: Insert the improved ECA attention mechanism into the inverted residual structure A built in step 2.1. The insertion position is between the ReLU activation function layer and the second pointwise convolution.

[0049] In this embodiment of the invention, the structure of the improved ECA attention mechanism includes a globally average pooling layer and a one-dimensional convolutional layer connected in sequence. After the one-dimensional convolutional layer, it is separated into two channels. The feature map input to the globally average pooling layer is denoted as feature map γ, and the feature maps output by the two channels are denoted as feature map α and feature map β, respectively. The height, width, and number of channels of feature map α, feature map β, and feature map γ are multiplied together to output a feature map F of [H×W×C]. This feature map F is the output of the improved ECA attention mechanism, where H, W, and C are the height, width, and number of channels of feature map F, respectively.

[0050] Figure 3 This is a structural diagram of the improved ECA attention mechanism of the present invention.

[0051] The inverse residual structure containing a channel-wise convolution with a stride of 2 and an improved ECA attention mechanism is denoted as inverse residual structure B, and the inverse residual structure containing a channel-wise convolution with a stride of 1 and an improved ECA attention mechanism is denoted as inverse residual structure C.

[0052] Step 2.3: The MobileNetv2 lightweight network is constructed using inverted residual structure B, inverted residual structure C, and intermediate convolution. The intermediate convolution is a convolution with a kernel size of 3×3, a stride of 2, and an output channel of 32.

[0053] In this embodiment of the invention, the specific steps are as follows:

[0054] S1, using in-convolution for downsampling;

[0055] S2, downsampling is performed using inverted residual structure B, and then passed through one set of inverted residual structure C;

[0056] S3, use the inverted residual structure B for downsampling, and then pass through 2 sets of inverted residual structures C;

[0057] S4, downsampling is performed using inverted residual structure B, and then the data is processed through 3 sets of inverted residual structures C;

[0058] S5, after passing through 3 sets of inverted residual structures C;

[0059] S6 uses inverted residual structure B for downsampling, and then passes through 3 sets of inverted residual structures C.

[0060] Step 3: Use the MobileNetv2 lightweight network built in Step 2 to extract features from the garbage images in the training set to obtain feature map M4 and feature map M5.

[0061] In this embodiment of the invention, the specific steps for extracting features from junk images in the training set using the MobileNetv2 lightweight network constructed in step 2 are as follows:

[0062] The garbage images in the training set are input into S1 for feature extraction, and the output feature map M1 is output. Feature map M1 is processed by S2 for feature extraction, and the output feature map M2 is output. Feature map M2 is processed by S3 for feature extraction, and the output feature map M3 is output. Feature map M3 is processed by S4 for feature extraction, and the output feature map M4 is output. Feature map M4 is processed by S5 and S6 for feature extraction, and the output feature map M5 is output.

[0063] Step 4: Perform multi-scale feature fusion on feature map M4 and feature map M5 obtained in step 3 to obtain feature map P5 and feature map P4. Specifically: convolve feature map M5 twice with a kernel size of 3×3 and a stride of 1 to obtain feature map P5; upsample feature map M5 and then add it to feature map M4 on the channel to obtain feature map P4.

[0064] Step 5: The feature maps P4 and P5 generated in Step 4 are fed into the detection head to extract the predicted bounding box information. Then, the predicted bounding box information and the target label information are fed into the loss function to iteratively train the initial improved YOLOv3 model and save the optimal weights.

[0065] The predicted bounding box information includes the predicted bounding box location coordinates of the feature map, the width and height of the bounding box, and the category information of the bounding box;

[0066] The target label information includes the true bounding box location coordinates of the feature map, the width and height of the bounding box, and the category information to which the bounding box belongs.

[0067] Step 6: Load the optimal weights saved in Step 5 into the initial improved YOLOv3 model to obtain the improved YOLOv3 model with the best performance, and define the improved YOLOv3 model with the best performance as the final improved YOLOv3 model.

[0068] Step 7: Feed the garbage images in the validation set into the final improved YOLOv3 model for validation and obtain the final detection results.

[0069] In this embodiment of the invention, the specific method for sending the garbage images in the validation set into the final improved YOLOv3 model for validation is as follows: the garbage images in the validation set are sent into the trained final improved YOLOv3 model, and two sets of feature images are output with sizes of [13, 13, 256] and [16, 16, 512], respectively. Then, the predicted detection box information is output through the detection head, and redundant detection boxes are removed by using non-maximum suppression operation to generate the final detection result.

[0070] The three values ​​in [] represent the height, width, and number of channels of the feature map, respectively.

[0071] In this embodiment of the invention, the expression for the loss function Loss in step 5 is as follows:

[0072] Loss = λ1Loss cls +λ2Loss obj +λ3Loss loc

[0073] Where λ1 is the classification balance coefficient, λ2 is the confidence balance coefficient, λ3 is the location balance coefficient, and Loss is the classification balance coefficient. cls For classification loss, Loss obj For confidence loss, Loss loc To pinpoint the loss.

[0074] Classification Loss cls The cross-entropy loss function is expressed as follows:

[0075]

[0076] Where S is the number of grids, and B is the number of bounding boxes generated by each grid; This indicates whether the j-th bounding box is responsible for predicting the target object in the i-th grid. If the bounding box is responsible for predicting the target, then... Equals 1, otherwise, Equal to 0; It is the predicted probability of the class c of the j-th bounding box in grid i. It represents the true probability, and classes represents the category to which the garbage in the garbage image belongs.

[0077] Confidence Loss obj The binary cross-entropy loss function is expressed as follows:

[0078]

[0079] Where, λ noobj It is the weight of the confidence error in the loss function when the predicted bounding box does not contain the predicted target. It is the prediction confidence of the j-th bounding box in grid i. It is the actual confidence level.

[0080] Location Loss loc The CIoU loss function is expressed as follows:

[0081]

[0082] Where IoU is the overlap ratio, which is the overlap ratio between the predicted bounding box (pred) and the ground truth bounding box (gt), ρ 2 (b,b st ) represents the Euclidean distance between the center points of the predicted bounding box and the ground truth bounding box, and b represents the center point of the predicted bounding box. st denoted by , c represents the diagonal distance of the smallest enclosed region that can simultaneously contain both the predicted and ground truth bounding boxes; α and v are set to penalize aspect ratio, with α being a positive trade-off parameter and v measuring the consistency of aspect ratio; the expressions for IoU, α, and v are as follows:

[0083]

[0084]

[0085]

[0086] Among them, Area pred It is the area of ​​the predicted bounding box, Area gt It is the area of ​​the actual bounding box, w gt It is the width of the actual bounding box, h gt w is the height of the ground truth bounding box, w is the width of the predicted bounding box, and h is the height of the predicted bounding box.

[0087] To verify the effectiveness of the present invention, the detection results of the improved YOLOv3 model of the present invention, the traditional YOLOv3 model, and the improved YOLOv3 model of the present invention without the improved ECA attention mechanism were compared. The specific data are shown in Table 1.

[0088]

[0089]

[0090] The comparison results show that this method greatly reduces the number of model parameters and improves the detection speed while ensuring detection accuracy. The improved ECA attention mechanism significantly improves detection accuracy while sacrificing a small number of parameters and detection speed.

Claims

1. A lightweight garbage detection method based on improved YOLOv3, characterized in that, Includes the following steps: Step 1: Obtain junk images and perform data augmentation on them. Then, combine the data-augmented junk images into a training set and a validation set in a 7:3 ratio, which will be used for training and validation of the initial improved YOLOv3 model, respectively. Step 2: Construct the MobileNetv2 lightweight network to replace the original Darknet53 backbone network in the traditional YOLOv3 model, obtaining the initial improved YOLOv3 model backbone network. The specific construction steps of the MobileNetv2 lightweight network are as follows: Step 2.1: Construct the inverse residual structure A using channel-wise convolution and pointwise convolution. Specifically: The channel-wise convolution is a convolution with a kernel size of 3×3, a stride of 1 or 2, and a number of groups equal to the number of output channels. The pointwise convolution is a convolution with a kernel size of 1×1, a stride of 1, and a number of groups of 1. The inverse residual structure A consists of a first pointwise convolution, a channel-wise convolution, a first batch normalization layer, a ReLU activation function layer, a second pointwise convolution, and a second batch normalization layer arranged in sequence. When the stride of the channel-wise convolution is 2, the output of the second batch normalization layer is the output of the inverse residual structure A. When the stride of the channel-wise convolution is 1, the sum of the output of the second batch normalization layer and the input of the first pointwise convolution is the output of the inverse residual structure A. Step 2.2: Insert the improved ECA attention mechanism into the inverted residual structure A built in step 2.

1. The insertion position is between the ReLU activation function layer and the second pointwise convolution. The inverse residual structure containing a channel-wise convolution with a stride of 2 and an improved ECA attention mechanism is denoted as inverse residual structure B, and the inverse residual structure containing a channel-wise convolution with a stride of 1 and an improved ECA attention mechanism is denoted as inverse residual structure C. Step 2.3: The MobileNetv2 lightweight network is constructed using inverted residual structure B, inverted residual structure C, and intermediate convolution. The intermediate convolution is a convolution with a kernel size of 3×3, a stride of 2, and an output channel of 32. The specific steps for constructing the MobileNetv2 lightweight network using inverted residual structure B, inverted residual structure C, and in-middle convolution are as follows: S1, using in-convolution for downsampling; S2, downsampling is performed using inverted residual structure B, and then passed through one set of inverted residual structure C; S3, use the inverted residual structure B for downsampling, and then pass through 2 sets of inverted residual structures C; S4, downsampling is performed using inverted residual structure B, and then the data is processed through 3 sets of inverted residual structures C; S5, after passing through 3 sets of inverted residual structures C; S6, use inverted residual structure B for downsampling, and then pass through 3 sets of inverted residual structures C; Step 3: Use the MobileNetv2 lightweight network to extract features from the garbage images in the training set to obtain feature map M4 and feature map M5; Step 4: Perform multi-scale feature fusion on feature map M4 and feature map M5 obtained in step 3 to obtain feature map P5 and feature map P4. Specifically: convolve feature map M5 twice with a kernel size of 3×3 and a stride of 1 to obtain feature map P5; upsample feature map M5 and then add it to feature map M4 on the channel to obtain feature map P4. Step 5: The feature maps P4 and P5 generated in Step 4 are fed into the detection head to extract the predicted bounding box information. Then, the predicted bounding box information and the target label information are fed into the loss function to iteratively train the initial improved YOLOv3 model and save the optimal weights. The predicted bounding box information includes the predicted bounding box location coordinates of the feature map, the width and height of the bounding box, and the category information of the bounding box; The target label information includes the true bounding box location coordinates of the feature map, the width and height of the bounding box, and the category information of the bounding box; Step 6: Load the optimal weights saved in Step 5 into the initial improved YOLOv3 model to obtain the improved YOLOv3 model with the best performance, and define the improved YOLOv3 model with the best performance as the final improved YOLOv3 model. Step 7: Feed the garbage images in the validation set into the final improved YOLOv3 model for validation and obtain the final detection results.

2. The lightweight waste detection method based on improved YOLOv3 according to claim 1, characterized in that, The garbage images mentioned in step 1 include garbage images from the Huawei competition's public garbage dataset and garbage images taken with a mobile phone; the data augmentation refers to adjusting the angle, saturation, and exposure of the garbage images. The data augmentation methods include cropping, translation, brightness enhancement, adding Gaussian noise, rotation, and mirroring. After adjustment, the garbage image has a pixel size of 416×416.

3. The lightweight waste detection method based on improved YOLOv3 according to claim 1, characterized in that, The improved ECA attention mechanism described in step 2 consists of a sequentially connected global average pooling layer and a one-dimensional convolutional layer. After the one-dimensional convolutional layer, it is separated into two channels. Let the feature map input to the global average pooling layer be feature map γ, and the feature maps output by the two channels be feature map α and feature map β, respectively. The height, width, and number of channels of feature map α, feature map β, and feature map γ are multiplied together to output a feature map of [H×W×C]. This feature map This is the output of the improved ECA attention mechanism, where H, W, and C are the feature maps, respectively. Height, width, and number of channels.

4. The lightweight waste detection method based on improved YOLOv3 according to claim 1, characterized in that, The specific method for feeding the garbage images in the validation set into the final improved YOLOv3 model for validation in step 7 is as follows: the garbage images in the validation set are fed into the trained final improved YOLOv3 model, and two sets of feature images are output with sizes of [13,13,256] and [16,16,512], respectively. Then, the predicted detection box information is output through the detection head, and redundant detection boxes are removed by non-maximum suppression operation to generate the final detection result. The three values ​​in [] represent the height, width, and number of channels of the feature map, respectively.

Images