python提取COCO数据集中特定的类

import os
from shutil import copyfile,move

src_path = './labels/train2017/'  # 标签
img_path = './images/train2017/'  # 图像
dst_label_path = './traffic_train_coco/labels/'
dst_img_path = './traffic_train_coco/images/'
cls_id = ['79']  # 牙刷

labels = os.listdir(src_path)
labels.sort()

for label in labels:
    if label[-1] != 't':
        continue
    # 存标签
    tmp = 0
    for line in open(src_path + '/' + label):
        str_list = line.split()
        # 被选类别的标签
        for i in range(len(cls_id)):
            if str_list[0] == cls_id[i]:
                # 改成自己的标签，这里是数组下标
                tmp =1
                break
    # 没有被选类别
    if tmp == 0:
        continue
    # 新的标签文件
    move(src_path + '/' + label, dst_label_path + '/' + label)
    # copyfile(src_path + '/' + label, dst_label_path + '/' + label)

    image = label[:-4] + '.jpg'
    # 拷贝有被选类别的图片
    move(img_path + '/' + image, dst_img_path + image)
    # copyfile(img_path + '/' + image, dst_img_path + image)

参考：链接