Python FAIR1M数据集 计算框的重叠 并去重 美化_polygon.intersection).area

from shapely.geometry import Polygon
import numpy as np
from xml.etree import ElementTree as ET
from xml.etree.ElementTree import ElementTree, Element
import os
iou_thr = 0.8
xmls_path = r"E:\data\FAIR1M\test\labelXmlPre"
xml_list = os.listdir(xmls_path)
save_dir = r"E:\data\FAIR1M\test\test"
 
 
# 美化xml
def prettyXml(element, indent, newline, level = 0): # elemnt为传进来的Elment类，参数indent用于缩进，newline用于换行
    if element:  # 判断element是否有子元素
        if element.text == None or element.text.isspace(): # 如果element的text没有内容
            element.text = newline + indent * (level + 1)
        else:
            element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * (level + 1)
    #else:  # 此处两行如果把注释去掉，Element的text也会另起一行
        #element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * level
    temp = list(element) # 将elemnt转成list
    for subelement in temp:
        if temp.index(subelement) < (len(temp) - 1): # 如果不是list的最后一个元素，说明下一个行是同级别元素的起始，缩进应一致
            subelement.tail = newline + indent * (level + 1)
        else:  # 如果是list的最后一个元素， 说明下一行是母元素的结束，缩进应该少一个
            subelement.tail = newline + indent * level
        prettyXml(subelement, indent, newline, level = level + 1) # 对子元素进行递归操作
 
 
# 创建一个xml文件，并带有基本信息
def create_xml(bbox_list, remove_bbox, save_dir, xml_name):
    annotation = Element("annotation")
    # 创建source及其子元素, 并将内容追加到annotation
    source = Element("source")
    file_name = Element("filename")
    file_name.text = xml_name[0:-4] + ".tif"
    origin = Element("origin")
    origin.text = "GF2/GF3"
    source.append(file_name)
    source.append(origin)
    annotation.append(source)
    # 创建research及其子元素
    research = Element("research")
    version = Element("version")
    version.text = '1.0'
    provider = Element("provider")
    provider.text = "AIRCAS"
    author = Element("author")
    author.text = "QLKT-UPUP"
    pluginname = Element("pluginname")
    pluginname.text = "FAIR1M"
    pluginclass = Element("pluginclass")
    pluginclass.text = "object detection"
    time = Element("time")
    time.text = "2021-09-02"
    research.append(version)
    research.append(provider)
    research.append(author)
    research.append(pluginname)
    research.append(pluginclass)
    research.append(time)
    annotation.append(research)
    # 创建objects
    objects = Element("objects")
    for i in range(len(bbox_list)):
        if i not in remove_bbox:
            # -------------------将剩余的object加入------------------
            obj = Element("object")
            coordinate = Element("coordinate")
            coordinate.text = "pixel"
            typ = Element("type")
            typ.text = "rectangle"
            description = Element("description")
            description.text = "None"
            possibleresult = Element("possibleresult")
            name = Element("name")
            name.text = bbox_list[i][0]
            probability = Element("probability")
            probability.text = str(bbox_list[i][1])
            possibleresult.append(name)
            possibleresult.append(probability)
            points = Element("points")
            point1 = Element("point")
            point1.text = str(bbox_list[i][2]) + ", " + str(bbox_list[i][3])
            point2 = Element("point")
            point2.text = str(bbox_list[i][4]) + ", " + str(bbox_list[i][5])
            point3 = Element("point")
            point3.text = str(bbox_list[i][6]) + ", " + str(bbox_list[i][7])
            point4 = Element("point")
            point4.text = str(bbox_list[i][8]) + ", " + str(bbox_list[i][9])
            point5 = Element("point")
            point5.text = str(bbox_list[i][2]) + ", " + str(bbox_list[i][3])
            points.append(point1)
            points.append(point2)
            points.append(point3)
            points.append(point4)
            points.append(point5)
            obj.append(coordinate)
            obj.append(typ)
            obj.append(description)
            obj.append(possibleresult)
            obj.append(points)
            objects.append(obj)
            # print(1111111111111111)
            # -----------------------------------------------------
    annotation.append(objects)
    # 美化xml文件的结果
    prettyXml(annotation, '\t', '\n')
    tree = ElementTree(annotation)
    # 将结果保存到文件
    tree.write(os.path.join(save_dir, xml_name), encoding='utf-8', xml_declaration=True)
 
 
def intersection(g, p):
    g = np.asarray(g)
    p = np.asarray(p)
    g = Polygon(g[0:8].reshape((4, 2)))
    p = Polygon(p[0:8].reshape((4, 2)))
    if not g.is_valid or not p.is_valid:
        return 0
    inter = Polygon(g).intersection(Polygon(p)).area
    union = g.area + p.area - inter
    if union == 0:
        return 0
    else:
        return inter / union
 
 
for xml_name in xml_list:
    print("-----------------------", xml_name, "----------------------------")
    file_path = os.path.join(xmls_path, xml_name)
    tree = ET.parse(file_path)
    annotation = tree.getroot()
    objects = annotation.find("objects")
    bbox_list = []
    for obj in objects:
        bbox = []
        possibleresult = obj.find('possibleresult')
        cls_name = possibleresult.find('name').text
        probability = float(possibleresult.find('probability').text)
        bbox.append(cls_name)
        bbox.append(probability)
        points = obj.find('points')
 
        for i in range(0, len(points) - 1):  # 只遍历前四个点
            point = points[i]
            xy = point.text.split(",")
            x = float(xy[0])
            y = float(xy[1])
            bbox.append(float('%.2f' % x))
            bbox.append(float('%.2f' % y))
        bbox_list.append(bbox)
    remove_bbox = []
    for i in range(len(bbox_list)):
        # print("____________________________")
        bbox1 = bbox_list[i]
        for j in range(i + 1, len(bbox_list)):
            bbox2 = bbox_list[j]
            iou = intersection(bbox1[2:10], bbox2[2:10])
            if iou > iou_thr:
                if bbox1[1] < bbox2[1] and bbox1[1] < 0.11:  # 比较两个目标的score
                    if i not in remove_bbox:
                        remove_bbox.append(i)
                elif bbox2[1] < 0.11:
                    if j not in remove_bbox:
                        remove_bbox.append(j)
    create_xml(bbox_list, remove_bbox, save_dir, xml_name)
 
        # print(iou)