Python手动修改FAIR1M数据集标签类别

import cv2
import os
import random
import xml.etree.ElementTree as ET
import tkinter as tk
from PIL import Image, ImageTk

imgs_path = r"E:\data\FAIR1M\test\images"
labels_path = r"E:\data\FAIR1M\test\test"
classes = ['Boeing737', 'Boeing747', 'Boeing777', 'Boeing787', 'C919', 'A220', 'A321', 'A330', 'A350', 'ARJ21',
           'other-airplane', 'Passenger Ship', 'Motorboat', 'Fishing Boat', 'Tugboat', 'Engineering Ship',
           'Liquid Cargo Ship', 'Dry Cargo Ship', 'Warship', 'other-ship', 'Small Car', 'Bus', 'Cargo Truck',
           'Dump Truck', 'Van', 'Trailer', 'Tractor', 'Excavator', 'Truck Tractor', 'other-vehicle', 'Basketball Court',
           'Tennis Court', 'Football Field', 'Baseball Field', 'Intersection', 'Roundabout', 'Bridge']
imgs_list = os.listdir(imgs_path)
# random.shuffle(imgs_list)
imgs_list.sort()
labels_list = os.listdir(labels_path)
print("请输入从第几个图片开始（上次处理处理到的图片）：")
start = int(input())
cv2.namedWindow("FAIR1M", 0)
# cv2.namedWindow("Obj", cv2.WINDOW_AUTOSIZE)

def mouse(event, x, y, flags, param):
    if event == cv2.EVENT_LBUTTONDOWN:
        print(filename)
        cv2.waitKey(4000)


def btn_click():
    print("11111")


for i in range(start, len(imgs_list)):
    filename = imgs_list[i]
    img_num = filename.split(".")[0]  # 获得图片的序号
    print("------------------------------------------------------------------------------")
    print("doing ", i, "-------", filename, "img.......")
    with open(os.path.join(labels_path, img_num + ".xml"), "r", encoding='UTF-8') as xml_file:
        tree = ET.parse(xml_file)
        root = tree.getroot()
        size = root.find('size')
        # w = int(size.find('width').text)  # 图片的宽
        # h = int(size.find('height').text)
        # width = image.shape[1]
        # hight = image.shape[0]

        for obj in root.iter('object'):

            possibleresult = obj.find('possibleresult')
            cls_name = possibleresult.find('name').text
            probability = possibleresult.find('probability').text
            points = obj.find('points')
            bbox = []
            # if cls_name in ["Small Car", "Van", "Dump Truck", "Cargo Truck", "other-ship", "other-airplane", "other-vehicle",
            #                 'Boeing737', 'Boeing747', 'Boeing777', 'Boeing787', 'C919', 'A220', 'A321', 'A330', 'A350', 'ARJ21']:
            if cls_name not in ["Bridge"]:
                continue
            _image = cv2.imread(os.path.join(imgs_path, img_num + ".tif"))
            # _image = image.copy()
            width, hight = _image.shape[1], _image.shape[0]
            x_l = []
            y_l = []
            for i in range(0, len(points) - 1):  # 只遍历前四个点
                point = points[i]
                xy = point.text.split(",")
                x = int(float(xy[0]))
                y = int(float(xy[1]))
                bbox.append((x, y))
                x_l.append(x)
                y_l.append(y)
            x1 = min(x_l)
            y1 = min(y_l)
            x2 = max(x_l)
            y2 = max(y_l)
            # cv2.resizeWindow("Obj", (x2 - x1) * 1, (y2 - y1) * 1)
            if x1 == x2 or y1 == y2 or x1 < 0 or x2 < 0 or y1 < 0 or y2 < 0:
                continue
            # obj_img = _image[y1:y2, x1:x2]
            cv2.putText(_image, cls_name, (200, 300), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 0), 2)
            # cv2.rectangle(_image, (x1, y1), (x2, y2), (0, 0, 255), 1, 4)
            cv2.putText(_image, probability[0:4], (x2, y2), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 3)
            cv2.line(_image, bbox[0], bbox[1], (0, 0, 255), 1, 4)
            cv2.line(_image, bbox[1], bbox[2], (0, 0, 255), 1, 4)
            cv2.line(_image, bbox[2], bbox[3], (0, 0, 255), 1, 4)
            cv2.line(_image, bbox[3], bbox[0], (0, 0, 255), 1, 4)
            cv2.resizeWindow("FAIR1M", width, hight)
            cv2.imshow("FAIR1M", _image)
            # cv2.imshow("Obj", obj_img)
            print("当前目标的类别是：", cls_name, "，准确度是：", probability)
            print(x1, y1, x2, y2)
            print("宽：", (x2 - x1), ",高：", (y2 - y1))
            print("按键功能：回车->查看下一个目标；+->跳至下一个图片；")
            print("请输入想修改的类别（0~36）：")
            cv2.waitKey(10)
            inp = input()
            if inp == "":
                print("已选择不修改...........")
                continue
            if inp == "+":
                break
            lab_num = int(inp)
            if lab_num < 0 or lab_num > 36:
                print("输入不合法，此目标修改作废.......")
            else:
                possibleresult.find('name').text = classes[lab_num]
                print("已经将该目标修改为：", classes[lab_num])
            # cv2.destroyWindow("Obj")
        tree.write(os.path.join(labels_path, img_num + ".xml"), encoding="UTF-8")  # 更新xml文件

    # cv2.resizeWindow("FAIR1M", 800, 800)
    # cv2.imshow("FAIR1M", image)
    cv2.setMouseCallback("FAIR1M", mouse)
    cv2.waitKey(2500)

cv2.destroyAllWindows()