Improved garbage classification method based on YOLOv8

Abstract

The invention provides an improved garbage classification method based on YOLOv8, and the method comprises the following steps: S1, obtaining garbage images, and constructing a data set; s2, an SPD-Conv module is introduced into the down-sampling operation of the backbone network of the YOLOv8 model, and an improved S-YOLOv8 model is obtained; s3, inputting the data set into an S-YOLOv8 model for training and evaluation, and obtaining an improved garbage classification detection model; and S4, when the trained S-YOLOv8 model is used for reasoning, SAHI is introduced for slice reasoning. According to the method, when the trained S-YOLOv8 model is used for reasoning, the SAHI is introduced for slice reasoning, when the improved garbage classification detection model is used for garbage classification recognition, the SAHI algorithm is combined for reasoning, the recognition accuracy of garbage of a small target can be improved, meanwhile, each slice is independently processed, and the recognition efficiency is improved. Under the condition that the slices can capture the target object, the detection quality is not influenced, and the detection quality is ensured.

Description

Technical Field

[0001] The present invention relates to the technical field of garbage classification, and in particular to an improved garbage classification method based on YOLOv8. Background Art

[0002] The goal of object detection is to detect the location and bounding box of the object in an image or video, and then classify or identify the object. This task requires the simultaneous completion of the object location and classification tasks. The output of object detection includes the location and category of the object. In simple terms, object detection enables the computer to find objects in an image and give their location and bounding box.

[0003] Object detection methods can generally be divided into two categories: machine learning-based and deep learning-based methods.

[0004] Machine learning-based object detection methods usually use traditional machine learning algorithms, such as support vector machines, AdaBoost, and random forests. The basic idea of ​​these methods is to extract image features and classify the features using a classifier, and then use an object detector to detect the object. These algorithms require manual selection and extraction of image features, and therefore require the knowledge and experience of domain experts.

[0005] Object detection methods based on deep learning usually use deep neural networks to automatically learn features and perform object detection. Currently, the more popular deep learning object detection methods include two categories: region extraction-based methods (two-stage detection methods) and single-stage detection methods. Among them, region extraction-based methods include R-CNN, Fast R-CNN, Faster R-CNN, and Mask R-CNN, etc. They mainly generate target candidate regions through candidate region extractors, and use CNN networks to extract features and classify each candidate region. Single-stage detection methods directly extract target location and category information from images, such as YOLO and SSD, which can achieve faster detection speeds.

[0006] Compared with the two-stage algorithm, the one-stage algorithm such as YOLO has a simple model structure, fast detection speed and sufficient accuracy, and has been widely used in the industry.

[0007] YOLO: "You Only Look Once" is an algorithm that uses convolutional neural networks for object detection. YOLO is one of the faster object detection algorithms. Although it is not the most accurate object detection algorithm, it is a good choice when real-time detection is required and accuracy does not need to be too high. However, studies have found that when the target object is small, the use of the YOLOv8 model for garbage classification has a low recognition rate. Summary of the invention

[0008] In order to solve the problems existing in the background technology, the present invention proposes an improved garbage classification method based on YOLOv8, which improves the recognition rate of small target objects in the garbage classification and recognition process.

[0009] The present invention provides the following technical solutions:

[0010] An improved garbage classification method based on YOLOv8 includes the following steps:

[0011] S1, obtain garbage images and construct a dataset, randomly dividing the dataset into training set, validation set and test set in a ratio of 3:1:1;

[0012] S2, introduces the SPD-Conv module in the downsampling operation of the YOLOv8 model backbone network to obtain the improved S-YOLOv8 model;

[0013] S3, input the data set into the S-YOLOv8 model for training and evaluation to obtain an improved garbage classification detection model;

[0014] S4, when using the trained S-YOLOv8 model for reasoning, SAHI is introduced for slice reasoning.

[0015] Further, the introduction of the SAHI comprises the following steps:

[0016] B1, install the SAHI library and import the corresponding modules;

[0017] B2, load the trained garbage classification detection model and instantiate the garbage classification detection model by defining model weights, paths and other parameters.

[0018] Furthermore, the SAHI reasoning method comprises the following steps:

[0019] C1, use the get_prediction function provided by SAHI to detect the image that needs to be predicted;

[0020] C2, performs slice reasoning by specifying the slice size and overlap ratio;

[0021] C3, convert the result into the format required by the garbage classification detection model through the PredictionResult object provided by SAHI;

[0022] C4, echoes the inference result on the predicted image.

[0023] Furthermore, the datasets include dataset A constructed using the open source dataset TACO, dataset B constructed using the open source dataset AquaTras, dataset C constructed using the open source dataset HGI-30, dataset D constructed using the open source dataset TrashNet-30, and dataset E constructed by collecting garbage images from life and annotating them using LabelImg.

[0024] Furthermore, the construction of the dataset E includes the following steps:

[0025] A1, labeling the garbage image through labeling software, including labeling the coordinate position information of the garbage object in the garbage image and the target name;

[0026] A2, obtain the annotation information and save it in the format required by the S-YOLOv8 model to obtain the data set E.

[0027] Furthermore, the SPD-Conv module includes a space-to-depth layer and a non-strided convolutional layer; the space-to-depth layer reduces the spatial dimension of the input feature map to the channel dimension while retaining all information in the channel dimension.

[0028] Compared with the prior art, the present invention has the following beneficial effects:

[0029] 1. The present invention discloses an improved garbage classification method based on YOLOv8. When using the trained S-YOLOv8 model for reasoning, SAHI is introduced for slice reasoning. When using the improved garbage classification detection model for garbage classification and recognition work, the SAHI algorithm is combined for reasoning, which can improve the recognition accuracy of small-target garbage. At the same time, each slice is processed independently. When the slice can capture the target object, the detection quality is not affected, thereby ensuring the detection quality.

[0030] 2. The present invention discloses an improved garbage classification method based on YOLOv8, which introduces SAHI for slice reasoning. It has the advantages of fast small image slice processing and low memory consumption, which is conducive to smooth operation on low-end hardware, reduces the computational burden, and can be easily extended to images of different sizes and resolutions, thereby enhancing scalability. BRIEF DESCRIPTION OF THE DRAWINGS

[0031] In order to more clearly illustrate the technical solutions of the embodiments of the present invention, the drawings required for use in the embodiments will be briefly introduced below. It should be understood that the following drawings only show certain embodiments of the present invention and therefore should not be regarded as limiting the scope. For ordinary technicians in this field, other related drawings can be obtained based on these drawings without paying creative work.

[0032] Figure 1 The figure is a flow chart of the method of the present invention.

[0033] Figure 2 It is a schematic diagram of the SAHI slice reasoning process of the present invention.

[0034] Figure 3 Schematic diagram of the image processing flow of the SPD-Conv module of the present invention. DETAILED DESCRIPTION

[0035] In order to make the purpose, technical solutions and advantages of the embodiments of the present invention clearer, the technical solutions in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. 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.

[0036] Therefore, the following detailed description of the embodiments of the present invention provided in the accompanying drawings is not intended to limit the scope of the invention claimed for protection, but merely represents selected embodiments of the present invention. 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.

[0037] An improved garbage classification method based on YOLOv8, reference Figure 1 , including the following steps:

[0038] Step 1: Obtain garbage images and construct a dataset, which is randomly divided into a training set, a validation set, and a test set in a ratio of 3:1:1;

[0039] Step 2, introduce the SPD-Conv module into the downsampling operation of the YOLOv8 model backbone network to obtain the improved S-YOLOv8 model;

[0040] Step 3: Input the data set into the S-YOLOv8 model for training and evaluation to obtain an improved garbage classification detection model;

[0041] Step 4: When using the trained S-YOLOv8 model for reasoning, introduce SAHI for slice reasoning.

[0042] Furthermore, the data set includes data set A constructed using the open source data set TACO, data set B constructed using the open source data set AquaTras, data set C constructed using the open source data set HGI-30, data set D constructed using the open source data set TrashNet-30, and data set E constructed by collecting garbage images in life and annotating them using LabelImg. Data set AE is input into the S-YOLOv8 model for training and evaluation in step 3 to obtain a garbage classification detection model; and is input into the garbage classification detection model for testing in step 4 to obtain a garbage image detection map.

[0043] Furthermore, the dataset E annotates the image by using Labeling software, obtains the annotation information and saves it in YOLO format to obtain the dataset, wherein the annotation includes the coordinate position information of the garbage object in the image and the name of the target;

[0044] The data set is randomly divided according to a proportion to obtain the training set, the validation set and the test set.

[0045] In step 2, the SPD-Conv module is introduced into the downsampling operation of the original YOLOv8 model backbone network, where the SPD-Conv module consists of a space-to-depth (SPD) layer followed by a non-strided convolution (Conv) layer; the SPD layer reduces the spatial dimension of the input feature map to the channel dimension while retaining all the information in the channel dimension, achieving downsampling of the spatial dimension without losing information in the channel. Compared with traditional CNN convolution, the SPD-Conv module has better performance in small target detection and low-resolution images. It avoids the fixed stride and pooling operations in traditional convolution, allowing the network to process low-resolution images and small objects more effectively;

[0046] The SPD-Conv module processing steps include the following steps:

[0047] F1, for the original feature map X of size S×S×C1, cut out a series of sub-feature maps;

[0048] F2, each sub-image downsamples the original feature map X by a certain scaling factor;

[0049] F3, connects these sub-feature maps along the channel dimension to obtain an intermediate feature map X';

[0050] F4, after the SPD feature transformation layer, adds a non-strided convolution layer with C_2 filters to further transform X' into X".

[0051] Specifically, for the original feature map X with size S×S×C1, a series of sub-feature maps will be cut out as follows:

[0052]

[0053] Generally speaking, for a given original feature map X, the sub-map (f_{x,y}) is formed by all entries X(i + y) divisible by (i + x), so that each sub-map downsamples X by a certain scaling factor. For example, when scale = 2), 4 sub-maps (f_{0,0}), (f_{1,0}), (f_{0,1}), (f_{1,1}) can be obtained, and their shapes are all (S / 2, S / 2, C_1), and X is downsampled by 2 times. Refer to Figure 3 .

[0054] Next, these sub-feature maps are concatenated along the channel dimension to obtain a feature map X', whose spatial dimension is reduced by a scaling factor and the channel dimension is increased by a scaling factor of 2. In other words, SPD transforms the feature map X(S, S, C_1) into the intermediate feature map X'(S / scale, S / scale, scale^{2}C_1).

[0055] After the SPD feature transformation layer, a non-strided (i.e., stride = 1) convolutional layer with C_2 filters is added, where C_2 < scale^{2}C_1, and X'(S / scale, S / scale, scale^{2}C_1) is further transformed into X”(S / scale, S / scale, C_2).

[0056] Furthermore, the introduction of the SAHI includes the following steps:

[0057] B1, Install the SAHI library and import the corresponding modules;

[0058] B2, Load the trained garbage classification detection model and instantiate the garbage classification detection model by defining the model weights, paths, and other parameters.

[0059] Furthermore, when using the trained S-YOLOv8 model for inference, SAHI is introduced for sliced inference.

[0060] During the inference process, to solve the problem of low recognition rate caused by small target objects, the present invention expands the dataset by extracting image patches from the image and fine-tunes the dataset. As Figure 2 shown, each image is cut into overlapping image patches with hyperparameters M and N sizes within the specified ranges of [Mmin, Mmax] and [Nmin, Nmax] Then during fine-tuning, the image patches are resized by maintaining the aspect ratio so that the image width is between 800 and 1333 pixels to obtain

[0061] ˊˊˊ

[0062] The enhanced images I1, I2, ..., I k , so that the relative object size is larger compared to the original image. Together with

[0063] Original Image (to facilitate the detection of large objects), Iˊ1, Iˊ2, ..., Iˊ k These images are used together during fine-tuning.

[0064] Slice-Assisted Hyper-Inference (SAHI): The slicing method is also used in the reasoning step, such as Figure 2 First, slice the original demand image I into l overlapping slices of size M×N Each slice image is then resized while maintaining the aspect ratio. Afterwards, the object detection forward pass is applied to each overlapping patch independently. An optional full-inference (FI) using the original image can be applied to detect larger objects. Finally, the overlapping predictions and FI results are merged back to the original size using NMS. During NMS, boxes with a higher intersection over union (IoU) than a predefined matching threshold Tm are matched, and for each match, detections with a lower detection probability than Td are removed.

[0065] The SAHI reasoning method comprises the following steps:

[0066] C1, use the get_prediction function provided by SAHI to detect the image that needs to be predicted;

[0067] C2, performs slice reasoning by specifying the slice size and overlap ratio;

[0068] C3, convert the result into the format required by the garbage classification detection model through the PredictionResult object provided by SAHI;

[0069] C4, echoes the inference result on the predicted image.

[0070] Change the original YOLOv8 model's reasoning method, replace the original model's standard prediction method with prediction using the SAHI library, and perform slice reasoning by specifying the slice size and overlap ratio. After reasoning is completed, the prediction object provided by SAHI is converted to the format required by the YOLO model, and the prediction result is echoed on the predicted image.

[0071] Specific steps: Install the SAHI library, import the corresponding module, load the trained garbage classification detection YOLO model, instantiate the detection model by defining the model weights, paths and other parameters, and use the get_prediction function provided by SAHI to detect the image that needs to be predicted. After the detection is completed, the PredictionResult object provided by SAHI is used to convert the result to the format required by the YOLO model, and then the detection result is drawn on the predicted image.

[0072] SAHI: "Slice-Assisted Super Inference". It is an innovative library designed to optimize object detection algorithms for large-scale and high-resolution images. Its core function is to split an image into manageable slices, perform object detection on each slice separately, and then stitch the results together. It is compatible with a variety of object detection models such as the YOLO series, optimizing the use of computing resources while ensuring flexibility.

[0073] Slice inference: Slice inference is to slice a large or high-resolution image into smaller slices, then perform object detection on these slices, and finally reassemble the slices to restore the object location on the original image. This technology is very valuable when computing resources are limited or processing high-resolution images is prone to memory issues.

[0074] By introducing SAHI for slice reasoning when using the trained S-YOLOv8 model for reasoning, and combining the SAHI algorithm for reasoning when using the improved garbage classification detection model for garbage classification and recognition, the recognition accuracy of small-target garbage can be improved. At the same time, each slice is processed independently. When the slice can capture the target object, the detection quality is not affected, thereby ensuring the detection quality.

[0075] The introduction of SAHI for slice inference has the advantages of fast processing of small image slices and low memory consumption, which is conducive to smooth operation on low-end hardware, reduces the computational burden, and can be easily extended to images of different sizes and resolutions, enhancing scalability.

[0076] The embodiments described above are only descriptions of the preferred embodiments of the present invention and are not intended to limit the scope of the present invention. Without departing from the design spirit of the present invention, various modifications and improvements made to the technical solutions of the present invention by ordinary technicians in this field should all fall within the protection scope determined by the claims of the present invention.

Claims

1. An improved garbage classification method based on YOLOv8, characterized in that: The following steps are involved: S1, obtain garbage images and construct a dataset, randomly dividing the dataset into training set, validation set and test set in a ratio of 3:1:1; S2, introduces the SPD-Conv module in the downsampling operation of the YOLOv8 model backbone network to obtain the improved S-YOLOv8 model; S3, input the data set into the S-YOLOv8 model for training and evaluation to obtain an improved garbage classification detection model; S4, when using the trained S-YOLOv8 model for reasoning, SAHI is introduced for slice reasoning.

2. An improved YOLOv8-based garbage classification method according to claim 1, characterized in that: The datasets include dataset A constructed using the open source dataset TACO, dataset B constructed using the open source dataset AquaTras, dataset C constructed using the open source dataset HGI-30, dataset D constructed using the open source dataset TrashNet-30, and dataset E constructed by collecting garbage images in daily life and annotating them using LabelImg.

3. An improved YOLOv8-based garbage classification method according to claim 2, characterized in that: The construction of the dataset E includes the following steps: A1, labeling the garbage image through labeling software, including labeling the coordinate position information of the garbage object in the garbage image and the target name; A2, obtain the annotation information and save it in the format required by the S-YOLOv8 model to obtain the data set E.

4. The improved YOLOv8-based garbage classification method according to claim 1, characterized in that: The SPD-Conv module includes a space-to-depth layer and a non-strided convolutional layer; the space-to-depth layer reduces the spatial dimension of the input feature map to the channel dimension while retaining all information in the channel dimension.

5. The improved YOLOv8-based garbage classification method according to claim 1, characterized in that: The introduction of SAHI comprises the following steps: B1, install the SAHI library and import the corresponding modules; B2, load the trained garbage classification detection model and instantiate the garbage classification detection model by defining model weights, paths and other parameters.

6. The improved YOLOv8-based garbage classification method according to claim 5, characterized in that: The SAHI reasoning method comprises the following steps: C1, use the get_prediction function provided by SAHI to detect the image that needs to be predicted; C2, performs slice reasoning by specifying the slice size and overlap ratio; C3, convert the result into the format required by the garbage classification detection model through the PredictionResult object provided by SAHI; C4, echoes the inference result on the predicted image.

Images