Here we will introduce how to use the YOLOv8 segmentation model to implement garbage identification. The training data used comes from the TACO garbage data set.
YOLOv8It is Ultralyticsa very popular open source AIalgorithm that currently supports tasks such as target detection, instance segmentation, and pose estimation.
This article will introduce how to use YOLOv8the segmentation model to achieve garbage identification, where the training data used comes from the TACO garbage data set.
Dataset introduction
TACOIt is a data set containing garbage images taken in different environments (indoors, woods, roads and beaches). The garbage objects in these images are carefully labeled with boxes and polygons. The annotation information is in the same format as the data set COCO. , there are 60a total of categories, but some categories are rarely or even not labeled. The following figure is TACOsome examples of annotations in the dataset:
If you need to download the data set, first execute the following command to pull the official GitHubwarehouse:
git clone https://github.com/pedropro/TACO.git
Then Pythonrun the script to download the data set:
python3 download.py
If the download is interrupted, simply re-execute downloadthe script to continue the download.
Training model
Convert label format
TACOThe original annotation information of the data set is saved in a annotations.jsonfile named. Before using the data set to train the YOLOv8segmentation model, the original annotation information needs to be converted into YOLOv8the required format. YOLOv8The annotation format required for segmentation model training is as follows:
<id> <x_1> <y_1> ... <x_n> <y_n>
The annotation information of an object is placed in one line. The first is the object category id(counted from the beginning), and then the sum 0of the pixel coordinates of each point of the polygon is arranged in order. The sum of the values xneeds to be divided by the width and height of the image for normalization. ization, all the annotation information of an image is placed in a file with the same name as the image.yxytxt
After format conversion, txtthe content in the file looks like this:
5 0.5183 0.4892 0.5480 0.4840 0.4840 0.5627 0.4840 0.5724 0.4853 0.5822 0.4879 0.5900
7 0.6227 0.5211 0.6232 0.5250 0.5074 0.6154 0.5081 0.6183 0.5107 0.5068 0.6120 0.6290
The key codes used for format conversion Pythonare as follows:
img = cv2.imread(image_path)
height, width, _ = img.shape
label_writer = open(label_path, "w")
for annotation in annotations:
category_id = annotation["category_id"]
seg_labels = []
for segmentation in annotation["segmentation"]:
points = np.array(segmentation).reshape((int(len(segmentation) / 2), 2))
for point in points:
x = point[0] / width
y = point[1] / height
seg_labels.append(x)
seg_labels.append(y)
label_writer.write(str(category_id) + " " + " ".join([str(a) for a in seg_labels]) + "\n")
label_writer.close()
Create configuration file
First, ultralytics/cfg/datasets/coco128-seg.yamlcreate a TACOdata set configuration file as follows taco-seg.yaml. The file content is as follows:
path: /home/test/TACO/data #数据集所在的目录
train: train.txt # 训练集路径，相对于path目录
val: val.txt # 验证集路径，相对于path目录
test: test.txt # 测试集路径，相对于path目录，可以不写
# 类别id和名称
names:
0: Aluminium foil
1: Battery
2: Aluminium blister pack
3: Carded blister pack
4: Other plastic bottle
5: Clear plastic bottle
6: Glass bottle
7: Plastic bottle cap
8: Metal bottle cap
9: Broken glass
10: Food Can
...
There are several ways to set up the data set. My method is to create imagestwo labelsdirectories to store images and txtannotation files, and then divide the data set 8:1:1into training set, verification set, and test set according to the proportion, and then divide the three The absolute paths of the data set images are written to three files train.txt: , val.txt, and . So the path set in the test.txtabove file is the directory where these three files are located. These three files contain the absolute paths to the pictures in the corresponding data set, similar to this :taco-seg.yamlpathtrain.txtval.txttest.txt
/home/test/TACO/data/images/batch_13/000077.jpg
/home/test/TACO/data/images/batch_11/000032.jpg
/home/test/TACO/data/images/batch_15/000073.jpg
After configuring the data set, you also need to set the model parameters. First make ultralytics/cfg/models/v8/yolov8-seg.yamla copy of the file, name it yolov8-seg-taco.yaml, and then change the number of categories in the file ncfrom 80that TACOof the data set 60:
...
# Parameters
nc: 60 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n-seg.yaml' will call yolov8-seg.yaml with scale 'n'
# [depth, width, max_channels]
n: [0.33, 0.25, 1024]
s: [0.33, 0.50, 1024]
m: [0.67, 0.75, 768]
l: [1.00, 1.00, 512]
x: [1.00, 1.25, 512]
...
Other parameters related to the model structure do not need to be modified if not necessary. whaosoft aiot http://143ai.com
train
Training YOLOv8can be Pythonimplemented using the command line or by writing code. Personally, I think it is more convenient to use the command line, so this article uses the command line for training. The calling commands are as follows:
yolo task=segment mode=train data=taco-seg.yaml model=yolov8n-seg-taco.yaml epochs=100 batch=16 imgsz=640 device=0 name=taco-seg
The parameters here dataare used to specify the data set configuration file, and modelthe parameters are used to specify the model configuration file. If you don’t know what parameters there are, you can refer to ultralytics/cfg/default.yamlthe file. This file contains all the required parameters. It should be noted that the model configuration file name I specified here is named yolov8n-seg-taco.yaml, but the file name I created earlier is named yolov8-seg-taco.yaml. Why is this? Because the model I want to use here is yolov8n. If I want to use the model, I can just yolov8xset the parameters during training .model=yolov8x-seg-taco.yaml
The results of the training are saved in runs/segment/taco-segthe directory, and the weights are saved in the weightsfolder under the directory.
result
After training is completed, we can call the command to test the effect of the model:
yolo task=segment mode=predict model=runs/segment/taco-seg/weights/best.pt source=/home/test/TACO/data/images/batch_9/000096.jpg show=True
Below are the results I tested on two images in the test set:
