目录
一、代码
直接用官方的代码
二、数据集转换(VOC 转 data 转 yolo)
1、标注生成原始数据集
经过rolabelimg标注后的数据集格式。可以直接将Annotations和JPEGImages两个文件夹放在新建的datasets下面。也可以和我一样将detect和obb做个分类。
2、voc转data
在项目根目录下创建一个roxml_to_dota.py文件,用于转换成dota能用的obb的txt文件。
# 文件名称 :roxml_to_dota.py
# 功能描述 :把rolabelimg标注的xml文件转换成dota能识别的xml文件,
# 再转换成dota格式的txt文件
# 把旋转框 cx,cy,w,h,angle,转换成四点坐标x1,y1,x2,y2,x3,y3,x4,y4
import os
import xml.etree.ElementTree as ET
import math
def edit_xml(xml_file, dotaxml_file):
"""
修改xml文件
:param xml_file:xml文件的路径
:return:
"""
tree = ET.parse(xml_file)
objs = tree.findall('object')
for ix, obj in enumerate(objs):
x0 = ET.Element("x0") # 创建节点
y0 = ET.Element("y0")
x1 = ET.Element("x1")
y1 = ET.Element("y1")
x2 = ET.Element("x2")
y2 = ET.Element("y2")
x3 = ET.Element("x3")
y3 = ET.Element("y3")
# obj_type = obj.find('bndbox')
# type = obj_type.text
# print(xml_file)
if (obj.find('robndbox') == None):
obj_bnd = obj.find('bndbox')
obj_xmin = obj_bnd.find('xmin')
obj_ymin = obj_bnd.find('ymin')
obj_xmax = obj_bnd.find('xmax')
obj_ymax = obj_bnd.find('ymax')
xmin = float(obj_xmin.text)
ymin = float(obj_ymin.text)
xmax = float(obj_xmax.text)
ymax = float(obj_ymax.text)
obj_bnd.remove(obj_xmin) # 删除节点
obj_bnd.remove(obj_ymin)
obj_bnd.remove(obj_xmax)
obj_bnd.remove(obj_ymax)
x0.text = str(xmin)
y0.text = str(ymax)
x1.text = str(xmax)
y1.text = str(ymax)
x2.text = str(xmax)
y2.text = str(ymin)
x3.text = str(xmin)
y3.text = str(ymin)
else:
obj_bnd = obj.find('robndbox')
obj_bnd.tag = 'bndbox' # 修改节点名
obj_cx = obj_bnd.find('cx')
obj_cy = obj_bnd.find('cy')
obj_w = obj_bnd.find('w')
obj_h = obj_bnd.find('h')
obj_angle = obj_bnd.find('angle')
cx = float(obj_cx.text)
cy = float(obj_cy.text)
w = float(obj_w.text)
h = float(obj_h.text)
angle = float(obj_angle.text)
obj_bnd.remove(obj_cx) # 删除节点
obj_bnd.remove(obj_cy)
obj_bnd.remove(obj_w)
obj_bnd.remove(obj_h)
obj_bnd.remove(obj_angle)
x0.text, y0.text = rotatePoint(cx, cy, cx - w / 2, cy - h / 2, -angle)
x1.text, y1.text = rotatePoint(cx, cy, cx + w / 2, cy - h / 2, -angle)
x2.text, y2.text = rotatePoint(cx, cy, cx + w / 2, cy + h / 2, -angle)
x3.text, y3.text = rotatePoint(cx, cy, cx - w / 2, cy + h / 2, -angle)
# obj.remove(obj_type) # 删除节点
obj_bnd.append(x0) # 新增节点
obj_bnd.append(y0)
obj_bnd.append(x1)
obj_bnd.append(y1)
obj_bnd.append(x2)
obj_bnd.append(y2)
obj_bnd.append(x3)
obj_bnd.append(y3)
tree.write(dotaxml_file, method='xml', encoding='utf-8') # 更新xml文件
# 转换成四点坐标
def rotatePoint(xc, yc, xp, yp, theta):
xoff = xp - xc
yoff = yp - yc
cosTheta = math.cos(theta)
sinTheta = math.sin(theta)
pResx = cosTheta * xoff + sinTheta * yoff
pResy = - sinTheta * xoff + cosTheta * yoff
return str(int(xc + pResx)), str(int(yc + pResy))
def totxt(xml_path, out_path):
# 想要生成的txt文件保存的路径,这里可以自己修改
files = os.listdir(xml_path)
for file in files:
tree = ET.parse(xml_path + os.sep + file)
root = tree.getroot()
name = file.strip('.xml')
output = out_path + name + '.txt'
file = open(output, 'w')
objs = tree.findall('object')
for obj in objs:
cls = obj.find('name').text
box = obj.find('bndbox')
x0 = int(float(box.find('x0').text))
y0 = int(float(box.find('y0').text))
x1 = int(float(box.find('x1').text))
y1 = int(float(box.find('y1').text))
x2 = int(float(box.find('x2').text))
y2 = int(float(box.find('y2').text))
x3 = int(float(box.find('x3').text))
y3 = int(float(box.find('y3').text))
file.write("{} {} {} {} {} {} {} {} {} 0\n".format(x0, y0, x1, y1, x2, y2, x3, y3, cls))
file.close()
print(output)
if __name__ == '__main__':
# -----**** 第一步:把xml文件统一转换成旋转框的xml文件 ****-----
roxml_path = "./datasets/obb/Annotations/" # 目录下保存的是需要转换的xml文件
dotaxml_path = './datasets/obb/dotaxml/'
out_path = './datasets/obb/labelTxt/'
filelist = os.listdir(roxml_path)
for path in (dotaxml_path, out_path):
if not os.path.exists(path):
os.makedirs(path)
for file in filelist:
edit_xml(os.path.join(roxml_path, file), os.path.join(dotaxml_path, file))
# -----**** 第二步:把旋转框xml文件转换成txt格式 ****-----
totxt(dotaxml_path, out_path)
运行完后生成两个文件夹
3、划分数据集
在项目根目录下新建一个splitdata.py文件,
# 将图片和标注数据按比例切分为 训练集和测试集
import shutil
import random
import os
import argparse
# 检查文件夹是否存在
def mkdir(path):
if not os.path.exists(path):
os.makedirs(path)
def main(image_dir, txt_dir, save_dir):
# 创建文件夹
mkdir(save_dir)
images_dir = os.path.join(save_dir, 'images')
labels_dir = os.path.join(save_dir, 'labels')
img_train_path = os.path.join(images_dir, 'train')
img_val_path = os.path.join(images_dir, 'val')
label_train_path = os.path.join(labels_dir, 'train_original')
label_val_path = os.path.join(labels_dir, 'val_original')
mkdir(images_dir)
mkdir(labels_dir)
mkdir(img_train_path)
mkdir(img_val_path)
mkdir(label_train_path)
mkdir(label_val_path)
# 数据集划分比例,训练集75%,验证集15%,测试集15%,按需修改
train_percent = 0.50
val_percent = 0.50
total_txt = os.listdir(txt_dir)
num_txt = len(total_txt)
list_all_txt = range(num_txt) # 范围 range(0, num)
num_train = int(num_txt * train_percent)
num_val = int(num_txt * val_percent)
train = random.sample(list_all_txt, num_train)
# 在全部数据集中取出train
val = [i for i in list_all_txt if not i in train]
# 再从val_test取出num_val个元素,val_test剩下的元素就是test
# val = random.sample(list_all_txt, num_val)
print("训练集数目:{}, 验证集数目:{}".format(len(train), len(val)))
for i in list_all_txt:
name = total_txt[i][:-4]
srcImage = os.path.join(image_dir, name + '.jpg')
srcLabel = os.path.join(txt_dir, name + '.txt')
if i in train:
dst_train_Image = os.path.join(img_train_path, name + '.jpg')
dst_train_Label = os.path.join(label_train_path, name + '.txt')
shutil.copyfile(srcImage, dst_train_Image)
shutil.copyfile(srcLabel, dst_train_Label)
elif i in val:
dst_val_Image = os.path.join(img_val_path, name + '.jpg')
dst_val_Label = os.path.join(label_val_path, name + '.txt')
shutil.copyfile(srcImage, dst_val_Image)
shutil.copyfile(srcLabel, dst_val_Label)
if __name__ == '__main__':
"""
python split_datasets.py --image-dir my_datasets/color_rings/imgs --txt-dir my_datasets/color_rings/txts --save-dir my_datasets/color_rings/train_data
"""
parser = argparse.ArgumentParser(description='split datasets to train,val,test params')
parser.add_argument('--image-dir', type=str,default='datasets/obb/JPEGImages', help='image path dir')
parser.add_argument('--txt-dir', type=str,default='datasets/obb/labelTxt' , help='txt path dir')
parser.add_argument('--save-dir', default='datasets/obb/meter', type=str, help='save dir')
args = parser.parse_args()
image_dir = args.image_dir
txt_dir = args.txt_dir
save_dir = args.save_dir
main(image_dir, txt_dir, save_dir)
运行完后,生成一个文件夹(meter),里面是划分好的数据集。
4、将labels的data格式转换成yolo格式
在项目根目录下新建一个data_to_yolo.py文件夹。
import sys
from ultralytics.data.converter import convert_dota_to_yolo_obb
# sys.path.append(r'/ultralytics')
convert_dota_to_yolo_obb(r'datasets/obb/meter')
注意修改路径!!
!!!!运行前需修改两个位置!!!!
ctrl+左键定位 这个函数 convert_dota_to_yolo_obb 直接可以定位到 converter.py 文件中
**将class_mapping修改成自己的类名**
**修改图片后缀**
然后运行data_to_yolo.py,运行完后在labels文件夹下会生成train和val两个文件夹。
至此,数据集(meter)完成,下一步可以用这个数据集训练
三、训练
3.1 在项目根目录下新建一个obb.yaml文件,复制以下内容。注意修改路径
names:
0: label
path: E:\python\yolov8\ultralytics_train_7_30\datasets/obb/meter
train: images/train
val: images/val
3.2 在ultralytics/models/yolo/obb下新建一个obb_train_my.py文件。复制以下内容
(需要提前下载yolov8n.pt 和yolov8n-obb.pt两个模型在这个目录下)前往官网下载
from ultralytics import YOLO
# Load a model
model = YOLO("yolov8n-obb.pt") # load a pretrained model (recommended for training)
# Train the model
results = model.train(task='obb',data="obb.yaml", model='yolov8n-obb.pt', epochs=300, imgsz=1024, batch=16, workers=0,
resume=False)
3.3 运行obb_train_my.py,开始训练
根据自己的情况修改参数,或者前往ultralytics/cfg/default.yaml,这个路径下修改这个yaml文件里面的参数,都是一样的
注意,参数设置优先是results = model.train(task='detect',data="datasets/detect/dataset.yaml", model='yolov8n.pt', epochs=300, imgsz=640, batch=16, workers=0, resume=False)这个里面的。
如果采用default.yaml,则需要将上面代码中的参数删除
eg:results = model.train()
3.2
训练后结果会保存在ultralytics/models/y/obb/runs下面