目标检测旋转框格式转换

概要

本篇博文将对目标检测旋转框格式转换做简单介绍，从带角度的旋转框转为四个角点的旋转框。

数据说明

所使用的数据集为HRSC2016。该数据为舰船数据集。

框的标注格式为（cx, cy, w, h, ang）

转换代码

根据需求，需要将其转为四个角点的旋转框（x1, y1, x2, y2, x3, y3, x4, y4），并保存为txt格式。

import xml.etree.cElementTree as ET
from xml.dom.minidom import Document
import xml.dom.minidom
import numpy as np
import os
import math
import sys
import glob


def coordinate_convert_r(box):
    w, h = box[2:-1]
    theta = -box[-1]
    x_lu, y_lu = -w/2, h/2
    x_ru, y_ru = w/2, h/2
    x_ld, y_ld = -w/2, -h/2
    x_rd, y_rd = w/2, -h/2

    x_lu_ = math.cos(theta)*x_lu + math.sin(theta)*y_lu + box[0]
    y_lu_ = -math.sin(theta)*x_lu + math.cos(theta)*y_lu + box[1]

    x_ru_ = math.cos(theta) * x_ru + math.sin(theta) * y_ru + box[0]
    y_ru_ = -math.sin(theta) * x_ru + math.cos(theta) * y_ru + box[1]

    x_ld_ = math.cos(theta) * x_ld + math.sin(theta) * y_ld + box[0]
    y_ld_ = -math.sin(theta) * x_ld + math.cos(theta) * y_ld + box[1]

    x_rd_ = math.cos(theta) * x_rd + math.sin(theta) * y_rd + box[0]
    y_rd_ = -math.sin(theta) * x_rd + math.cos(theta) * y_rd + box[1]

    convert_box = [x_lu_, y_lu_, x_ru_, y_ru_, x_rd_, y_rd_, x_ld_, y_ld_]

    return convert_box


def read_xml_gtbox_and_label(xml_path):
    """
    :param xml_path: the path of voc xml
    :return: a list contains gtboxes and labels, shape is [num_of_gtboxes, 9],
           and has [x1, y1, x2, y2, x3, y3, x4, y4, label] in a per row
    """

    tree = ET.parse(xml_path)
    root = tree.getroot()
    img_width = None
    img_height = None
    box_list = []
    for child_of_root in root:
     
        if child_of_root.tag == 'Img_SizeWidth':
            img_width = int(child_of_root.text)
        if child_of_root.tag == 'Img_SizeHeight':
            img_height = int(child_of_root.text)
        if child_of_root.tag == 'HRSC_Objects':
            box_list = []
            for child_item in child_of_root:
                if child_item.tag == 'HRSC_Object':
                    label = 1
                    # for child_object in child_item:
                    #     if child_object.tag == 'Class_ID':
                    #         label = NAME_LABEL_MAP[child_object.text]
                    tmp_box = [0., 0., 0., 0., 0.]
                    for node in child_item:
                        if node.tag == 'mbox_cx':
                            tmp_box[0] = float(node.text)
                        if node.tag == 'mbox_cy':
                            tmp_box[1] = float(node.text)
                        if node.tag == 'mbox_w':
                            tmp_box[2] = float(node.text)
                        if node.tag == 'mbox_h':
                            tmp_box[3] = float(node.text)
                        if node.tag == 'mbox_ang':
                            tmp_box[4] = float(node.text)

                    tmp_box = coordinate_convert_r(tmp_box)
                        # assert label is not None, 'label is none, error'
                    tmp_box.append(label)
                    # if len(tmp_box) != 0:
                    box_list.append(tmp_box)
            # box_list = coordinate_convert(box_list)
            # print(box_list)
    gtbox_label = np.array(box_list, dtype=np.int32)

    return img_height, img_width, gtbox_label

def mkdir(path):
    if not os.path.exists(path):
        os.mkdir(path)


if __name__ == '__main__':
    import matplotlib.pyplot as plt
    import cv2 as cv
    from shutil import copyfile
    src_image_path = './AllImages'
    src_xml_path = './Annotations'
    txt_path = './labelTxt'
    out_img_path = './images'
    mkdir(txt_path)
    mkdir(out_img_path)

    src_imgs = glob.glob(f'{src_image_path}/*.bmp')

    for img_path in src_imgs:
        try:
            ori_image = cv.imread(img_path)
            x_path = img_path[:-3].replace('AllImages', 'Annotations') + 'xml'
            img_height, img_width, gtbox_labels = read_xml_gtbox_and_label(x_path)
            if len(gtbox_labels) == 0:
                continue

            # for gtbox_label in gtbox_labels:
            #     tl = np.asarray([gtbox_label[0], gtbox_label[1]], np.float32)
            #     tr = np.asarray([gtbox_label[2], gtbox_label[3]], np.float32)
            #     br = np.asarray([gtbox_label[4], gtbox_label[5]], np.float32)
            #     bl = np.asarray([gtbox_label[6], gtbox_label[7]], np.float32)
            #     box = np.asarray([bl, tl, tr, br], np.float32)
            #     box = np.int0(box)
            #     cv.drawContours(ori_image, [box], 0, (255, 255, 255), 1)
            # plt.imshow(ori_image)
            # plt.show()
            # print(gtbox_labels)
            name = os.path.split(img_path)[-1][:-4]
            for ann in gtbox_labels:
                with open(f"{txt_path}/{name}.txt", 'a+') as f:
                    f.write(f"{ann[0]} {ann[1]} {ann[2]} {ann[3]} {ann[4]} {ann[5]} {ann[6]} {ann[7]} ship 0\n")
            cv.imwrite(f'{out_img_path}/{name}.png', ori_image)
        except:
            print(img_path)