一、环境配置
- Ubuntu 22.04
- Torch 1.10
- CUDA 11.3
- python 3.9
环境配置参考下面链接(建议Linux系统)yolov5_obb/install.md at master · hukaixuan19970627/yolov5_obb (github.com)https://github.com/hukaixuan19970627/yolov5_obb/blob/master/docs/install.md
- 创建虚拟环境
conda create -n Py39_Torch1.10_cu11.3 python=3.9 -y
source activate Py39_Torch1.10_cu11.3
- 查看CUDA版本(确保CUDA runtime api version ≤ CUDA driver version)
nvcc -V
nvidia-smi
- 安装PyTorch和torchvision
pip3 install torch==1.10.1+cu113 torchvision==0.11.2+cu113 torchaudio==0.10.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
nvcc -V
python
>>> import torch
>>> torch.version.cuda
>>> exit()
- 克隆yolov7-obb项目(也可直接在GitHub上下载)
git clone https://github.com/Egrt/yolov7-obb.git
- 配置yolov7-obb所需环境
cd yolov7-obb
pip install -r requirements.txt
cd utils/nms_rotated
python setup.py develop #安装非极大值抑制库
运行 python setup.py develop时报错
可以看出时gcc版本问题,检测已有版本发现gcc 12 > gcc 11
故安装低版本的gcc,并创建软链接(此处要对应改为自己的位置)
sudo apt-get install gcc-10
sudo apt-get install g++-10
sudo ln -s /usr/bin/gcc-10 /usr/local/cuda-11.6/bin/gcc #创建软链接
再次运行 python setup.py develop
至此,环境配置完成!!!
二、数据集制作
- 下载官方源码
https://github.com/cgvict/roLabelImg
- 进入下载好的roLabellmg-master文件夹内 ,在终端打开
pyrcc5 -oresources.py resources.qrc
python roLabelImg.py
- 快捷键
w | 创建矩形框 |
e | 创建旋转矩形框 |
d |
下一张 |
a | 上一张 |
zxcv | 旋转矩形框 |
Ctrl + s | 保存 |
- 标注的XML格式
<annotation verified="yes">
<folder>hsrc</folder>
<filename>100000001</filename>
<path>/Users/haoyou/Library/Mobile Documents/com~apple~CloudDocs/OneDrive/hsrc/100000001.bmp</path>
<source>
<database>Unknown</database>
</source>
<size>
<width>1166</width>
<height>753</height>
<depth>3</depth>
</size>
<segmented>0</segmented>
<object>
<type>bndbox</type>
<name>ship</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>178</xmin>
<ymin>246</ymin>
<xmax>974</xmax>
<ymax>504</ymax>
</bndbox>
</object>
</annotation>
- 数据集格式转换(把旋转框 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/Annotations" # 目录下保存的是需要转换的xml文件
dotaxml_path = './datasets/dotaxml'
filelist = os.listdir(roxml_path)
for file in filelist:
edit_xml(os.path.join(roxml_path, file), os.path.join(dotaxml_path, file))
三、训练自己的数据集
本文参考yolov7-obb项目如下Egrt/yolov7-obb: 在YOLOv7的基础上使用KLD损失修改为旋转目标检测yolov7-obb (github.com)https://github.com/Egrt/yolov7-obb
- 数据集格式 (其中Annotations为上面格式转换过的xml文件)
VOCdevkit/VOC2007
├── Annotations
├── 0001.xml
├── 0002.xml
.
.
├── ImageSets
├── Main
├── JPEGImages
├── 0001.xml
├── 0002.xml
.
.
- 运行voc_annotation.py文件,生成2007_train.txt和2007_val.txt文件
- 开始训练