旋转目标检测复现-yolov5-obb
复现源码：
https://github.com/hukaixuan19970627/yolov5_obb
亲测可行
安装流程：
按照https://github.com/hukaixuan19970627/yolov5_obb/blob/master/docs/install.md
确保安装过程不报错，否则影响后续训练
安装成功即可准备数据集
hf_txt存放划分好的训练集、测试集、验证集，里面内容为数据图像文件名，
images存放要训练的图像
labelTxt存放将xml转换后的txt标签文件
hf.py数据集划分;
# -*- coding: utf-8 -*-
import os
import random
trainval_percent = 0.9
train_percent = 0.9
xmlfilepath = 'xml'
txtsavepath = 'images'
total_xml = os.listdir(xmlfilepath)
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
ftrainval = open('hf_txt/trainval.txt', 'w')
ftest = open('hf_txt/test.txt', 'w')
ftrain = open('hf_txt/train.txt', 'w')
fval = open('hf_txt/val.txt', 'w')
for i in list:
name = total_xml[i][:-4] + '\n'
if i in trainval:
ftrainval.write(name)
if i in train:
ftrain.write(name)
else:
fval.write(name)
else:
ftest.write(name)
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()
xml转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 = "/root/autodl-tmp/yolov5_obb/dataset/dataset_demo/xml/" # 目录下保存的是需要转换的xml文件
dotaxml_path = '/root/autodl-tmp/yolov5_obb/dataset/dataset_demo/1xml/'
out_path = '/root/autodl-tmp/yolov5_obb/dataset/dataset_demo/labelTxt/'
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)
voc_label.py划分训练集，验证集，测试集路径：
# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import os
from os import getcwd
sets = ['train', 'val', 'test']
classes = ["large_car","small_car"]
abs_path = os.getcwd()
wd = getcwd()
for image_set in sets:
if not os.path.exists('labelTxt/'):
os.makedirs('labelTxt/')
image_ids = open('hf_txt/%s.txt' % (image_set)).read().strip().split()
list_file = open('%s.txt' % (image_set), 'w')
for image_id in image_ids:
list_file.write(abs_path + '/images/%s.jpg\n' % (image_id))
list_file.close()
修改相应数据集访问路劲
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
#path: ./dataset # dataset root dir
train: dataset/dataset_demo/train.txt #images # train images (relative to 'path')
val: dataset/dataset_demo/val.txt #images # val images (relative to 'path')
#test: dataset_demo/images #images # test images (optional)
# Classes
nc: 2 # number of classes
names: ['large_car','small_car'] # class names
# Download script/URL (optional)
# download: https://ultralytics.com/assets/coco128.zip
最后修改训练文件train.py,修改成对应的文件路径即可
上述都没问题即可训练
python train.py
扩展：部署yolov5-obb：
https://blog.csdn.net/qq_41043389/article/details/127777272