代码收藏家 基于MMRotate训练自定义数据集 做旋转目标检测 2022-3-30
本文目录
简述
MMRotate 是一款基于 PyTorch 的旋转框检测的开源工具箱，是 OpenMMLab 项目的成员之一。里面包含了rcnn、faster rcnn、r3det等各种旋转目标的检测模型，适合于遥感图像领域的目标检测。
1.MMrotate下载
MMrotate包下载：下载链接
目录结构如下：
2.环境安装
所需要的依赖环境：
Linux & Windows
Python 3.7+
PyTorch 1.6+
CUDA 9.2+
GCC 5+
mmcv-full 1.4.5+
mmdet 2.19.0+
安装流程：
#首先，创建虚拟一个虚拟环境 mmrotate
conda create -n mmrotate python=3.7 -y
conda activate mmrotate
# 1.安装 pytorch及cuda
conda install pytorch==1.7.0 torchvision==0.8.0 cudatoolkit=10.1 -c pytorch
# 2.安装最新版本的 mmcv
pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu101/torch1.7.0/index.html
# 3.安装 mmdetection
pip install mmdet
# 4.安装 mmrotate
cd mmrotate # 转到mmrotate目录下
pip install -r requirements/build.txt
pip install -v -e . # or "python setup.py develop"
测试安装的环境是否正常：
打开demo文件夹下的image_demo.py：修改相应的参数：
– -img : 测试的图片
– -config：选择测试的模型配置文件
– -checkpoint: 对应的预训练模型文件，可自行下载：MMrotate模型库 百度网盘连接：rcnn 权重 提取码：ur0b
– -output: 检测输出的图片
然后，执行并得到如下结果，说明环境安装正确：
python demo/image_demo.py
3.自定义数据集制作
要训练自己的模型，自定义数据集制作这部分其实是最麻烦的。MMrotate所使用的数据集格式是dota类型的，图片为.png格式且尺寸是 n×n 的（方形），不过不用担心，该项目中有相应的工具包可自动转换
DOTA的标签格式：
x1, y1, x2, y2, x3, y3, x4, y4, category: 目标名字 difficult：表示标签检测的难易程度 （1表示困难，0表示不困难）
x1,y1为左上角的坐标，然后顺时针排列4个坐标
3.2 生成DOTA数据集格式的标签
利用上一步中xml标签 转换为DOTA的标签，转换的代码：
import os
import xml.etree.ElementTree as ET
import math
import cv2 as cv
def voc_to_dota(xml_path, xml_name):
txt_name = xml_name[:-4] + '.txt'
txt_path = xml_path + '/txt_label'
if not os.path.exists(txt_path):
os.makedirs(txt_path)
txt_file = os.path.join(txt_path, txt_name)
file_path = os.path.join(xml_path, file_list[i])
tree = ET.parse(os.path.join(file_path))
root = tree.getroot()
# print(root[6][0].text)
image_path = 'data/dota/images/'
out_path = 'data/dota/outputImg/'
filename = image_path + xml_name[:-4] + '.jpg'
img = cv.imread(filename)
with open(txt_file, "w+", encoding='UTF-8') as out_file:
# out_file.write('imagesource:null' + '\n' + 'gsd:null' + '\n')
for obj in root.findall('object'):
name = obj.find('name').text
difficult = obj.find('difficult').text
# print(name, difficult)
robndbox = obj.find('robndbox')
cx = float(robndbox.find('cx').text)
cy = float(robndbox.find('cy').text)
w = float(robndbox.find('w').text)
h = float(robndbox.find('h').text)
angle = float(robndbox.find('angle').text)
# print(cx, cy, w, h, angle)
p0x, p0y = rotatePoint(cx, cy, cx - w / 2, cy - h / 2, -angle)
p1x, p1y = rotatePoint(cx, cy, cx + w / 2, cy - h / 2, -angle)
p2x, p2y = rotatePoint(cx, cy, cx + w / 2, cy + h / 2, -angle)
p3x, p3y = rotatePoint(cx, cy, cx - w / 2, cy + h / 2, -angle)
# 找最左上角的点
dict = {p0y:p0x, p1y:p1x, p2y:p2x, p3y:p3x }
list = find_topLeftPopint(dict)
#print((list))
if list[0] == p0x:
list_xy = [p0x, p0y, p1x, p1y, p2x, p2y, p3x, p3y]
elif list[0] == p1x:
list_xy = [p1x, p1y, p2x, p2y, p3x, p3y, p0x, p0y]
elif list[0] == p2x:
list_xy = [p2x, p2y, p3x, p3y, p0x, p0y, p1x, p1y]
else:
list_xy = [p3x, p3y, p0x, p0y, p1x, p1y, p2x, p2y]
# 在原图上画矩形 看是否转换正确
cv.line(img, (int(list_xy[0]), int(list_xy[1])), (int(list_xy[2]), int(list_xy[3])), color=(255, 0, 0), thickness= 3)
cv.line(img, (int(list_xy[2]), int(list_xy[3])), (int(list_xy[4]), int(list_xy[5])), color=(0, 255, 0), thickness= 3)
cv.line(img, (int(list_xy[4]), int(list_xy[5])), (int(list_xy[6]), int(list_xy[7])), color=(0, 0, 255), thickness= 2)
cv.line(img, (int(list_xy[6]), int(list_xy[7])), (int(list_xy[0]), int(list_xy[1])), color=(255, 255, 0), thickness= 2)
data = str(list_xy[0]) + " " + str(list_xy[1]) + " " + str(list_xy[2]) + " " + str(list_xy[3]) + " " + \
str(list_xy[4]) + " " + str(list_xy[5]) + " " + str(list_xy[6]) + " " + str(list_xy[7]) + " "
data = data + name + " " + difficult + "\n"
out_file.write(data)
cv.imwrite(out_path + xml_name[:-4] + '.jpg', img)
def find_topLeftPopint(dict):
dict_keys = sorted(dict.keys()) # y值
temp = [dict[dict_keys[0]], dict[dict_keys[1]]]
minx = min(temp)
if minx == temp[0]:
miny = dict_keys[0]
else:
miny = dict_keys[1]
return [minx, miny]
# 转换成四点坐标
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
# pRes = (xc + pResx, yc + pResy)
# 保留一位小数点
return float(format(xc + pResx, '.1f')), float(format(yc + pResy, '.1f'))
# return xc + pResx, yc + pResy
if __name__ == '__main__':
root_path = 'data/dota/xml'
file_list = os.listdir(root_path)
for i in range(0, len(file_list)):
if ('.xml' in file_list[i]) or ('.XML' in file_list[i]):
voc_to_dota(root_path, file_list[i])
print('----------------------------------------{}{}----------------------------------------'
.format(file_list[i], ' has Done!'))
else:
print(file_list[i] + ' is not xml file')
其中转换的原理就是基于大佬的知乎文章中转换矩阵，生成的DOTA标签格式如下：
目前所得到文件夹结构如下：
images: 存放的.jpg图像（我只有200张）
xml: rolabelImg打的标签(这时已没啥用了，可删除)
txt_Label: 生成的DOTA标签
然后将上述数据集分成 train、test、val 、trainval等几部分以便于训练，其中每个部分的文件夹下都包含有 images(图像) 和 labels(对应的txt标签)
我这里是手动划分的：train 80%，test 10%，val 10%。
3.3 数据集裁剪（split）
然后将得到的 train、test、val 中的图片进行裁剪，裁剪的原因 是由于我的图像是1920×1080的矩形，而且也不是png格式的，所以找到：mmrotate-main/tools/data/dota/split/ 路径下img_split.py文件(裁剪脚本) 以及 mmrotate-main/tools/data/dota/split/split_configs/ 路径下的配置文件，其文件内容就是img_split.py的配置信息，我们需要修改其中的参数，让其加载上述的train、test、val中的图像及标签，并进行裁剪。
… 
修改好上述几个关键的参数之后，即可运行img_split.py进行裁剪了：
python mmrotate-main/tools/data/dota/split/img_split.py --base_json mmrotate-main/tools/data/dota/split/split_configs/ss_train.json
其他 test、val、trainval数据集的裁剪同理，只需修改 – -base_json参数为对应的配置文件即可。
裁剪的结果：
4.修改配置文件
数据集准备好之后，接下来，需要修改训练相关的配置信息：主要就是 tools/tain.py 中的相关参数修改，train.py文件如下：
两个主要的参数： – -config: 使用的模型文件 （我使用的是 faster rcnn) ； – -work-dir：训练得到的模型及配置信息保存的路径
由于使用的模型是：rotated_faster_rcnn_r50_fpn_x1_dota_le90.py，在该文件中实现了：训练相关参数的配置，数据集信息的配置，以及faster-rcnn网络架构的搭建；打开该文件修改其中的目标类别数为自己数据集的类别数，我的类别为1， 所以修改：num_classes = 1,
同时，修改 mmrotate-main/mmrotate/datasets/dota.py 文件中的类别名字（CLASSES）， 具体如下：
最后，修改训练使用的数据集路径：找到并打开 mmrotate-main/configs/base/datasets/dotav1.py 文件，修改其中的 data_root 路径为自己裁剪的数据集路径，该路径包含了上述划分的train、test、val、trainval数据集
5.训练并测试
一切准备就绪后 即可开始愉快的训练了！！！执行：
python tools/train.py
训练结果：
保存训练的结果在 run 目录下：
如果遇到 Cuda out of memory错误：可将 mmrotate-main/configs/base/datasets/dotav1.py文件中的samples_per_gpu 和 workers_per_gpu 改小一点。
测试效果：
来源：YD-阿三
