旋转框四点坐标转成roLabelImg的角度格式

甲方发过来的标签是四点坐标格式的,我们需要用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
更多坐标和角度转换可以参考杨学小哥的代码,我这个也是在上面略做了点修改。

  • 3
    点赞
  • 33
    收藏
    觉得还不错? 一键收藏
  • 2
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值