甲方发过来的标签是四点坐标格式的,我们需要用roLabelImg检查一遍图片和标签,这就需要转成角度格式,roLabelImg跟别的表示方法不一样,是逆时针为正的,范围是0-π。脚本里用到了OpenCV的最小外接矩形函数cv2.minAreaRect,以及根据长短边来判断角度在哪个象限。具体代码和详细注释如下:
# *_* coding : UTF-8 *_*
# 开发人员: csu·pan-_-||
# 开发时间: 2020/11/17 16:24
# 文件名称: four_to_theta.py
# 开发工具: PyCharm
# 功能描述: 将四点坐标x1,y1,x2,y2,x3,y3,x4,y4转成旋转框 cx,cy,w,h,angle
# 以便rolabelimg能够识别
import os
import xml.etree.ElementTree as ET
import math
import numpy as np
import cv2
def edit_xml(xml_file):
"""
修改xml文件
:param xml_file:xml文件的路径
:return:
"""
tree = ET.parse(xml_file)
objs = tree.findall('object') # 获取节点
for ix, obj in enumerate(objs):
cx = ET.Element("cx") # 创建节点
cy = ET.Element("cy")
w = ET.Element("w")
h = ET.Element("h")
angle = ET.Element("angle")
type = ET.Element("type")
print(xml_file)
type.text = 'robndbox'
obj.append(type) # 新增节点
obj_bnd = obj.find('bndbox')
obj_bnd.tag = 'robndbox' # 修改节点名
obj_x0 = obj_bnd.find('x0')
obj_y0 = obj_bnd.find('y0')
obj_x1 = obj_bnd.find('x1')
obj_y1 = obj_bnd.find('y1')
obj_x2 = obj_bnd.find('x2')
obj_y2 = obj_bnd.find('y2')
obj_x3 = obj_bnd.find('x3')
obj_y3 = obj_bnd.find('y3')
x0 = float(obj_x0.text) # 获取节点的值
y0 = float(obj_y0.text)
x1 = float(obj_x1.text)
y1 = float(obj_y1.text)
x2 = float(obj_x2.text)
y2 = float(obj_y2.text)
x3 = float(obj_x3.text)
y3 = float(obj_y3.text)
# 用OpenCV的最小矩形转换成-90到0的角度
boxes = backward_convert([x0,y0,x1,y1,x2,y2,x3,y3])
# 根据长短边转成 0 到 180
new_boxes = coordinate_present_convert(boxes)
# 转成弧度:
new_boxes[0][-1] = new_boxes[0][-1] * math.pi/180
new_boxes = new_boxes.astype(np.str)
cx.text,cy.text,w.text,h.text,angle.text = new_boxes[0]
obj_bnd.remove(obj_x0) # 删除节点
obj_bnd.remove(obj_y0)
obj_bnd.remove(obj_x1)
obj_bnd.remove(obj_y1)
obj_bnd.remove(obj_x2)
obj_bnd.remove(obj_y2)
obj_bnd.remove(obj_x3)
obj_bnd.remove(obj_y3)
obj_bnd.append(cx) # 新增节点
obj_bnd.append(cy)
obj_bnd.append(w)
obj_bnd.append(h)
obj_bnd.append(angle)
tree.write(xml_file, method='xml', encoding='utf-8') # 更新xml文件
def coordinate_present_convert(coords, shift=True):
"""
:param coords: shape [-1, 5]
:param shift: [-90, 90) --> [-180, 0)
:return: shape [-1, 5]
"""
# angle range from [-90, 0) to [0,180)
w, h = coords[:, 2], coords[:, 3]
remain_mask = np.greater(w, h)
convert_mask = np.logical_not(remain_mask).astype(np.int32)
remain_mask = remain_mask.astype(np.int32)
remain_coords = coords * np.reshape(remain_mask, [-1, 1])
coords[:, [2, 3]] = coords[:, [3, 2]]
coords[:, 4] += 90
convert_coords = coords * np.reshape(convert_mask, [-1, 1])
coords_new = remain_coords + convert_coords
if shift:
if coords_new[:, 4] >= 0:
coords_new[:, 4] = coords_new[:, 4] -180
return np.array(coords_new, dtype=np.float32)
def backward_convert(coordinate):
"""
:param coordinate: format [x1, y1, x2, y2, x3, y3, x4, y4]
:return: format [x_c, y_c, w, h, theta, (label)]
"""
boxes = []
box = np.int0(coordinate)
box = box.reshape([4, 2])
rect1 = cv2.minAreaRect(box)
x, y, w, h, theta = rect1[0][0], rect1[0][1], rect1[1][0], rect1[1][1], rect1[2]
if theta == 0:
w, h = h, w
theta -= 90
boxes.append([x, y, w, h, theta])
return np.array(boxes, dtype=np.float32)
if __name__ == '__main__':
dir=r"E:\Projects\xml_four"
filelist = os.listdir(dir)
for file in filelist:
edit_xml(os.path.join(dir,file))
参考:coordinate_convert.py
更多坐标和角度转换可以参考杨学小哥的代码,我这个也是在上面略做了点修改。