Yolov5 lite 重新训练COCO数据(只训练COCO含Person部分,其他当负样本)

一、数据准备

  1. 本文使用的YoLov5 lite 对COCO进行重新训练。这里需要提出的出,YOLO针对训练样本的数据划分,是不能直接采我们下载的COCO数据划分方式的的,需要将COCO中的json文件划分成txt文件(这里目标检测,就针对instances)。txt中存放的是每张图片中含有目标对象的类别和anchor的坐标。YOLO的数据集的划分目录结构如下:
    images:
    train2017:
    xxxxx.jpg
    xxxxx.jpg

    val2017:
    xxxxx.jpg
    xxxxx.jpg

    labels:
    train2017:
    xxxxx.txt
    xxxxx.txt

    val2017:
    xxxxx.txt
    xxxxx.txt

一、 修改文件

1.我们把数据集划分好了以后,并且放在了你一般放数据集的地方后,我们需要去修改一些配置文件,让我们训练的时候可以找到这些数据集。我们下载好yolo项目以后,我们找到data文件夹下的coco.yaml。我们修改里面的参数

train: ../../data/object/COCO/Yolo/images/train2017  # 118287 images
val: ../../data/object/COCO/Yolo/images/val2017  # 5000 images

这里修改的是我们数据集的路径,需要注意的是这里填的是图像的路径,label训练代码会自己找到。

# number of classes
nc: 1

这个是类别的数量。我这里是要Person这个类,所有类别为1

# class names
names: [ 'person' ]

类别名称

2. 修改完数据部分以后,我们需要去修改models文件夹下的文件,这个文件夹下有几个版本权重的配置文件,我们针对你使用的权重版本,我这里用的是v5Lite-s.pt,所以我修改v5Lite-s.yaml这个配置文件中的 nc 参数,我只有一个类,所以修改为1。

nc: 1  # number of classes

修改完配置文件以后,因为我们只要针对Person进行识别,我们只需要Personlabel信息,其他类别的信息就不要了,直接为负样本。针对这里,我们需要去修改utils文件夹下的dataset的代码。

a. 因为YOLO会对数据做一个cache为方便读取数据,那我们在cache阶段就将我们Person单独取出来。我们找到cache_labels这个函数。

 # verify labels
        # 这里就是读trian/label的地方
        if os.path.isfile(lb_file):
            nf += 1  # label found
            with open(lb_file, 'r') as f:
                person_list = []
                l = [x.split() for x in f.read().strip().splitlines()]

                # 只将含有person的数据放入cache
                for ll in l:
                    if ll[0] == '1':
                        person_list.append(ll)
                l = person_list

我们新建一个person_list 用来存person的信息,我们读label的时候,是针对一个文件一行一行的读取(因为一个txt中有多个类型的信息),每一行的数据构成是 【类别id,anchor的坐标信息一共四个值】 。例如: 【58 0.389578 0.416103 0.038594 0.163146】 ,因为我需要的是person,person的类别id是1,所以我只保留为1的数据,这里的代码就非常的简单了。

这里修改好以后,就可以按照正常的要求进行训练啦!

你可能感兴趣的:(深度学习,目标检测,计算机视觉)