A method and system for classifying garbage images

Abstract

The application discloses a kind of rubbish image classification method and system, the method includes the following steps: step 1, obtain rubbish dataset, and divide rubbish dataset into training set and test set;Step 2, improve YOLOv5 model: using CBAM attention mechanism as feature extraction module, using BiFPN structure as feature fusion structure, using boundary box loss function LCIoU;Step 3, the prior anchor frame of training set is obtained using K-means clustering method, and the improved YOLOv5 model is trained using the prior anchor frame of training set, boundary box loss function LCIoU;Step 4, the test set is detected using the trained YOLOv5 model, and the correct rubbish category is obtained by outputting rubbish target detection result. The rubbish recognition accuracy can be improved by the method of the application.

Description

Technical Field

[0001] This invention relates to the field of image technology, and in particular to a method and system for classifying garbage images. Background Technology

[0002] At present, garbage sorting mainly relies on people's awareness and manual sorting by garbage station staff. However, these methods have obvious drawbacks: due to the large variety of garbage types and the difficulty in remembering them, people who sort garbage may not be able to sort it accurately; manual sorting has certain limitations, is inefficient, and the staff's energy is limited, making it impossible to effectively sort large amounts of garbage.

[0003] Therefore, how to improve the accuracy of waste identification and classify waste more accurately is a technical problem that needs to be solved by those skilled in the art. Summary of the Invention

[0004] This invention provides a method and system for garbage image classification, which improves the accuracy of garbage identification and enables more accurate garbage classification. To achieve the above effects, the technical solution of this invention is as follows:

[0005] A method for classifying garbage images includes the following steps:

[0006] Step 1: Obtain the garbage dataset and divide it into a training set and a test set;

[0007] Step 2: Improve the YOLOv5 model: Use the CBAM attention mechanism as the feature extraction module, the BiFPN structure as the feature fusion structure, and the bounding box loss function LCIoU;

[0008] Step 3: Use the K-means clustering method to obtain the prior anchor boxes of the training set, and use the prior anchor boxes of the training set and the bounding box loss function LCIoU to train the improved YOLOv5 model.

[0009] Step 4: Use the trained YOLOv5 model to detect garbage targets on the test set, output the garbage target detection results, and obtain the correct garbage category.

[0010] The purpose of this invention is to improve the accuracy of waste sorting and disposal, making waste sorting more intelligent. It utilizes YOLOv5 to identify various forms of different types of waste, combined with relevant machine vision knowledge, to efficiently sort waste and effectively identify waste types even in complex environments with numerous piles of waste in different forms.

[0011] YOLOv5 is an object recognition and localization algorithm based on deep neural networks. Its biggest advantages are its fast running speed and significantly improved accuracy. It can be used to identify various types of junk targets.

[0012] Furthermore, step 1 specifically involves: obtaining a garbage dataset, labeling the garbage dataset using VoTT or Imglabeling tools, and randomly dividing the labeled dataset into training and test sets.

[0013] Furthermore, in step 1, a batch is extracted from the garbage dataset, and four images are randomly selected from it. These images are then randomly cropped and stitched together to form a new image. This process is repeated until mosaic data augmentation and label smoothing are performed.

[0014] Furthermore, the label smoothing process is a binary cross-entropy label smoothing process, and its formula is:

[0015] y*(1.0-ε)+0.5*ε

[0016] In the formula, y is the one-hot encoding of the garbage image data label; ε is the smoothing factor.

[0017] Furthermore, in step 2, the CBAM attention mechanism is used as the feature extraction module by replacing the first CBH convolutional module in the backbone module with a Focus module, which is used to slice the input image; and replacing the original C3 module of the backbone module with the CBAM attention mechanism, which includes channel attention and spatial attention.

[0018] Furthermore, step 2, which uses the BiFPN structure as the feature fusion structure, involves: removing nodes with only one input edge from the original BiFPN structure; increasing the number of input feature maps on the feature maps with lower resolution; using a concatenation operation during upsampling and a superposition operation during downsampling.

[0019] Furthermore, step 2 includes setting the activation function of the convolutional layer of the YOLOv5 model to the Hardswish function.

[0020] Furthermore, the specific steps of step 3 are as follows:

[0021] Step 3.1: After performing data augmentation on the training set described in Step 1, input the training set into the YOLOv5 model to obtain the predicted values ​​of the object detection labels;

[0022] Step 3.2: Calculate the loss function value using the true value and predicted value of the target detection label;

[0023] Step 3.3: Update the model parameters using the loss function value;

[0024] Step 3.4: Input the test set into the improved YOLOv5 model to obtain the predicted object detection labels;

[0025] Step 3.5: Calculate the bounding box loss function value and test set accuracy using the actual and predicted values ​​of the object detection labels;

[0026] Step 3.6: Determine if the test set accuracy is greater than the maximum accuracy. If it is, save the YOLOv5 model, update the test set accuracy, and proceed to the next step; otherwise, proceed directly to the next step.

[0027] Step 3.7: Determine whether the YOLOv5 model has converged. If it has, proceed to the next step; otherwise, return to step 3.1.

[0028] Step 3.8: Determine if the maximum number of training rounds has been reached. If yes, end the process and output the completed YOLOv5 model; otherwise, return to step 3.1.

[0029] Furthermore, the bounding box loss function LCIoU is defined as follows:

[0030]

[0031]

[0032] In the formula, IOU represents the overlap ratio between the predicted bounding box and the ground truth bounding box, i.e., the union of their intersection and union; ρ 2 p represents the Euclidean distance between the centers of the predicted bounding box and the ground truth bounding box; p, p gt , where are the center coordinates of the predicted bounding box and the ground truth bounding box, respectively; c represents the diagonal length of the minimum bounding box of the predicted bounding box and the ground truth bounding box; α represents the weight coefficient of υ; υ represents the parameter that measures the consistency of aspect ratio; the specific calculation formulas for υ and α are as follows:

[0033]

[0034] In the formula, ω gt and h gt ω represents the width and height of the ground truth bounding box; h and ω represent the width and height of the predicted bounding box.

[0035]

[0036] A garbage image classification system, comprising:

[0037] The image acquisition module is used to acquire the garbage dataset and divide the garbage dataset into training set and test set;

[0038] The model improvement module is used to improve the YOLOv5 model: it uses the CBAM attention mechanism as the feature extraction module, the BiFPN structure as the feature fusion structure, and the bounding box loss function LCIoU.

[0039] The training module is used to obtain the prior anchor boxes of the training set using the K-means clustering method, and to train the improved YOLOv5 model using the prior anchor boxes of the training set and the bounding box loss function LCIoU.

[0040] The output module is used to detect garbage targets on the test set using the trained YOLOv5 model, and output the garbage target detection results to obtain the correct garbage category.

[0041] Compared with the prior art, the beneficial effects of the technical solution of the present invention are:

[0042] This invention improves the YOLOv5 model by using the CBAM attention mechanism as a feature extraction module, the BiFPN structure as a feature fusion structure, and the bounding box loss function LCIoU, thereby increasing the accuracy of garbage identification and classifying garbage more accurately. Attached Figure Description

[0043] The accompanying drawings are for illustrative purposes only and should not be construed as limiting the invention. To better illustrate this embodiment, some components in the drawings may be omitted, enlarged, or reduced, and do not represent the actual product dimensions. It is understandable to those skilled in the art that some well-known structures and their descriptions may be omitted in the drawings.

[0044] Figure 1 This is a flowchart of a garbage image classification method provided in an embodiment of the present invention;

[0045] Figure 2a , 2b These are schematic diagrams of the YOLOv5 model before and after the improvement provided in the embodiments of the present invention.

[0046] Figure 3 This is a schematic diagram of mosaic data enhancement provided in an embodiment of the present invention;

[0047] Figure 4 These are schematic diagrams of the predicted bounding box and the ground truth bounding box provided in the embodiments of the present invention;

[0048] Figure 5 This is a diagram showing the model training results provided in an embodiment of the present invention. Detailed Implementation

[0049] To make the objectives, technical solutions, and advantages of the embodiments of the present invention clearer, the technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0050] Before introducing the embodiments of this application, the relevant terms involved in the embodiments of this application are first explained as follows:

[0051] VoTT: A tool for annotating images;

[0052] Example 1

[0053] For easier understanding, please refer to Figure 1 An embodiment of a garbage image classification method provided by the present invention includes the following steps:

[0054] Step 1: Obtain the spam dataset and divide it into a training set and a test set in a 9:1 ratio. The spam dataset consists of images or video streams to be detected as spam. Step 1 specifically involves:

[0055] The garbage dataset team used cameras to capture images of everyday garbage and collected various types of garbage photos online. They then used VoTT or Imglabeling tools to label the garbage dataset, resulting in a JSON file. The labeled dataset was then randomly divided into independent and non-repeating training and test sets.

[0056] like Figure 3 As shown, a batch is extracted from the garbage dataset, and four images are randomly selected from it. These images are then randomly cropped and stitched together to form a new image. This process is repeated until mosaic data augmentation and label smoothing are performed.

[0057] Manual labeling can lead to mislabeling, which can significantly impact the results. Label smoothing is a common method to address this issue. However, if one-hot labels are used to calculate the loss during training, only the loss for the position labeled 1 is calculated, excluding the loss for the position labeled 0. This results in good performance on the training set, but poor prediction performance during the prediction process because the data source is not the training set. Therefore, label smoothing is necessary for the dataset.

[0058] Specifically, the label smoothing process is a binary cross-entropy label smoothing (smooth_BCE) process, and its formula is as follows:

[0059] y*(1.0-ε)+0.5*ε

[0060] In the formula, y is the one-hot encoding of the garbage image data labels, generally representing a vector; ε is the smoothing factor. After smoothing, the labels of the data in the training set, regardless of whether they are correct or not, are taken into account in the sample loss function, resulting in increased loss and improved learning ability of the model.

[0061] Step 2: Improve the YOLOv5 model: Use the CBAM attention mechanism as the feature extraction module, the BiFPN structure as the feature fusion structure, and the bounding box loss function LCIoU;

[0062] Specifically, the improved YOLOv5 model structure is as follows: Figure 2b As shown; Step 2 uses the CBAM attention mechanism as the feature extraction module as follows: The first CBH convolution module in the backbone module is replaced with the Focus module, which is used to slice the input image; The original C3 module of the backbone module (feature extraction part) is replaced with the CBAM attention mechanism, which includes channel attention and spatial attention;

[0063] The channel attention mechanism uses 1×1×n pooling to process the information between channels of the feature map extracted by the backbone module into an n-dimensional vector, and then uses the improved YOLOv5 model to calculate this vector. This method combines the data information of each channel. The spatial attention mechanism focuses on the length and width of the feature map and pays more attention to the global attributes of the feature map. Therefore, the use of the CBAM attention mechanism can achieve better feature extraction results.

[0064] Understandably, during model training, for example, if the input garbage dataset consists of 4×4×1 images, after the Focus module's slicing operation, a 2×2×4 structure as shown in the right figure is obtained. The image [b,3,640,640] is input into the improved YOLOv5 model, and after Focus slicing, a feature map [b,12,320,320] is obtained, where b represents batch, i.e., the number of images input each time, 3 and 12 represent the number of channels of the image, and 640 and 320 represent the length and width of the image.

[0065] The input garbage dataset is processed by the Focus module, reducing the model's computational load and saving computer and time resources. Simultaneously, it avoids losing important information, preserving image downsampling information and facilitating subsequent feature extraction. The convolutional downsampling layer in the CBAM attention mechanism is used to reduce the layer size, both to reduce computation and to increase the receptive field, allowing subsequent convolutional layers to learn more global information. The CBAM attention mechanism can efficiently adjust the neural network dynamically and learn linear and nonlinear modes, leveraging the advantages of multi-layer networks in feature extraction and enabling it to better adapt to complex and varied garbage shapes and occlusion situations.

[0066] Specifically, the use of the BiFPN structure as the feature fusion structure is as follows: on the basis of the original BiFPN structure, nodes with only one input edge are removed; the number of input feature maps is increased on the feature maps with smaller resolution; a concatenation operation is used in the upsampling process, and a superposition operation is used in the downsampling process.

[0067] When removing nodes with only one input edge, this edge does not contribute much to feature fusion and has little impact on the feature network that fuses different features. Removing this input edge has little impact on the improved YOLOv5 model. While removing nodes with only one input edge, an additional edge is added between the original input node and the output node to facilitate the fusion of more feature map information without increasing the cost too much.

[0068] Set the activation function of the convolutional layer of the YOLOv5 model to the Hardswish function, which has better detection performance.

[0069] Step 3: Use the statistical K-means clustering method to obtain the prior anchor boxes of the training set, and use the prior anchor boxes of the training set and the bounding box loss function LCIoU to train the improved YOLOv5 model.

[0070] Specifically, the bounding box loss function LCIoU is defined as follows:

[0071]

[0072]

[0073] In the formula, IOU represents the overlap ratio between the predicted bounding box (A in the figure) and the ground truth bounding box (B in the figure), that is, the union of their intersection and union; for example Figure 4 , ρ 2 (d) in the figure represents the Euclidean distance between the center points of the predicted bounding box and the ground truth bounding box; p, p gt, where are the center coordinates of the predicted bounding box and the ground truth bounding box, respectively; c represents the diagonal length of the minimum bounding box of the predicted bounding box and the ground truth bounding box; α represents the weight coefficient of υ, which does not need to be differentiated; υ represents the parameter that measures the consistency of aspect ratio; the specific calculation formulas for υ and α are as follows:

[0074]

[0075] In the formula, ω gt and h gt ω represents the width and height of the ground truth bounding box; h and ω represent the width and height of the predicted bounding box.

[0076]

[0077] In the YOLOv5 model parameters, set the input image size, prior bounding box size, labeled class names, initial learning rate, and learning rate adjustment strategy. Train the model using the training set; use the validation set for validation during training; and save the final weight file after the network model converges.

[0078] Compared to GIOU, the CIOU selected in this invention not only possesses the advantages of GIOU, but also takes into account the Euclidean distance between center points and the aspect ratio.

[0079] Step 4: Use the trained YOLOv5 model to detect garbage targets on the test set, output the garbage target detection results, and obtain the correct garbage category.

[0080] This invention uses a machine vision sensor and a YOLOv5 neural network to perform inference on input images and videos. It can detect over a wide range, has rich original information about the detected targets, and has a good recognition effect.

[0081] This invention achieves intelligent waste sorting; while reducing the waste of human resources and the difficulties of waste disposal, the improved YOLOv5 model increases the accuracy of waste detection, ultimately reaching an accuracy of 93%. The training results are as follows... Figure 5 As shown.

[0082] Example 2

[0083] Specifically, based on Example 1, the solution will be explained in conjunction with specific implementation examples to further demonstrate the technical effects of this solution. Step 3, the training process, is as follows:

[0084] Step 3.1: After performing data augmentation on the training set described in Step 1, input the training set into the YOLOv5 model to obtain the predicted values ​​of the object detection labels;

[0085] Step 3.2: Calculate the loss function value using the true value and predicted value of the target detection label;

[0086] Step 3.3: Update the model parameters using the loss function value;

[0087] Step 3.4: Input the test set into the improved YOLOv5 model to obtain the predicted object detection labels;

[0088] Step 3.5: Calculate the bounding box loss function value and test set accuracy using the actual and predicted values ​​of the object detection labels;

[0089] Step 3.6: Determine if the test set accuracy is greater than the maximum accuracy. If it is, save the YOLOv5 model, update the test set accuracy, and proceed to the next step; otherwise, proceed directly to the next step.

[0090] Step 3.7: Determine whether the YOLOv5 model has converged. If it has, proceed to the next step; otherwise, return to step 3.1.

[0091] Step 3.8: Determine if the maximum number of training rounds has been reached. If yes, end the process and output the completed YOLOv5 model; otherwise, return to step 3.1.

[0092] Step 3.4 specifically involves: inputting the test set into the Backbone module of the improved YOLOv5 model. After feature extraction via a Focus module, four CBA modules, four CBAM attention mechanism modules, and one SPPF module, corresponding feature maps are obtained. These feature maps are then input into the Neck module. In the Neck, weighted feature fusion using BiFPN is employed to learn the importance of different input features, allowing for differentiated fusion of different input features. A line is added between the initial features and the final node in the same layer to facilitate faster fusion of the lowest-level features, resulting in more reliable fusion results. After Neck fusion, three feature maps of different sizes are obtained. These three feature maps are then processed by a CBL module and a Conv convolutional layer to obtain spatial vectors with widths and heights of 76, 38, and 19, respectively, yielding the predicted object detection labels after processing by the improved YOLOv5 model.

[0093] Example 3

[0094] The present invention also provides a garbage image classification system, comprising:

[0095] The image acquisition module is used to acquire the garbage dataset and divide the garbage dataset into training set and test set;

[0096] The model improvement module is used to improve the YOLOv5 model: it uses the CBAM attention mechanism as the feature extraction module, the BiFPN structure as the feature fusion structure, and the bounding box loss function LCIoU.

[0097] The training module is used to obtain the prior anchor boxes of the training set using the K-means clustering method, and to train the improved YOLOv5 model using the prior anchor boxes of the training set and the bounding box loss function LCIoU.

[0098] The output module is used to detect garbage targets on the test set using the trained YOLOv5 model, and output the garbage target detection results to obtain the correct garbage category.

[0099] Obviously, the above embodiments of the present invention are merely examples for clearly illustrating the present invention, and are not intended to limit the implementation of the present invention. Those skilled in the art can make other variations or modifications based on the above description. It is neither necessary nor possible to exhaustively describe all embodiments here. Any modifications, equivalent substitutions, and improvements made within the spirit and principles of the present invention should be included within the scope of protection of the claims of the present invention.

Claims

1. A method for classifying a junk image, the method comprising: Comprising the following steps: Step 1, obtaining a garbage dataset, dividing the garbage dataset into a training set and a test set; Step 2, improving YOLOv5 model: using CBAM attention mechanism as feature extraction module, using BiFPN structure as feature fusion structure, using bounding box loss function LCIoU ; The CBAM attention mechanism is used as the feature extraction module: the first CBH convolution module in the backbone module is replaced with a Focus module, which is used for slicing the input image; the original C3 module of the backbone module is replaced with a CBAM attention mechanism, which includes channel attention and spatial attention; The BiFPN structure is used as the feature fusion structure: on the basis of the original BiFPN structure, nodes with only one input edge are removed; the number of input feature maps is increased on feature maps with smaller resolution; in the upsampling process, a splicing operation is adopted, and in the downsampling process, a stacking operation is adopted; Step 2 further comprises setting the activation function of the convolution layer of the YOLOv5 model to a Hardswish function; Step 3, obtain the prior anchor frame of the training set by using the K-means clustering method, and use the prior anchor frame of the training set and the bounding box loss function LCIoU Train the improved YOLOv5 model; Step 4, using the trained YOLOv5 model to detect the test set, outputting the garbage target detection result, and obtaining the correct garbage class.

2. The method of claim 1, wherein, In step 1, the garbage dataset is obtained, and the VoTT or Imglabeling tool is used for labeling the garbage dataset; the labeled dataset is randomly divided into a training set and a test set.

3. The method of claim 2, wherein, In step 1, the data set in the garbage data set is taken out batch, and then 4 images are randomly taken out from it, which are randomly positioned and cropped to form a new image, and then the image is subjected to mosaic data enhancement processing and label smoothing processing.

4. The method of claim 3, wherein, The label smoothing processing is a binary cross-entropy label smoothing processing, and its formula is: In the formula, y is a one-hot encoding of the garbage image data label; is a smoothing factor.

5. The method of claim 1, wherein, The specific steps of step 3 are: Step 3.1: After data enhancement on the training set in step 1, input the training set into the YOLOv5 model to obtain the target detection label prediction value; Step 3.2: Calculate the loss function value using the target detection label true value and the target detection label prediction value; Step 3.3: Update the model parameters using the loss function value; Step 3.4: Input the test set into the improved YOLOv5 model to obtain the target detection label prediction value; Step 3.5: Calculate the bounding box loss function value and the test set accuracy using the target detection label true value and the target detection label prediction value; Step 3.6: Determine whether the test set accuracy is greater than the maximum accuracy, if yes, save the YOLOv5 model, update the test set accuracy, and go to the next step; otherwise, go to the next step directly; Step 3.7: Determine whether the YOLOv5 model converges, if yes, go to the next step; otherwise, return to step 3.1; Step 3.8: Determine whether the maximum training round is reached, if yes, end and output the trained YOLOv5 model; otherwise, return to step 3.

1.

6. The method of claim 5, wherein, The bounding box loss function LCIoU is defined as: In the formula, IOU represents the intersection rate of the predicted frame and the real frame, that is, the union of their intersection and union; p 2 represents the Euclidean distance between the center points of the predicted frame and the real frame; p, p gt are the center point coordinates of the predicted frame and the real frame, respectively; c represents the diagonal length of the minimum bounding box of the predicted frame and the real frame; α represents the weight coefficient of υ; represents a parameter for measuring the consistency of the aspect ratio; , α The specific calculation formula is as follows: In the formula, denotes the width, height of the real box; denotes the width, height of the predicted box; 。 7. A spam image classification system, characterized by, Comprising: An image acquisition module for acquiring a garbage dataset, dividing the garbage dataset into a training set and a test set; A model improvement module is configured to improve the YOLOv5 model: using a CBAM attention mechanism as a feature extraction module, using a BiFPN structure as a feature fusion structure, and using a bounding box loss function LCIoU ; Step 2, improving the YOLOv5 model: using CBAM attention mechanism as a feature extraction module, using BiFPN structure as a feature fusion structure, using bounding box loss function LCIoU; Using CBAM attention mechanism as a feature extraction module: replacing the first CBH convolution module in the backbone module with a Focus module, which is used for slicing the input image; replacing the original C3 module in the backbone module with a CBAM attention mechanism, which includes channel attention and spatial attention; The use of BiFPN structure as a feature fusion structure: removing nodes with only one input edge based on the original BiFPN structure; increasing the number of input feature maps on feature maps with smaller resolutions; using concatenation operation in the up-sampling process and using superposition operation in the down-sampling process; Step 2 also includes setting the activation function of the convolution layer of the YOLOv5 model to the Hardswish function; The training module is configured to obtain prior anchor boxes of the training set by using a K-means clustering method, and obtain the improved YOLOv5 model by using the prior anchor boxes of the training set, a bounding box loss function LCIoU The improved YOLOv5 model is trained. The output module is used to detect the test set using the trained YOLOv5 model, output the garbage target detection result, and get the correct garbage category.

Images