200字范文,内容丰富有趣,生活中的好帮手!
200字范文 > 基于MMRotate训练自定义数据集 做旋转目标检测 -3-30

基于MMRotate训练自定义数据集 做旋转目标检测 -3-30

时间:2019-07-23 08:29:32

相关推荐

基于MMRotate训练自定义数据集 做旋转目标检测 -3-30

本文目录

简述1.MMrotate下载2.环境安装3.自定义数据集制作3.1 roLabelImg 打标签3.2 生成DOTA数据集格式的标签3.3 数据集裁剪(split)4.修改配置文件5.训练并测试6 公开数据集及对应代码下载网址参考博客

简述

MMRotate 是一款基于 PyTorch 的旋转框检测的开源工具箱,是 OpenMMLab 项目的成员之一。里面包含了rcnn、faster rcnn、r3det等各种旋转目标的检测模型,适合于遥感图像领域的目标检测。

1.MMrotate下载

MMrotate包下载:下载链接

目录结构如下

2.环境安装

所需要的依赖环境

Linux & WindowsPython 3.7+PyTorch 1.6+CUDA 9.2+GCC 5+mmcv-full 1.4.5+mmdet 2.19.0+

安装流程:

#首先,创建虚拟一个虚拟环境 mmrotateconda create -n mmrotate python=3.7 -yconda activate mmrotate# 1.安装 pytorch及cudaconda install pytorch==1.7.0 torchvision==0.8.0 cudatoolkit=10.1 -c pytorch# 2.安装最新版本的 mmcvpip install mmcv-full -f /mmcv/dist/cu101/torch1.7.0/index.html# 3.安装 mmdetectionpip install mmdet# 4.安装 mmrotatecd mmrotate # 转到mmrotate目录下pip install -r requirements/build.txtpip 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.1 roLabelImg 打标签

roLabelImg工具算是labelme的进阶版吧,它能把矩形框进行旋转,制作旋转目标的标签。操作时:A键 D键 Z键 X键 C键 V键 这几个常用快捷键用得多点

生成的xml标签文件如下

其中相比普通的标签文件多了一个角度信息(angle),后续转为DOTA的标签时就是利用了该角度(这里强烈建议先看大佬的文章:旋转框的转换问题),所以在用roLabelImg打标签时一定要注意,roLabelImg中的angle是以旋转框的w边 与x轴顺时针之间所形成的夹角。

生成的目录结构如下:

3.2 生成DOTA数据集格式的标签

利用上一步中xml标签 转换为DOTA的标签,转换的代码(只需要给定3个路径参数就可以运行了):

import osimport xml.etree.ElementTree as ETimport mathimport cv2 as cvdef voc_to_dota(xml_dir, xml_name, img_dir, savedImg_dir):txt_name = xml_name[:-4] + '.txt' # txt文件名字:去掉xml 加上.txttxt_path = xml_dir + '/txt_label' # txt文件目录:在xml目录下创建的txtl_label文件夹if not os.path.exists(txt_path):os.makedirs(txt_path)txt_file = os.path.join(txt_path, txt_name) # txt完整的含名文件路径img_name = xml_name[:-4] + '.jpg'# 图像名字img_path = os.path.join(img_dir, img_name) # 图像完整路径img = cv.imread(img_path) # 读取图像xml_file = os.path.join(xml_dir, xml_name) tree = ET.parse(os.path.join(xml_file)) # 解析xml文件 然后转换为DOTA格式文件root = tree.getroot() 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').textdifficult = 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)if not os.path.exists(savedImg_dir):os.makedirs(savedImg_dir)out_img = os.path.join(savedImg_dir, xml_name[:-4]+'.jpg')cv.imwrite(out_img, 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 - xcyoff = yp - yccosTheta = math.cos(theta)sinTheta = math.sin(theta)pResx = cosTheta * xoff + sinTheta * yoffpResy = - sinTheta * xoff + cosTheta * yoff# pRes = (xc + pResx, yc + pResy)# 保留一位小数点return float(format(xc + pResx, '.1f')), float(format(yc + pResy, '.1f'))# return xc + pResx, yc + pResyimport argparsedef parse_args():parser = argparse.ArgumentParser(description='数据格式转换')parser.add_argument('--xml-dir', default='./buildings/train/label', help='original xml file dictionary')parser.add_argument('--img-dir', default='./buildings/train', help='original image dictionary')parser.add_argument('--outputImg-dir', default='./buildings/train/out', help='saved image dictionary after dealing ')args = parser.parse_args()return argsif __name__ == '__main__':args = parse_args()xml_path = args.xml_dirxmlFile_list = os.listdir(xml_path)print(xmlFile_list)for i in range(0, len(xmlFile_list)):if ('.xml' in xmlFile_list[i]) or ('.XML' in xmlFile_list[i]):voc_to_dota(xml_path, xmlFile_list[i], args.img_dir, args.outputImg_dir)print('----------------------------------------{}{}----------------------------------------'.format(xmlFile_list[i], ' has Done!'))else:print(xmlFile_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 改小一点。

测试效果:

6 公开数据集及对应代码下载网址

比较全面的数据集网站: 根据关键词搜索基本上都能找到相关的数据集

参考博客

AI目标检测:VOC格式数据集转换为DOTA类型数据集

使用MMrotate的Kfiou进行身份证分割

本内容不代表本网观点和政治立场,如有侵犯你的权益请联系我们处理。
网友评论
网友评论仅供其表达个人看法,并不表明网站立场。