Improved YOLOv8-based lightweight garbage sorting robot target detection method

Abstract

The invention provides a lightweight garbage sorting robot target detection method based on improved YOLOv8, and aims to improve the detection speed and accuracy and reduce the consumption of computing resources. According to the method, firstly, a depth camera is used for collecting multi-angle junk images, a data set is expanded through a generative adversarial network (GAN), and the generalization ability of a model is enhanced; and then, an improved YOLOv8 lightweight target detection network model is adopted, the model takes Mobile NetV3 as a backbone network, and an improved SPPF, a parameter-free SimAM attention mechanism and a dysample upsampling module are integrated, so that the target detection performance is improved. In the model training stage, the learning rate is adjusted by dynamically adjusting the batch size and adopting a cosine annealing strategy, and the model performance is optimized. And finally, the trained model is converted into a TensorRT engine, and the TensorRT engine is deployed on NVIDIA Jetson TX2, so that efficient operation in real-time application is ensured. The method is suitable for embedded equipment with limited resources, and the detection speed and accuracy of the garbage sorting robot can be effectively improved.

Description

Technical Field

[0001] The invention belongs to the field of robot vision, and in particular relates to a lightweight target detection method for a garbage sorting mobile robot. Background Art

[0002] With the acceleration of urbanization, garbage disposal has become a major challenge facing the world. Traditional garbage sorting mainly relies on manual operation, which is not only inefficient but also costly, and also poses health and safety risks. In order to improve the efficiency and accuracy of garbage disposal, automated garbage sorting technology has emerged. In recent years, the development of machine vision technology has provided new possibilities for realizing automated garbage sorting.

[0003] Initially, automated waste sorting systems relied mainly on simple sensor technology, which performed preliminary classification by detecting the physical characteristics of waste, such as weight and volume. However, this method cannot effectively identify the specific type and material of waste. With the advancement of machine learning and artificial intelligence, especially the emergence of deep learning algorithms, waste sorting systems based on image recognition have begun to develop gradually. These systems obtain images of waste through cameras and use convolutional neural networks (CNN) to analyze and classify them.

[0004] Deep learning algorithms such as YOLO (You Only Look Once) and Faster R-CNN have performed well in object detection tasks. However, these algorithms are usually computationally complex, have large number of parameters and high hardware requirements, and are not suitable for running on resource-constrained embedded devices. This results in slow detection speed and difficulty in meeting real-time requirements. In addition, the wide variety and different forms of garbage have increased the difficulty of detection. Summary of the invention

[0005] To solve the above technical problems, the purpose of an embodiment of the present invention is to provide a lightweight garbage sorting robot target detection method based on improved YOLOv8, which is specially designed for resource-constrained garbage sorting mobile robots to improve detection speed and accuracy and reduce the consumption of computing resources.

[0006] The present invention adopts the following technical solution:

[0007] A lightweight garbage sorting robot target detection method based on improved YOLOv8 includes the following steps:

[0008] Step S1: Use a depth camera (RGB-D) to collect image information of garbage, obtain garbage images of multiple angles and types, and form a preliminary data set;

[0009] Step S2: Use Generative Adversarial Network (GAN) to expand the preliminary data set. GAN generates diverse junk images to enhance the diversity of the data set and improve the generalization ability of the model.

[0010] Step S3: Use labelimg to perform manual labeling, and divide the training set and the test set into 8:2 to form the final data set;

[0011] Step S4, based on the lightweight improvement of the YOLOv8 target detection model, mainly includes lightweight design of its backbone network and neck network. First, the MobileNetV3 network is used to replace the original feature extraction backbone network. The lightweight design of MobileNetV3 reduces the computational complexity of the model and is suitable for running on resource-constrained embedded devices. Secondly, the improved SPPF-S is used to replace the original SPPF spatial pyramid pooling module to improve the model reasoning speed. Finally, the parameter-free SimAM attention mechanism is introduced after the spatial pyramid pooling module of the model to form a complete backbone network. At the same time, the SimAM attention mechanism is also introduced in the neck of the model, and the lightweight dynamic upsampling dysample module is used to replace the original upsampling module in the neck to reduce the number of model parameters and calculations, and further improve network performance;

[0012] Step S5: Use the improved YOLOv8 model to train the labeled data set. Adjust the batch size, learning rate and other related parameters to optimize the model performance of network training;

[0013] Step S6: Convert the trained pt model to a trt model, use the TensorRT inference optimizer for inference, optimize the inference speed and efficiency of the model on the embedded device, and ensure efficient operation in real-time applications;

[0014] Step S7: deploy the model on the main control chip NVIDIA Jetson TX2 of the garbage sorting mobile robot for actual environment testing to verify the performance and stability of the model in real-time sorting tasks;

[0015] A further technical solution is that in step 1, the depth camera is used to collect image information including five common types of garbage in life: beverage bottles, plastic bags, napkins, cans, and plush toys.

[0016] A further technical solution is that in step S2, the use of a generative adversarial network (GAN) to expand the preliminary data set includes the following specific implementation steps:

[0017] Step S21, prepare a group of junk images and perform standardization processing;

[0018] Step S22: Build a GAN model, where the generator is responsible for generating new images, and the discriminator is used to distinguish between real and generated images. By alternately training the two networks, the generator gradually learns to generate realistic garbage images;

[0019] Step S23: Screen out the generated images with higher quality and merge them with the original data to form a richer data set;

[0020] A further technical solution is that the specific structure of the lightweight improved model in step S4 is: including an input image input module, a backbone feature extraction network, a neck feature fusion module, and a head detection result output module;

[0021] A further technical solution is that the Input image input module specifically performs image input on the final data set formed after step S3, and the specific image input size is 3*640*640;

[0022] A further technical solution is that the Backbone feature extraction backbone network module is specifically a feature extraction backbone network that uses a layer of CBS modules and eight layers of lightweight MobileNetV3 modules for deep stacking;

[0023] The CBS module mentioned above is mainly composed of a 3*3 convolutional layer, a batch normalization layer, and an h-swish activation function layer. The specific form of the h-swish activation function is as follows:

[0024]

[0025] The MobileNetV3 module mentioned above is mainly composed of a 1*1 convolution module, a 3*3 depth-separable convolution module, a SE attention mechanism module, and an inverted residual module;

[0026] A further technical solution is that the neck feature fusion module specifically introduces an improved SPPF-S module, and introduces lightweight dynamic upsampling dysample upsampling in the feature extraction PANet to replace the original upsampling and SimAM attention mechanism;

[0027] A further technical solution is that the Head detection result output module specifically generates three detection frames for each detection target after passing through the Neck feature fusion layer, and its intersection-over-union ratio threshold is set to 0.6, and the detection frame above the threshold is output;

[0028] A further technical solution is that the YOLOv8 model training in step S5 specifically includes setting the training batch size to 8, using the cosine annealing strategy to adjust the learning rate, which is conducive to the model approaching the optimal solution more carefully in the later stage of training, and using Dropout in the training process to reduce overfitting and improve the generalization performance of the model;

[0029] A further technical solution is that the model conversion of the trained model in step S6 specifically includes the following specific implementation steps:

[0030] Step S61: Convert the trained deep learning model PT format to ONNX format to facilitate the subsequent TensorRT engine construction;

[0031] Step S62: Use TensorRT's builder and parser to convert the ONNX format model into the TensorRT engine to optimize the model's reasoning performance;

[0032] Step S63: Using the built TensorRT engine, perform model reasoning on the JetsonNX platform by allocating input and output buffers, creating an execution context, and executing the reasoning process;

[0033] Step S64, further optimize the reasoning performance by setting the TensorRT engine precision mode to INT8; BRIEF DESCRIPTION OF THE DRAWINGS

[0034] Figure 1 is a flow chart of the steps of the present invention;

[0035] Figure 2 This is a structural diagram of a lightweight target detection model of the present invention;

[0036] Figure 3 It is the structural diagram of the CBS model of the present invention;

[0037] Figure 4 This is a structural diagram of the MobileNetV3 model of the present invention;

[0038] Figure 5 This is the structural diagram of the improved model SPPF-S of the present invention;

[0039] Figure 6 It is the SimAM attention mechanism module of the present invention; DETAILED DESCRIPTION

[0040] In order to make the purpose, technical solution and advantages of the present invention clearer, the technical solution of the present invention is described clearly and completely below. Obviously, the described embodiments are 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 creative work are within the scope of protection of the present invention.

[0041] In the detailed description of the present invention, it should be noted that the terms “include” and “comprising” refer to the existence of the described features, components, steps, operations, elements, etc., but do not mean to exclude the existence or possibility of adding other features, components, steps, operations, elements, etc.

[0042] In addition, the terms in the present invention are only used to describe specific embodiments and are not intended to limit the scope of the present invention. In the description of the present invention and the corresponding claims, unless otherwise clearly stated, the words "a", "an" or "the" in the singular form should be understood to include their plural forms.

[0043] Furthermore, the term "and / or" in the present invention refers to any combination of the listed items, including a single combination and all possible combinations, and does not exclude the possibility of these combinations.

[0044] The accompanying drawings show several schematic diagrams of the structures of the embodiments of the present invention. These figures are not drawn to actual scale, some details may be enlarged for clarity, and some details may be omitted. The shapes of the regions and layers shown in the figures and their relative sizes and positional relationships are schematic, and may differ in actual applications due to manufacturing tolerances or technical limitations. Those skilled in the art can design regions / layers with different shapes, sizes, and relative positions according to actual needs.

[0045] The present invention provides a lightweight garbage sorting robot target detection method based on improved YOLOv8, which is specially designed for resource-constrained garbage sorting mobile robots to improve detection speed and accuracy and reduce the consumption of computing resources. The technical solution of the present invention is specifically described below in conjunction with the accompanying drawings.

[0046] like Figure 1 The method specifically comprises the following steps:

[0047] Step S1: Use a depth camera (RGB-D) to collect image information of garbage, obtain garbage images of multiple angles and types, and form a preliminary data set;

[0048] Step S2: Use Generative Adversarial Network (GAN) to expand the preliminary data set. GAN generates diverse junk images to enhance the diversity of the data set and improve the generalization ability of the model.

[0049] Step S3: Use labelimg to perform manual labeling, and divide the training set and the test set into 8:2 to form the final data set;

[0050] Step S4, based on the lightweight improvement of the YOLOv8 target detection model, mainly includes lightweight design of its backbone network and neck network. First, the MobileNetV3 network is used to replace the original feature extraction backbone network. The lightweight design of MobileNetV3 reduces the computational complexity of the model and is suitable for running on resource-constrained embedded devices. Secondly, the improved SPPF-S is used to replace the original SPPF spatial pyramid pooling module to improve the model reasoning speed. Finally, the parameter-free SimAM attention mechanism is introduced after the spatial pyramid pooling module of the model to form a complete backbone network. At the same time, the SimAM attention mechanism is also introduced in the neck of the model, and the lightweight dynamic upsampling dysample module is used to replace the original upsampling module in the neck to reduce the number of model parameters and calculations, and further improve network performance. The specific model structure is as follows: Figure 2 As shown;

[0051] Step S5: Use the improved YOLOv8 model to train the labeled data set. Adjust the batch size, learning rate and other related parameters to optimize the model performance of network training;

[0052] Step S6: Convert the trained pt model to a trt model, use the TensorRT inference optimizer for inference, optimize the inference speed and efficiency of the model on the embedded device, and ensure efficient operation in real-time applications;

[0053] Step S7: deploy the model on the main control chip NVIDIA Jetson TX2 of the garbage sorting mobile robot for actual environment testing to verify the performance and stability of the model in real-time sorting tasks;

[0054] A further technical solution is that in step 1, the depth camera is used to collect image information including five common types of garbage in life: beverage bottles, plastic bags, napkins, cans, and plush toys. The resolution of the collected images is a high-definition format of 1280x720 pixels, and 50 images are collected for each detection target at different viewing angles;

[0055] A further technical solution is that in step S2, the use of a generative adversarial network (GAN) to expand the preliminary data set includes the following specific implementation steps:

[0056] Step S21, prepare a group of junk images and perform standardization processing;

[0057] Step S22: Build a GAN model, where the generator is responsible for generating new images, and the discriminator is used to distinguish between real and generated images. By alternately training the two networks, the generator gradually learns to generate realistic garbage images;

[0058] Step S23: Filter out the generated images with higher quality and merge them with the original data to form a richer data set.

[0059] A further technical solution is that the standardization process in step S21 is specifically as follows:

[0060] Step S211, resize the image to a uniform resolution of 640*640 to meet the needs of model training;

[0061] Step S212: normalize the pixel values ​​to between [0, 1] to accelerate model convergence.

[0062] A further technical solution is that the specific structure of the lightweight improved model in step S4 is: including an input image input module, a backbone feature extraction network, a neck feature fusion module, and a head detection result output module;

[0063] A further technical solution is that the Input image input module specifically performs image input on the final data set formed after step S3, and the specific image input size is 3*640*640;

[0064] A further technical solution is that the Backbone feature extraction backbone network module is specifically a feature extraction backbone network that uses a layer of CBS modules and eight layers of lightweight MobileNetV3 modules for deep stacking;

[0065] The CBS module mentioned above is mainly composed of a 3*3 convolutional layer, a batch normalization layer, and an h-swish activation function layer. The specific form of the h-swish activation function is as follows:

[0066]

[0067] The CBS module is as follows Figure 3 As shown;

[0068] The MobileNetV3 module mentioned above is mainly composed of a 1*1 convolution module, a 3*3 depth-separable convolution module, a SE attention mechanism module, and an inverted residual module;

[0069] The MobileNetV3 module is as follows Figure 4 As shown;

[0070] A further technical solution is that the Neck feature fusion module specifically introduces an improved SPPF-S module, and introduces lightweight dynamic upsampling DySample upsampling in the feature extraction PANet to replace the original upsampling and SimAM attention mechanism;

[0071] A further technical solution is that the improved SPPF-S module is specifically formed by converting the 5*5 average pooling module into two 3*3 average pooling modules. The original SPPF module is formed by sequentially splicing three layers of 5*5 pooling modules, and the result of each layer of pooling is concat spliced ​​with the original input. Now the three-layer 5*5 pooling module is changed into a 6-layer 3*3 pooling module, and the second pooling module, the fourth pooling module, and the sixth pooling module are output to the original input for concat splicing, which reduces the number of model parameters and the amount of calculation. At the same time, the smaller pooling core can capture more local characteristics and improve the model performance. At the same time, the SPPF-S input head activation function RELU is changed to the activation function RELU6. Compared with the traditional RELU, RELU6 is more suitable for deployment on the mobile terminal and maintains good numerical resolution at low precision. The specific form of its RELU6 activation function is as follows:

[0072]

[0073] The improved SPPF-S module is as follows Figure 5 As shown;

[0074] A further technical solution is to introduce the SimAM attention mechanism into the neck. SimAM is a parameter-free, lightweight attention mechanism that generates attention weights by calculating the similarity between each pixel in the feature map and its adjacent pixels, and then applies the Sigmoid activation function and multiplies it with the original feature map to enhance important features and suppress unimportant features, thereby improving model performance.

[0075] The improved SimAM module is as follows Figure 6 As shown;

[0076] A further technical solution is that the DySample module introduced in PANnet is an ultra-lightweight and efficient dynamic upsampler that achieves upsampling by learning the location of sampling points instead of relying on dynamic convolution kernels. This method avoids the time-consuming dynamic convolution and additional subnetworks in traditional dynamic upsamplers, thereby reducing the amount of computation and resource consumption. DySample does not require a customized CUDA package and has fewer parameters, floating-point operations, GPU memory and latency. The specific steps are as follows:

[0077] S41, input feature map: an input feature map of size C×H×W;

[0078] S42, generating an offset of 2s^2×H×W, where the offset represents the displacement of the sampling point relative to the original grid point;

[0079] S43, PixelShuffle operation: convert the offset into a suitable spatial resolution through the PixelShuffle operation;

[0080] A further technical solution is that the Head detection result output module specifically generates three detection frames for each detection target after passing through the Neck feature fusion layer, and its intersection-over-union ratio threshold is set to 0.6, and the detection frame above the threshold is output;

[0081] A further technical solution is that the YOLOv8 model training in step S5 specifically includes setting the training batch size to 8, using the cosine annealing strategy to adjust the learning rate, which is conducive to the model approaching the optimal solution more carefully in the later stage of training, and using Dropout in the training process to reduce overfitting and improve the generalization performance of the model;

[0082] The core idea of ​​the above cosine annealing strategy is to simulate the periodic changes of the cosine function to dynamically adjust the learning rate to help the model converge better and improve the generalization performance. The basic form is to let the learning rate gradually decrease along the cosine curve. The learning rate starts from the maximum value and gradually decreases to the minimum value along the cosine curve in multiple periods;

[0083] Dropout is used in the above training process to prevent overfitting of the neural network. The core idea of ​​Dropout technology is to randomly "discard" some neurons in the network during the training process to reduce the complex co-adaptation relationship between neurons and improve the generalization ability of the model.

[0084] A further technical solution is that the model conversion of the trained model in step S6 specifically includes the following specific implementation steps:

[0085] Step S61: Convert the trained deep learning model PT format to ONNX format to facilitate the subsequent TensorRT engine construction;

[0086] Step S62: Use TensorRT's builder and parser to convert the ONNX format model into the TensorRT engine to optimize the model's reasoning performance;

[0087] Step S63: Using the built TensorRT engine, perform model reasoning on the JetsonNX platform by allocating input and output buffers, creating an execution context, and executing the reasoning process;

[0088] Step S64, further optimize the reasoning performance by setting the TensorRT engine precision mode to INT8;

[0089] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, rather than to limit it. Although the present invention has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or make equivalent replacements for some of the technical features therein. However, these modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the embodiments of the present invention.

Claims

1. A lightweight garbage sorting robot target detection method based on improved YOLOv8, characterized in that: The steps include: Step S1: Use a depth camera (RGB-D) to collect image information of garbage, obtain garbage images of multiple angles and types, and form a preliminary data set. Step S2: Use a generative adversarial network (GAN) to expand the preliminary data set. GAN generates diverse junk images to enhance the diversity of the data set and improve the generalization ability of the model. Step S3: Use label i mg for manual annotation and divide the training set and test set into 8:2 to form the final data set. Step S4, based on the lightweight improvement of the YOLOv8 target detection model, mainly includes lightweight design of its backbone network and neck network. First, the improved MobileNetV3 network is used to replace the original feature extraction backbone network. The improved MobileNetV3 lightweight design reduces the computational complexity of the model and is suitable for running on resource-constrained embedded devices. Secondly, the improved SPPF-S is used to replace the original SPPF spatial pyramid pooling module to improve the model reasoning speed. Finally, the parameter-free SimAM attention mechanism is introduced after the spatial pyramid pooling module of the model to form a complete backbone network. At the same time, the SimAM attention mechanism is also introduced in the neck of the model, and the lightweight dynamic upsampling dysample module is used to replace the original upsampling module in the neck to reduce the number of model parameters and the amount of calculation, thereby further improving network performance. Step S5: Use the improved YOLOv8 model to train the labeled data set, adjust the batch size, learning rate and other related parameters, and optimize the model performance of the network training. Step S6: Convert the trained pt model to a trt model, use the TensorRT inference optimizer for inference, optimize the inference speed and efficiency of the model on the embedded device, and ensure efficient operation in real-time applications. Step S7: Deploy the model on the main control chip NVIDIA Jetson TX2 of the garbage sorting mobile robot for actual environment testing to verify the performance and stability of the model in real-time sorting tasks.

2. A lightweight garbage sorting robot target detection method based on improved YOLOv8 according to claim 1, characterized in that: The specific steps of S2 are: Step S21, prepare a group of junk images and perform standardization processing; Step S22: construct a GAN model, in which the generator is responsible for generating new images, and the discriminator is used to distinguish between real and generated images. By alternately training the two networks, the generator gradually learns to generate realistic garbage images. Step S23: Filter out the generated images with higher quality and merge them with the original data to form a richer data set.

3. A lightweight garbage sorting robot target detection method based on improved YOLOv8 according to claim 1, characterized in that: The specific structure of the S4 lightweight improved model is: Input image input module, Backbone feature extraction backbone network, Neck neck feature fusion module, Head detection result output module.

4. A lightweight garbage sorting robot target detection method based on improved YOLOv8 according to claim 3, characterized in that: The Input image input module specifically inputs the image of the final data set formed after step S3, and its specific image input size is 3*640*640. The Backbone feature extraction backbone network module specifically adopts a layer of CBS module and eight layers of lightweight MobileNetV3 modules to perform a deeply stacked feature extraction backbone network. The Neck neck feature fusion module specifically introduces an improved SPPF-S module, and introduces lightweight dynamic upsampling DySample upsampling in the feature fusion PANet instead of the original upsampling and SimAM attention mechanism.

5. A lightweight garbage sorting robot target detection method based on improved YOLOv8 according to claim 4, characterized in that: The improved SPPF-S module is specifically formed by converting the 5*5 average pooling module into two 3*3 average pooling modules. The original SPPF module is formed by sequentially splicing three layers of 5*5 pooling modules, and the result of each layer of pooling is concat spliced ​​with the original input. Now the three layers of 5*5 pooling modules are changed into 6 layers of 3*3 pooling modules, and the second pooling module, the fourth pooling module, and the sixth pooling module are output to the original input for concat splicing, which reduces the number of model parameters and the amount of calculation. At the same time, the smaller pooling core can capture more local characteristics and improve the model performance. At the same time, the SPPF-S input head activation function RELU is changed to the activation function RELU6. Compared with the traditional RELU, RELU6 is more suitable for deployment on the mobile terminal and maintains good numerical resolution under low precision. The specific form of its RELU6 activation function is as follows:

6. A lightweight garbage sorting robot target detection method based on improved YOLOv8 according to claim 4, characterized in that: The DySample upsampling module introduced in PANnet is an ultra-lightweight and efficient dynamic upsampler. It achieves upsampling by learning the location of sampling points instead of relying on dynamic convolution kernels. The specific steps are as follows: S41, input feature map: an input feature map of size C×H×W; S42, generating an offset of 2s^2×H×W, where the offset represents the displacement of the sampling point relative to the original grid point; S43, Pixel Shuffle operation: The offset is converted into a suitable spatial resolution through the Pixel Shuffle operation.

7. A lightweight garbage sorting robot target detection method based on improved YOLOv8 according to claim 1, characterized in that: The step S6 of converting the trained model into a model specifically includes the following specific implementation steps: Step S61: Convert the trained deep learning model PT format to ONNX format to facilitate the subsequent TensorRT engine construction; Step S62: Use TensorRT's builder and parser to convert the ONNX format model into the TensorRT engine to optimize the model's reasoning performance; Step S63: Using the built TensorRT engine, perform model reasoning on the Jetson NX platform by allocating input and output buffers, creating an execution context, and executing the reasoning process; Step S64: Further optimize the inference performance by setting the TensorRT engine precision mode to I NT8.

Images