YOLOv8~垃圾识别_yolov8 垃圾检测

这里将介绍如何使用YOLOv8的分割模型实现垃圾识别，其中所使用的训练数据来自TACO垃圾数据集。

YOLOv8是Ultralytics开源的一个非常火的AI算法，目前支持目标检测、实例分割、姿态估计等任务。

本文将介绍如何使用YOLOv8的分割模型实现垃圾识别，其中所使用的训练数据来自TACO垃圾数据集。

数据集介绍

TACO是一个包含在不同环境下（室内、树林、道路和海滩）拍摄的垃圾图像数据集，这些图像中的垃圾对象被精细地用方框和多边形进行了标注，标注信息采用与COCO数据集一样的格式，总共有60个类别，不过有的类别标注得很少甚至没有。下图是TACO数据集中的一些标注示例：

如果需要下载数据集，先执行下面的命令拉取官方的GitHub仓库：

git clone https://github.com/pedropro/TACO.git  

然后用Python运行脚本即可下载数据集：

python3 download.py  

如果下载过程中被中断了，只需重新执行download脚本即可继续下载。

训练模型

转换标注格式

TACO数据集原始的标注信息被保存在一个名为annotations.json的文件中，在使用该数据集训练YOLOv8分割模型前，需要先把原始的标注信息转换为YOLOv8要求的格式。YOLOv8分割模型训练时需要的标注格式如下：

<id> <x_1> <y_1> ... <x_n> <y_n>  

一个对象的标注信息放在一行，首先是该对象类别的id（从0开始算），接着将多边形各点像素坐标的x和y值依次排列，其中x和y的值需要分别除以图像的宽度和高度进行归一化，一幅图像的所有标注信息放在一个与图像同名的txt文件中。

进行格式转换后，txt文件中的内容类似于这样：

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  

用于格式转换的关键Python代码如下：

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()  

创建配置文件

首先仿照ultralytics/cfg/datasets/coco128-seg.yaml创建一个TACO数据集的配置文件taco-seg.yaml，文件内容如下：

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  
  ...  

数据集的设置的方式有几种形式，我的方式是建立images和labels两个目录，分别用于存放图像和txt标注文件，然后把数据集按照8:1:1的比例划分训练集、验证集、测试集，再把三个数据集图片的绝对路径分别写入train.txt、val.txt和test.txt三个文件中。所以上面的taco-seg.yaml文件中设置的路径path就是train.txt、val.txt和test.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  

配置好数据集后，还要设置模型参数。首先将ultralytics/cfg/models/v8/yolov8-seg.yaml文件拷贝一份，命名为yolov8-seg-taco.yaml，然后把文件中的类别数量nc从80改为TACO数据集的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]  
  
...  

其他有关模型结构的参数如果没有必要就不需要修改了。  whaosoft aiot http://143ai.com  

训练

训练YOLOv8可以使用命令行也可以编写Python代码实现，个人觉得还是使用命令行比较方便，所以本文采用命令行的方式进行训练，调用的命令如下：

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  

这里data参数用于指定数据集配置文件，model参数用于指定模型配置文件，如果不知道有哪些参数可以参考ultralytics/cfg/default.yaml文件，这个文件里面包含所有需要的参数。需要注意的是，我这里指定的模型配置文件名为yolov8n-seg-taco.yaml，但是前面我创建的文件名为yolov8-seg-taco.yaml，这是为什么呢？因为我这里想使用的模型是yolov8n。假如我想使用yolov8x模型，那么训练的时候设置参数model=yolov8x-seg-taco.yaml就可以了。

训练的结果保存在runs/segment/taco-seg目录下，其中权重保存在该目录下的weights文件夹中。

结果

训练完成后，我们可以调用命令测试一下模型的效果：

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  

下面是我在测试集的两张图片上测试的结果：