将cub数据集转换为yolov5所需格式

莫言过

已于 2023-04-27 10:31:11 修改

阅读量1.6k

点赞数 5

文章标签： YOLO

于 2023-04-26 11:03:21 首次发布

本文链接：https://blog.youkuaiyun.com/qq_43979221/article/details/130360949

版权

文章介绍了CUB鸟类数据集的结构，并详细阐述了如何将其转换为YOLOv5所需的格式，包括处理yaml配置文件，创建标签文件以及将数据划分为训练集和验证集。接着，作者提到了使用YOLOv5进行模型训练的情况，以及模型检测的结果和准确性评估。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

cub数据集简介：

cub数据集（Caltech-UCSD Birds-200-2011）由加州理工大学提供（https://www.vision.caltech.edu/datasets/cub_200_2011/），共包含200种鸟类，共11788张图片，下载后包含如下文件（具体信息见README）：
在这里插入图片描述

images文件夹：文件夹内有200个子文件夹（每种鸟类一个文件夹），每个文件夹中包含约60张同种鸟类图片；
在这里插入图片描述

classes.txt：每行为种类序号及其对应种类，共200条信息；
在这里插入图片描述

images.txt：每行为图片序号及其images文件夹下路径及图片名字，共11788条信息；
在这里插入图片描述

bounding_boxes.txt：每行为图片序号及bounding_box信息（x、y、weith、height）（x，y为框左上角坐标），共11788条(目测标注的不是很准，待验证)；
在这里插入图片描述

image_class_labels.txt：每行为图片序号及其所属类别，共11788条；
在这里插入图片描述

train_test_split.txt：每行为图片序号及其是否属于训练集，1代表是，0代表否；
在这里插入图片描述

yolov5所需格式介绍

首先需要一个yaml文件（应该是存放在 ./yolov5/data下），描述数据路径及种类序号（详情见https://docs.ultralytics.com/yolov5/tutorials/train_custom_data/#train-on-custom-data）
在这里插入图片描述
然后每张图片要有一个描述文档，每行为种类序号（从0开始）及bounding_box信息（需转换至0-1的范围）

存放目录：datasets和yolov5同级，images和labels下面分别有train和val文件夹（在yaml文件中写明路径）

数据处理

首先处理yaml文件：

import shutil
import cv2
path='./CUBdata'
labpath='./datasets/cub/labels'
imapath='./datasets/cub/images'
def get_yaml():
    #get cub.yaml
    with open(path+'/classes.txt', 'r') as fo,open('cub.yaml', "a") as fi:  #r只读，w写，a追加写
        for num, line in enumerate(fo):		#enumerate为枚举，num为从0开始的序号，line为每一行的信息
            s='  '+str(num)+': '+line.split(" ")[-1]  #以空格分隔，去掉末尾的换行
            fi.write(s)						#追加写入目标文件

处理后得到目标文件（前一部分为复制coco128的yaml文件）：
在这里插入图片描述
然后处理label及图片：

def get_alllab():
    dataall={}     # 字典用于存放txt中的各种信息
    with open(path + '/images.txt', 'r') as imagesall,open(path + '/image_class_labels.txt', 'r') as classall,\
        open(path + '/train_test_split.txt', 'r') as splitall,open(path + '/bounding_boxes.txt', 'r') as boxall:
        for num, line in enumerate(imagesall):   # 值用列表存储，方便后续添加元素，-1去掉末尾的/n
            s=line.split(" ")
            dataall[s[0]]=[s[1][:-1]]
        for num, line in enumerate(classall):
            s=line.split(" ")
            dataall[s[0]].append(s[1][:-1])
        for num, line in enumerate(splitall):
            s=line.split(" ")
            dataall[s[0]].append(s[1][:-1])
        for num, line in enumerate(boxall):
            s=line.split(" ")
            dataall[s[0]].extend([s[1],s[2],s[3],s[4][:-1]])
    print('dataall have got...')
    for item in dataall:
        na = item.rjust(12, '0')   #item为字典的键，左侧扩充0，改为所需名字格式（未看到明确要求）
        image = cv2.imread(path + '/images/' + dataall[item][0]) # 读取图片，使用shape获取图片宽高
        # 这两行代码为验证boundingbox信息，手动画框，图片存储至test文件夹，左上角和右下角坐标,image.shape[1] 宽度  image.shape[0] 高度
        # cv2.rectangle(image, (int(float(dataall[item][3])), int(float(dataall[item][4]))), (int(float(dataall[item][3]))+\
        #             int(float(dataall[item][5])),int(float(dataall[item][4]))+int(float(dataall[item][6]))), (0, 0, 255), 3)
        # cv2.imwrite(imapath+'/test2017/' + na + '.jpg', image)  # 带小数的str需先转为float才能转为int
        x = (float(dataall[item][3]) + float(dataall[item][5]) / 2) / image.shape[1]
        y = (float(dataall[item][4]) + float(dataall[item][6]) / 2) / image.shape[0]
        w = float(dataall[item][5]) / image.shape[1]
        h = float(dataall[item][6]) / image.shape[0]
        s = str(int(dataall[item][1]) - 1) + ' ' + str('%.6f' % x) + ' ' + str('%.6f' % y) + \
            ' ' + str('%.6f' % w) + ' ' + str('%.6f' % h) #将cub的boundingbox转换为yolov5格式

        if dataall[item][2]=='1':  #划分训练集验证集，shutil.copy（a，b）为复制图片，a为原路径，b为目标路径（带名字则自动重命名）
            with open(labpath +'/train2017/{}.txt'.format(na), 'w') as lab:
                lab.write(s)       # 写入文件 已存在就覆盖，没有就生成
            shutil.copy(path + '/images/' + dataall[item][0], imapath+'/train2017/' + na + '.jpg')

        elif dataall[item][2] == '0':
            with open(labpath + '/val2017/{}.txt'.format(na), 'w') as lab:
                lab.write(s)
            shutil.copy(path + '/images/' + dataall[item][0], imapath+'/val2017/' + na + '.jpg')

模型训练及检测

处理完数据就可以训练模型了

nohup python train.py --data cub.yaml --epochs 300 --weights yolov5n.pt  --batch-size 128 >> cunnohup.out 2>&1 &

我训练用的5n，没有指定batchsize，我看了下是默认16，总共用了两个多小时好像。然后使用图片检测效果还可以

python detect.py --weights best.pt --source test.jpg --save-txt --save-conf

使用上述指令可以将一些数据保存到txt文件（师兄看的detect.py代码发现的方法），我复制了val中的图片用作检测，结果如下：
在这里插入图片描述

运行结果会保存到 ./run/detect，本来是会显示为该种类的概率的，但是类别名太长挡住了，不知道有没有什么办法改进，（后面只能将数据保存到txt）
在这里插入图片描述
数据应该分别是类别序号（处理的时候统一-1了，所以序号为1）、四项boundingbox信息、概率。

后来我摘出来了几张图片用作测试，发现准确率一般，可能最少还是得用5s训练。

为了验证cub给的boundingbox信息，我根据他的txt手动画了框，看着感觉还是挺准确的。（画框代码在上面注释掉了，只训练数据不用管他，没有影响）

在这里插入图片描述
以上，如有不正确的地方还望不吝赐教。