快速使用YOLOv5进行训练VOC格式的数据集

jupiter
2021-09-15 / 0 评论 / 811 阅读 / 正在检测是否收录...
温馨提示:
本文最后更新于2021年12月07日,已超过809天没有更新,若内容或图片失效,请留言反馈。

训练步骤

STEP1:下载官方YOLOv5的代码并配置环境

git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

STEP2:准备VOC格式的数据集

  • 数据放置格式
├──train_data_VOC
    ├── Annotations
        ├──放置xml文件
    ├── JPEGImages
        ├──防止img文件

STEP3:将数据集转为YOLOv5所需要的COCO格式

mkdir train_data_COCO
vim VOC2COCO.py
import os
import shutil
import random
import xmltodict
from progressbar import *

#================================================================================================================
# 函数定义区

# 函数-将voc xml中的object转化为对应的一条yolo数据
def get_yolo_data(obj,img_width,img_height):
    # 获取voc格式的数据信息 
    name = obj['name']
    xmin = float(obj['bndbox']['xmin'])
    xmax = float(obj['bndbox']['xmax'])
    ymin = float(obj['bndbox']['ymin'])
    ymax = float(obj['bndbox']['ymax'])

    # 计算yolo格式的数据信息
    class_idx = class_names.index(name)
    x_center,y_center = (xmin+xmax)/2,(ymin+ymax)/2
    box_width = xmax - xmin
    box_height = ymax - ymin

    yolo_data = "{} {} {} {} {}\n".format(class_idx,x_center/img_width,y_center/img_height,box_width/img_width,box_height/img_height)
    
    return yolo_data

# 函数-将xml文件转为txt文件
def convert_annotations(image_name):
    in_file = xml_file_path + image_name + '.xml'
    out_file = txt_file_path + image_name + '.txt'

    yolo_data = ""
    
    with open(in_file) as f:
        xml_str = f.read()
    
     # 转为字典
    xml_dic = xmltodict.parse(xml_str)
    
    # 获取图片的width、height
    img_width = float(xml_dic["annotation"]["size"]["width"])
    img_height = float(xml_dic["annotation"]["size"]["height"])
    
    # 获取xml文件中的object
    objects = xml_dic["annotation"]["object"]
    
    if isinstance(objects,list): # xml文件中包含多个object
        for obj in objects:
            yolo_data += get_yolo_data(obj,img_width,img_height)   
    else: # xml文件中包含1个object
        obj = objects
        yolo_data += get_yolo_data(obj,img_width,img_height)
    
    with open(out_file,'w') as f:
        f.write(yolo_data)

# 函数-创建最终用于训练的COCO格式数据集的文件夹
def create_dir():
    if not os.path.exists('train_data_COCO/images/'):
        os.makedirs('train_data_COCO/images/')
    if not os.path.exists('train_data_COCO/labels/'):
        os.makedirs('train_data_COCO/labels/')

    if not os.path.exists('train_data_COCO/images/train/'):
        os.makedirs('train_data_COCO/images/train')
    if not os.path.exists('train_data_COCO/images/val/'):
        os.makedirs('train_data_COCO/images/val/')
    if not os.path.exists('train_data_COCO/images/test/'):
        os.makedirs('train_data_COCO/images/test/')

    if not os.path.exists('train_data_COCO/labels/train/'):
        os.makedirs('train_data_COCO/labels/train/')
    if not os.path.exists('train_data_COCO/labels/val/'):
        os.makedirs('train_data_COCO/labels/val/')
    if not os.path.exists('train_data_COCO/labels/test/'):
        os.makedirs('train_data_COCO/labels/test/')

    return

#================================================================================================================
# 功能实现区

"""
STEP1:准备工作:数据准备+创建各种所需的文件夹
"""

# 对应的VOC数据集的路径参数+类别参数
xml_file_path = './train_data_VOC/Annotations/'       # 检查和自己的xml文件夹名称是否一致
images_file_path = './train_data_VOC/JPEGImages/'  # 检查和自己的图像文件夹名称是否一致
class_names = ['Person', 'BridgeVehicle', 'LuggageVehicle', 'Plane', 'RefuelVehicle', 'FoodVehicle', 'RubbishVehicle', 'WaterVehicle', 'PlatformVehicle', 'TractorVehicle']


# 创一个临时文件夹用来存放xml文件转换出来的对应的txt文件
if not os.path.exists('train_data_COCO/temp_labels/'):
    os.makedirs('train_data_COCO/temp_labels/')
txt_file_path = 'train_data_COCO/temp_labels/'

# 执行xml到txt的转换,存储到一个临时文件夹
total_xml = os.listdir(xml_file_path)
num_xml = len(total_xml)  # XML文件总数
for i in range(num_xml):
    name = total_xml[i][:-4]
    convert_annotations(name)

# 创建COCO格式的数据所需要的各种文件夹
create_dir()

# 读取所有的txt文件
total_txt = os.listdir(txt_file_path)

print("数据准备工作完成,开始进行数据分配")

"""
STEP2:数据分配:按比例对数据集进行划分
"""

# 设置数据集划分比例,训练集75%,验证集15%,测试集15%
train_percent = 0.8
val_percent = 0.15
test_percent = 0.05


# 计算train,val,test每一类的数据数量
num_txt = len(total_txt)
num_train = int(num_txt * train_percent)
num_val = int(num_txt * val_percent)
num_test = num_txt - num_train - num_val

# 根据计算出的每类的数据数量计算出进行数据分配的索引
list_all_txt = range(num_txt)  # 范围 range(0, num)
train = random.sample(list_all_txt, num_train)# train从list_all_txt取出num_train个元素
val_test = [i for i in list_all_txt if not i in train]# 所以list_all_txt列表只剩下了这些元素:val_test
val = random.sample(val_test, num_val)# 再从val_test取出num_val个元素,val_test剩下的元素就是test


# 根据采样的索引结果进行文件分配工作
print("训练集数目:{}, 验证集数目:{},测试集数目:{}".format(len(train), len(val), len(val_test) - len(val)))

#进度条功能
widgets = ['VOC2COCO: ',Percentage(), ' ', Bar('#'),' ', Timer(),' ', ETA()]
pbar = ProgressBar(widgets=widgets, maxval=num_txt).start()
count = 0

for i in list_all_txt:
    name = total_txt[i][:-4]

    srcImage = images_file_path + name + '.jpg'
    srcLabel = txt_file_path + name + '.txt'

    if i in train:
        dst_train_Image = 'train_data_COCO/images/train/' + name + '.jpg'
        dst_train_Label = 'train_data_COCO/labels/train/' + name + '.txt'
        shutil.copyfile(srcImage, dst_train_Image)
        shutil.copyfile(srcLabel, dst_train_Label)
    elif i in val:
        dst_val_Image = 'train_data_COCO/images/val/' + name + '.jpg'
        dst_val_Label = 'train_data_COCO/labels/val/' + name + '.txt'
        shutil.copyfile(srcImage, dst_val_Image)
        shutil.copyfile(srcLabel, dst_val_Label)
    else:
        dst_test_Image = 'train_data_COCO/images/test/' + name + '.jpg'
        dst_test_Label = 'train_data_COCO/labels/test/' + name + '.txt'
        shutil.copyfile(srcImage, dst_test_Image)
        shutil.copyfile(srcLabel, dst_test_Label)
    
    #更新进度条
    count += 1
    pbar.update(count)

#释放进度条
pbar.finish()
print("数据分配工作完成,开始释放临时文")

"""
STEP3:释放临时文件
"""
shutil.rmtree(txt_file_path)
print("临时文件释放完成,VOC2COCO执行结束")
python VOC2COCO.py

STEP4:在data下创建与数据对应的data.yaml文件

文件内容按照数据的数据情况填写

path: train_data_COCO # root
train: # train images (relative to 'path') 
  - images/train
val: # val images (relative to 'path') 
  - images/val
test: # test images (optional)
  - images/test

# Classes
nc: 10  # number of classes
names: ['Person', 'BridgeVehicle', 'LuggageVehicle', 'Plane', 'RefuelVehicle', 'FoodVehicle', 'RubbishVehicle', 'WaterVehicle', 'PlatformVehicle', 'TractorVehicle']  # class names

STEP5:下载预训练模型

mkdir weights
cd weights
wget https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt

STEP6:开始训练

python train.py --data data/data.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt  --batch-size 64 --epochs 60 

参考资料

  1. https://github.com/ultralytics/yolov5
0

评论 (0)

打卡
取消