训练步骤
STEP1:下载官方YOLOv5的代码并配置环境
git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txtSTEP2:准备VOC格式的数据集
- 数据放置格式
├──train_data_VOC
├── Annotations
├──放置xml文件
├── JPEGImages
├──防止img文件STEP3:将数据集转为YOLOv5所需要的COCO格式
mkdir train_data_COCO
vim VOC2COCO.pyimport 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.pySTEP4:在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 namesSTEP5:下载预训练模型
mkdir weights
cd weights
wget https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.ptSTEP6:开始训练
python train.py --data data/data.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --batch-size 64 --epochs 60
评论 (0)