一.数据集制作
1.标注软件:roLabelImg
roLabelImg是基于labelImg改进的,是用来标注为VOC格式的数据,但是在labelImg的基础上增加了能够使标注的框进行旋转的功能。
2.数据格式转换
2.1 xml转txt
# 文件名称 :roxml_to_dota.py
# 功能描述 :把rolabelimg标注的xml文件转换成dota能识别的xml文件,
# 再转换成dota格式的txt文件
# 把旋转框 cx,cy,w,h,angle,或者矩形框cx,cy,w,h,转换成四点坐标x1,y1,x2,y2,x3,y3,x4,y4
import os
import xml.etree.ElementTree as ET
import math
cls_list=['mouse']
def edit_xml(xml_file, dotaxml_file):
"""
修改xml文件
:param xml_file:xml文件的路径
:return:
"""
#dxml_file = open(xml_file,encoding='gbk')
#tree = ET.parse(dxml_file).getroot()
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 = max(float(obj_xmin.text),0)
ymin = max(float(obj_ymin.text),0)
xmax = max(float(obj_xmax.text),0)
ymax = max(float(obj_ymax.text),0)
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)
i=0
for file in files:
tree = ET.parse(xml_path + os.sep + file)
root = tree.getroot()
name = file.split('.')[0]
output = out_path +'\\'+name + '.txt'
file = open(output, 'w')
i=i+1
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))
if x0<0:
x0=0
if x1<0:
x1=0
if x2<0:
x2=0
if x3<0:
x3=0
if y0<0:
y0=0
if y1<0:
y1=0
if y2<0:
y2=0
if y3<0:
y3=0
for cls_index,cls_name in enumerate(cls_list):
if cls==cls_name:
file.write("{} {} {} {} {} {} {} {} {} {}\n".format(x0, y0, x1, y1, x2, y2, x3, y3, cls,cls_index))
file.close()
# print(output)
print(i)
if __name__ == '__main__':
# -----**** 第一步:把xml文件统一转换成旋转框的xml文件 ****-----
roxml_path = 'data_mouse_ro_1/org_xml'
dotaxml_path = 'data_mouse_ro_1/dotaxml'
out_path = 'data_mouse_ro_1/dotatxt'
filelist = os.listdir(roxml_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)
2.2 dota_to_yolo_obb
dota数据格式:937.0 913.0 921.0 912.0 923.0 874.0 940.0 875.0 small-vehicle 0
yolo_obb格式:class_index, x1, y1, x2, y2, x3, y3, x4, y4
sys.path.append('/path/to/ultralytics')
from ultralytics.data.converter import convert_dota_to_yolo_obb
convert_dota_to_yolo_obb('/home/fut/project/ultralytics-main/ultralytics/datasets_ro')
跳转到convert_dota_to_yolo_obb.py函数,对class_mapping进行修改
2.3 分割数据集
数据集文件分布格式如下:
datasets
--images
--train
--val
--labelTxt
--trian
--val
分割代码:
import os
import random
import shutil
# 设置随机数种子
random.seed(42)
# 数据集文件夹路径和输出文件夹路径
data_folder = 'data_mouse_ro_1'
img_folder = 'data_mouse_ro_1/dataset/images'
label_folder = 'data_mouse_ro_1/dataset/labels'
# 计算每个子集的大小
total_files = len(os.listdir(os.path.join(data_folder, 'img')))
train_size = int(total_files * 0.9)
test_size = int(total_files - train_size)
# 获取所有图像文件的文件名列表
image_files = os.listdir(os.path.join(data_folder, 'img'))
random.shuffle(image_files)
# 复制图像和标注文件到相应的子集文件夹中
for i, image_file in enumerate(image_files):
base_file_name = os.path.splitext(image_file)[0]
image_path = os.path.join(data_folder, 'img', image_file)
label_path = os.path.join(data_folder, 'dotatxt', base_file_name + '.txt')
if i < train_size:
print(image_path)
#print(os.path.join(img_folder, 'train'))
shutil.copy(image_path, os.path.join(img_folder, 'train'))
shutil.copy(label_path, os.path.join(label_folder, 'train_original'))
else:
shutil.copy(image_path, os.path.join(img_folder, 'val'))
shutil.copy(label_path, os.path.join(label_folder, 'val_original'))
二.开始训练
(1)下载预训练权重
(2)创建dota8-obb.yaml,修改相关参数
(3)修改yolov8-obb.yaml参数,修改nc
(4)训练
from ultralytics import YOLO
def main():
model = YOLO('yolov8n-obb.yaml').load('yolov8n-obb.pt') # build from YAML and transfer weights
model.train(data='dota8-obb.yaml', epochs=100, imgsz=640, batch=4, workers=4)
if __name__ == '__main__':
main()
参考:
全网首发!Yolov8_obb旋转框训练、测试、推理手把手教学(DOTA1.0数据集map50已达80%)
Yolov8_obb(prob loss) 基于anchor_free的旋转框目标检测,剪枝,跟踪(ByteTracker)
YOLOv8-OBB推理详解及部署实现
roLabelImg的使用
关于旋转框定义的一些理解和感想