利用DOTA的脚本绘制HRSC数据集的真实框，首先将HRSC的标注文件转换为DOTA格式，然后利用DOTA的脚本绘制目标框

新建文件

进入到HRSC2016的Test目录，新建两个文件夹：

mkdir DOTA_labels DOTA_labels_drawed

新建3个py文件

dota_utils.py：

import sys
import codecs
import numpy as np
import shapely.geometry as shgeo
import os
import re
import math
"""some basic functions which are useful for process DOTA data
"""wordname_18 = ['airport','small-vehicle','large-vehicle','plane','storage-tank','ship','harbor','ground-track-field','soccer-ball-field','tennis-court','swimming-pool','baseball-diamond','roundabout','basketball-court','bridge','helicopter','container-crane','helipad']def custombasename(fullname):return os.path.basename(os.path.splitext(fullname)[0])def GetFileFromThisRootDir(dir,ext = None):allfiles = []needExtFilter = (ext != None)for root,dirs,files in os.walk(dir):for filespath in files:filepath = os.path.join(root, filespath)extension = os.path.splitext(filepath)[1][1:]if needExtFilter and extension in ext:allfiles.append(filepath)elif not needExtFilter:allfiles.append(filepath)return allfilesdef TuplePoly2Poly(poly):outpoly = [poly[0][0], poly[0][1],poly[1][0], poly[1][1],poly[2][0], poly[2][1],poly[3][0], poly[3][1]]return outpolydef parse_dota_poly(filename):"""parse the dota ground truth in the format:[(x1, y1), (x2, y2), (x3, y3), (x4, y4)]"""objects = []#print('filename:', filename)f = []if (sys.version_info >= (3, 5)):fd = open(filename, 'r')f = fdelif (sys.version_info >= 2.7):fd = codecs.open(filename, 'r')f = fd# count = 0while True:line = f.readline()# count = count + 1# if count < 2:#     continueif line:splitlines = line.strip().split(' ')object_struct = {}### clear the wrong name after check all the data#if (len(splitlines) >= 9) and (splitlines[8] in classname):if (len(splitlines) < 9):continueif (len(splitlines) >= 9):object_struct['name'] = splitlines[8]if (len(splitlines) == 9):object_struct['difficult'] = '0'elif (len(splitlines) >= 10):# if splitlines[9] == '1':# if (splitlines[9] == 'tr'):#     object_struct['difficult'] = '1'# else:object_struct['difficult'] = splitlines[9]# else:#     object_struct['difficult'] = 0object_struct['poly'] = [(float(splitlines[0]), float(splitlines[1])),(float(splitlines[2]), float(splitlines[3])),(float(splitlines[4]), float(splitlines[5])),(float(splitlines[6]), float(splitlines[7]))]gtpoly = shgeo.Polygon(object_struct['poly'])object_struct['area'] = gtpoly.area# poly = list(map(lambda x:np.array(x), object_struct['poly']))# object_struct['long-axis'] = max(distance(poly[0], poly[1]), distance(poly[1], poly[2]))# object_struct['short-axis'] = min(distance(poly[0], poly[1]), distance(poly[1], poly[2]))# if (object_struct['long-axis'] < 15):#     object_struct['difficult'] = '1'#     global small_count#     small_count = small_count + 1objects.append(object_struct)else:breakreturn objectsdef dots4ToRecC(poly, img_w, img_h):xmin, ymin, xmax, ymax = dots4ToRec4(poly)x = (xmin + xmax)/2y = (ymin + ymax)/2w = xmax - xminh = ymax - yminreturn x/img_w, y/img_h, w/img_w, h/img_hdef parse_dota_poly2(filename):"""parse the dota ground truth in the format:[x1, y1, x2, y2, x3, y3, x4, y4]"""objects = parse_dota_poly(filename)for obj in objects:obj['poly'] = TuplePoly2Poly(obj['poly'])obj['poly'] = list(map(int, obj['poly']))return objectsdef parse_dota_rec(filename):"""parse the dota ground truth in the bounding box format:"xmin, ymin, xmax, ymax""""objects = parse_dota_poly(filename)for obj in objects:poly = obj['poly']bbox = dots4ToRec4(poly)obj['bndbox'] = bboxreturn objects
## bounding box transfer for varies formatdef dots4ToRec4(poly):xmin, xmax, ymin, ymax = min(poly[0][0], min(poly[1][0], min(poly[2][0], poly[3][0]))), \max(poly[0][0], max(poly[1][0], max(poly[2][0], poly[3][0]))), \min(poly[0][1], min(poly[1][1], min(poly[2][1], poly[3][1]))), \max(poly[0][1], max(poly[1][1], max(poly[2][1], poly[3][1])))return xmin, ymin, xmax, ymax
def dots4ToRec8(poly):xmin, ymin, xmax, ymax = dots4ToRec4(poly)return xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax#return dots2ToRec8(dots4ToRec4(poly))
def dots2ToRec8(rec):xmin, ymin, xmax, ymax = rec[0], rec[1], rec[2], rec[3]return xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymaxdef groundtruth2Task1(srcpath, dstpath):filelist = GetFileFromThisRootDir(srcpath)# names = [custombasename(x.strip())for x in filelist]filedict = {}for cls in wordname_15:fd = open(os.path.join(dstpath, 'Task1_') + cls + r'.txt', 'w')filedict[cls] = fdfor filepath in filelist:objects = parse_dota_poly2(filepath)subname = custombasename(filepath)pattern2 = re.compile(r'__([\d+\.]+)__\d+___')rate = re.findall(pattern2, subname)[0]for obj in objects:category = obj['name']difficult = obj['difficult']poly = obj['poly']if difficult == '2':continueif rate == '0.5':outline = custombasename(filepath) + ' ' + '1' + ' ' + ' '.join(map(str, poly))elif rate == '1':outline = custombasename(filepath) + ' ' + '0.8' + ' ' + ' '.join(map(str, poly))elif rate == '2':outline = custombasename(filepath) + ' ' + '0.6' + ' ' + ' '.join(map(str, poly))filedict[category].write(outline + '\n')def Task2groundtruth_poly(srcpath, dstpath):thresh = 0.1filedict = {}Tasklist = GetFileFromThisRootDir(srcpath, '.txt')for Taskfile in Tasklist:idname = custombasename(Taskfile).split('_')[-1]# idname = datamap_inverse[idname]f = open(Taskfile, 'r')lines = f.readlines()for line in lines:if len(line) == 0:continue# print('line:', line)splitline = line.strip().split(' ')filename = splitline[0]confidence = splitline[1]bbox = splitline[2:]if float(confidence) > thresh:if filename not in filedict:# filedict[filename] = codecs.open(os.path.join(dstpath, filename + '.txt'), 'w', 'utf_16')filedict[filename] = codecs.open(os.path.join(dstpath, filename + '.txt'), 'w')# poly = util.dots2ToRec8(bbox)poly = bbox#               filedict[filename].write(' '.join(poly) + ' ' + idname + '_' + str(round(float(confidence), 2)) + '\n')# print('idname:', idname)# filedict[filename].write(' '.join(poly) + ' ' + idname + '_' + str(round(float(confidence), 2)) + '\n')filedict[filename].write(' '.join(poly) + ' ' + idname + '\n')def polygonToRotRectangle(bbox):""":param bbox: The polygon stored in format [x1, y1, x2, y2, x3, y3, x4, y4]:return: Rotated Rectangle in format [cx, cy, w, h, theta]"""bbox = np.array(bbox,dtype=np.float32)bbox = np.reshape(bbox,newshape=(2,4),order='F')angle = math.atan2(-(bbox[0,1]-bbox[0,0]),bbox[1,1]-bbox[1,0])center = [[0],[0]]for i in range(4):center[0] += bbox[0,i]center[1] += bbox[1,i]center = np.array(center,dtype=np.float32)/4.0R = np.array([[math.cos(angle), -math.sin(angle)], [math.sin(angle), math.cos(angle)]], dtype=np.float32)normalized = np.matmul(R.transpose(),bbox-center)xmin = np.min(normalized[0,:])xmax = np.max(normalized[0,:])ymin = np.min(normalized[1,:])ymax = np.max(normalized[1,:])w = xmax - xmin + 1h = ymax - ymin + 1return [float(center[0]),float(center[1]),w,h,angle]

hrsc2dota.py

import xml.etree.ElementTree as ET
import os
import math
import cv2
import numpy as np
def get_label(xml_path):in_file = open(xml_path)tree=ET.parse(in_file)root = tree.getroot()labels = []for obj in root.iter('HRSC_Object'):difficult = obj.find('difficult').textclass_id = int(obj.find('Class_ID').text) % 100# class_id = 0 # 标签对应关系自行修改# if int(difficult) == 1:#     continuembox_cx, mbox_cy, mbox_w, mbox_h, mbox_ang = (float(obj.find('mbox_cx').text),float(obj.find('mbox_cy').text),float(obj.find('mbox_w').text),float(obj.find('mbox_h').text),float(obj.find('mbox_ang').text))labels.append([class_id,mbox_cx, mbox_cy, mbox_w, mbox_h,mbox_ang])return labels
# 计算旋转框四个顶点的坐标
def get_rotated_box_vertices(labels,label_path):with open(label_path,'w') as f:for i in range(len(labels)):class_id,mbox_cx, mbox_cy, mbox_w, mbox_h,angle_rad= labels[i]rotation_matrix = np.array([[np.cos(angle_rad), -np.sin(angle_rad)],[np.sin(angle_rad), np.cos(angle_rad)]])box_half_width = mbox_w / 2box_half_height = mbox_h / 2box_vertices = np.array([[-box_half_width, -box_half_height],[box_half_width, -box_half_height],[box_half_width, box_half_height],[-box_half_width, box_half_height]])rotated_vertices = np.dot(box_vertices, rotation_matrix.T)rotated_vertices[:, 0] += mbox_cxrotated_vertices[:, 1] += mbox_cyrotated_vertices = np.round(rotated_vertices).astype(np.int32)# print(rotated_vertices)# f.write(" ".join([str(a) for a in rotated_vertices]) + '\n')rotated_vertices = rotated_vertices.reshape(-1)f.write(" ".join([str(a) for a in rotated_vertices]) + " " + str(class_id) + '\n')# return rotated_vertices_listxml_root = r"Annotations"
txt_root = r"DOTA_labels"xml_name = os.listdir(xml_root)
# print(len(xml_name))
for i in range(len(xml_name)):xml_path = os.path.join(xml_root,xml_name[i])if xml_path == "Annotations/.ipynb_checkpoints":continuetxt_path = os.path.join(txt_root,xml_name[i].split('.')[0]+'.txt')get_rotated_box_vertices(get_label(xml_path),txt_path)

dota_drawed.py

import xml.etree.ElementTree as ET
import os
import math
import cv2
import numpy as np
import dota_utils as util
import random# 手动输入cx cy w h angle进行绘制
# from HRSC_to_DOTA import get_rotated_box_vertices
# cx = 569.5045
# cy = 263.4875
# w = 261.0578
# h = 65.08137
# angle = -1.562451
# vertices = get_rotated_box_vertices(cx, cy, w, h, angle)
# vertices = np.array(vertices,dtype=np.int32)
# img = cv2.imread(r'AllImages\100000640.bmp')
# cv2.polylines(img,[vertices], isClosed=True, color=(255, 0, 0), thickness=2)
# cv2.imshow('test',img)
# cv2.waitKey(0)
# cv2.destroyAllWindows()img_root = r"AllImages"
label_root = r"DOTA_labels"
drawed_img_root = r"DOTA_labels_drawed"
image_name = os.listdir(img_root)
for i in range(len(image_name)):img_path = os.path.join(img_root,image_name[i])label_path = os.path.join(label_root,image_name[i].split('.')[0]+'.txt')drawed_img_path = os.path.join(drawed_img_root,image_name[i])objects = util.parse_dota_poly(label_path)print(objects)img = cv2.imread(img_path)poly = []for i in range(len(objects)):poly.append(np.array(objects[i]['poly'],dtype=np.int32))print(poly)cv2.polylines(img,poly, isClosed=True, color=(35, 37, 133), thickness=2)drawed_img_path = drawed_img_path.replace('.bmp','.png')cv2.imwrite(drawed_img_path,img)

运行

conda activate mmrotate

python hrsc2dota.py

可以看到已经转换好的DOTA格式标注文件

python dota_drawed.py

可以看到绘制好的图片

说明：在dota_drawed.py文件中，40行我改了一下图片后缀，默认是bmp格式
在39行可以改标注框的颜色：在线取色器

参考：
《HRSC2016_dataset 旋转框 + 划分数据集》
《深度学习|dota格式的txt文件转化为yolo格式的txt文件》