python yolov8半自动标注

首先标注一部分图片,进行训练,生成模型,标注文件为xml方便后面统一做处理。

1、标注数据(文件为xml, 转为txt用于训练,保留xml标签文件)

2、模型训练(训练配置、训练代码、)

3、使用训练好的模型进行预标注 (生成标注文件 xml)

4、检测标注文件工具:

单类别拆分、

合并所有类别xml、

合并指定几个类别、

固定矩形位置增加类别与固定位置绘制矩形检测是否重合删除类别(增删类别)、

固定矩形位置生成xml、

1、标注数据

注意:标注使用voc

这样标注文件是xml文件。

生成文件结构:放入当前目录下运行

python 复制代码
import os

# 创建文件夹结构
# 创建单层目录,如果目录存在则报错
path = "./VOCdevkit"
os.mkdir(path)
# 创建多级目录,如果目录存在则报错
p1 = os.path.join(path, "VOC2007")
os.makedirs(p1)
p2 = os.path.join(p1, "Annotations")
os.makedirs(p2)
p3 = os.path.join(p1, "YOLOLabels")
os.makedirs(p3)

xml转txt标注文件:

将xml文件放入:\VOCdevkit\VOC2007\Annotations 文件夹中

将图片放入:\VOCdevkit\VOC2007\JPEGImages 文件夹中

可指定训练与测试的百分比

cpp 复制代码
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import random
from shutil import copyfile

classes = ['clascc1', 'class2']  # 注意与自己的类对应,对应,对应,不然转好的txt文件是空的
# classes=["ball"]

TRAIN_RATIO = 80  # 按自己的要求划分,这里代表是train:test=8:2

# 创建文件夹结构
# 创建单层目录,如果目录存在则报错
path = "./VOCdevkit"
os.mkdir(path)
# 创建多级目录,如果目录存在则报错
p1 = os.path.join(path, "VOC2007")
os.makedirs(p1)
p2 = os.path.join(p1, "Annotations")
os.makedirs(p2)
p3 = os.path.join(p1, "YOLOLabels")
os.makedirs(p3)


def clear_hidden_files(path):
    dir_list = os.listdir(path)
    for i in dir_list:
        abspath = os.path.join(os.path.abspath(path), i)
        if os.path.isfile(abspath):
            if i.startswith("._"):
                os.remove(abspath)
        else:
            clear_hidden_files(abspath)


def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


def convert_annotation(image_id):
    in_file = open('VOCdevkit/VOC2007/Annotations/%s.xml' % image_id, encoding='utf-8')
    out_file = open('VOCdevkit/VOC2007/YOLOLabels/%s.txt' % image_id, 'w', encoding='utf-8')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        # difficult = obj.find('difficult').text
        cls = obj.find('name').text
        # if cls not in classes or int(difficult) == 1:
        #     continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
    in_file.close()
    out_file.close()


wd = os.getcwd()
data_base_dir = os.path.join(wd, "VOCdevkit/")
if not os.path.isdir(data_base_dir):
    os.mkdir(data_base_dir)
work_sapce_dir = os.path.join(data_base_dir, "VOC2007/")
if not os.path.isdir(work_sapce_dir):
    os.mkdir(work_sapce_dir)
annotation_dir = os.path.join(work_sapce_dir, "Annotations/")
if not os.path.isdir(annotation_dir):
    os.mkdir(annotation_dir)
clear_hidden_files(annotation_dir)
image_dir = os.path.join(work_sapce_dir, "JPEGImages/")
if not os.path.isdir(image_dir):
    os.mkdir(image_dir)
clear_hidden_files(image_dir)
yolo_labels_dir = os.path.join(work_sapce_dir, "YOLOLabels/")
if not os.path.isdir(yolo_labels_dir):
    os.mkdir(yolo_labels_dir)
clear_hidden_files(yolo_labels_dir)
yolov5_images_dir = os.path.join(data_base_dir, "images/")
if not os.path.isdir(yolov5_images_dir):
    os.mkdir(yolov5_images_dir)
clear_hidden_files(yolov5_images_dir)
yolov5_labels_dir = os.path.join(data_base_dir, "labels/")
if not os.path.isdir(yolov5_labels_dir):
    os.mkdir(yolov5_labels_dir)
clear_hidden_files(yolov5_labels_dir)
yolov5_images_train_dir = os.path.join(yolov5_images_dir, "train/")
if not os.path.isdir(yolov5_images_train_dir):
    os.mkdir(yolov5_images_train_dir)
clear_hidden_files(yolov5_images_train_dir)
yolov5_images_test_dir = os.path.join(yolov5_images_dir, "val/")
if not os.path.isdir(yolov5_images_test_dir):
    os.mkdir(yolov5_images_test_dir)
clear_hidden_files(yolov5_images_test_dir)
yolov5_labels_train_dir = os.path.join(yolov5_labels_dir, "train/")
if not os.path.isdir(yolov5_labels_train_dir):
    os.mkdir(yolov5_labels_train_dir)
clear_hidden_files(yolov5_labels_train_dir)
yolov5_labels_test_dir = os.path.join(yolov5_labels_dir, "val/")
if not os.path.isdir(yolov5_labels_test_dir):
    os.mkdir(yolov5_labels_test_dir)
clear_hidden_files(yolov5_labels_test_dir)

train_file = open(os.path.join(wd, "yolov5_train.txt"), 'w')
test_file = open(os.path.join(wd, "yolov5_val.txt"), 'w')
train_file.close()
test_file.close()
train_file = open(os.path.join(wd, "yolov5_train.txt"), 'a')
test_file = open(os.path.join(wd, "yolov5_val.txt"), 'a')
list_imgs = os.listdir(image_dir)  # list image files
prob = random.randint(1, 100)
print("Probability: %d" % prob)
for i in range(0, len(list_imgs)):
    path = os.path.join(image_dir, list_imgs[i])
    if os.path.isfile(path):
        image_path = image_dir + list_imgs[i]
        voc_path = list_imgs[i]
        (nameWithoutExtention, extention) = os.path.splitext(os.path.basename(image_path))
        (voc_nameWithoutExtention, voc_extention) = os.path.splitext(os.path.basename(voc_path))
        annotation_name = nameWithoutExtention + '.xml'
        annotation_path = os.path.join(annotation_dir, annotation_name)
        label_name = nameWithoutExtention + '.txt'
        label_path = os.path.join(yolo_labels_dir, label_name)
    prob = random.randint(1, 100)
    print("Probability: %d" % prob)
    if (prob < TRAIN_RATIO):  # train dataset
        if os.path.exists(annotation_path):
            train_file.write(image_path + '\n')
            convert_annotation(nameWithoutExtention)  # convert label
            copyfile(image_path, yolov5_images_train_dir + voc_path)
            copyfile(label_path, yolov5_labels_train_dir + label_name)
    else:  # test dataset
        if os.path.exists(annotation_path):
            test_file.write(image_path + '\n')
            convert_annotation(nameWithoutExtention)  # convert label
            copyfile(image_path, yolov5_images_test_dir + voc_path)
            copyfile(label_path, yolov5_labels_test_dir + label_name)
train_file.close()
test_file.close()

2、模型训练

将下面文件结构复制到训练的路径(服务器)

yol;ov8 数据集文件结构:当前文件结构可以直接用于训练

训练配置:

配置图片的路径

python 复制代码
train: /data/libowen/yolov8_1/ultralytics/data/VOCdevkit/images/train
val: /data/libowen/yolov8_1/ultralytics/data/VOCdevkit/images/val
# number of classes
nc: 2
# class names
names: ['class1', 'class2']

训练代码:

python 复制代码
from ultralytics import YOLO

# Load a model
# model = YOLO("yolov8n.yaml")  # build a new model from YAML
# 目标检测 n s m l x 
model = YOLO("yolov8m.pt")  # load a pretrained model (recommended for training)
# 图像分类
# model = YOLO("yolov8n-cls.pt")  # load a pretrained model (recommended for training)
# model = YOLO("dataset.yaml").load("yolov8n.pt")  # build from YAML and transfer weights

# Train the model
results = model.train(data="dataset.yaml", epochs=40, imgsz=640)  # 40次 输入图像缩放大小640
# results = model.train(data="D:/yolo_/mu_biao_gen_zong/data", epochs=40, imgsz=640)

3、预标注--生成标注文件 xml

参数:读取图片的路径、保存xml的路径、需要标注的类别、需要标注的文件后缀、模型的路径+名称、图像的分辨率信息、

cpp 复制代码
import cv2
from ultralytics import YOLO
import numpy as np
import time
import os
import xml.etree.ElementTree as ET
from xml.dom import minidom


def save_xml_add(root, lei_bie_name, xmin, ymin, xmax, ymax):
    object_elem1 = ET.SubElement(root, "object")
    ET.SubElement(object_elem1, "name").text = lei_bie_name
    bndbox1 = ET.SubElement(object_elem1, "bndbox")
    ET.SubElement(bndbox1, "xmin").text = str(xmin)
    ET.SubElement(bndbox1, "ymin").text = str(ymin)
    ET.SubElement(bndbox1, "xmax").text = str(xmax)
    ET.SubElement(bndbox1, "ymax").text = str(ymax)


def create_xml_annotation(image_path, coord1, xml_save_path, size_3):
    root = ET.Element("annotation")

    folder = os.path.dirname(image_path)
    ET.SubElement(root, "folder").text = folder

    filename = os.path.basename(image_path)
    ET.SubElement(root, "filename").text = filename

    size = ET.SubElement(root, "size")
    ET.SubElement(size, "width").text = size_3[0]
    ET.SubElement(size, "height").text = size_3[1]
    ET.SubElement(size, "depth").text = size_3[2]

    for i in coord1:
        save_xml_add(root, str(i[0]), str(i[1]), str(i[2]), str(i[3]), str(i[4]))

    with open(xml_save_path, 'w', encoding='utf-8') as xml_file:
        # 将 XML 元素树转换为字节串,编码为 utf-8
        rough_string = ET.tostring(root, 'utf-8')
        # 使用 minidom 模块解析生成的字节串,得到一个可操作的 XML 对象
        reparsed = minidom.parseString(rough_string)
        # 将重新解析后的 XML 对象转换为格式打印(pretty-print)的字符串形式,
        # 其中 indent="  "表示使用两个空格作为缩进
        string_ = reparsed.toprettyxml(indent="  ")
        xml_file.write(string_)
    print("已保存:", xml_save_path)


    

class Yolov8Detector():
    def __init__(self, mode_name):
        # 加载目标检测模型
        # self.model = YOLO(mode_name)
        self.model = YOLO(mode_name)
        self.model_cls_name = self.model.names

    def run(self, frame, image_size):
        # 判断图像类型是否正确
        if (isinstance(frame, np.ndarray)):
            boxes_list = []
            # 目标检测
            results = self.model.predict(frame, imgsz=image_size)
            # results = self.model(frame)
            result = results[0]
            # print(result, "这是")
            boxes = result.boxes.cpu()  # Boxes对象用于边界框输出
            for cls_, box in zip(boxes.cls, boxes):
                cls_1 = int(np.array(cls_))
                x1 = int(np.array(box.xyxy)[0][0])
                y1 = int(np.array(box.xyxy)[0][1])
                x2 = int(np.array(box.xyxy)[0][2])
                y2 = int(np.array(box.xyxy)[0][3])
                boxes_list.append((self.model_cls_name[cls_1], x1, y1, x2, y2))
            
            # 返回格式
            # [(类别1, x1, y1, x2, y2), (类别2, x1, y1, x2, y2)]
            return boxes_list
    
    def model_cls_name(self):
        # 获取模型内类别名
        return self.model_cls_name

# 加载检测模型
v8D = Yolov8Detector("zhayaoku_20240828.pt")


# 图片路径
img_path = "./img_lab_1000/img_1000/"
# 保存xml的路径
xml_save_path = "./img_lab_1000/img_xml/"


# 图片分辨率信息
size_list = ["1920", "1080", "3"]

# 需要标注的类别  # head_no_hat  head_with_hat
class_name_list_xml_save = ['helmet']

for img_name in os.listdir(img_path):

    if img_name.endswith(".jpg"):
        image_path = os.path.join(img_path, img_name)

        frame = cv2.imread(image_path)
        coord1 = v8D.run(frame, 640) 
        coord2 = []

        for i_ in coord1:
            for name_i in class_name_list_xml_save:
                if i_[0] == name_i:
                    coord2.append((i_[0], i_[1], i_[2], i_[3], i_[4]))
        # xml文件名
        xml_save_path_name = os.path.join(xml_save_path, os.path.splitext(img_name)[0] + ".xml")
        # 保存xml
        create_xml_annotation(image_path, coord2, xml_save_path_name, size_list)

检测标注文件工具:

单类别拆分

参数:全部类别名、xml读取路径、xml保存路径、

python 复制代码
import xml.etree.ElementTree as ET
import os
from xml.dom import minidom

# 全部类别
list_class = ['helmet', 'human_backward', 'human_forward',
              'closed_door', 'opened_door', 'covered_door',
              'head_with_hat', 'head_no_hat']
# xml标签路径 (绝对路径)
path_lab = "E:/zyk_lab/炸药库ataset240820/ce/lab/"

# 拆分保存xml的路径 (绝对路径)
save_xml_path = "E:/zyk_lab/炸药库ataset240820/ce/ce/"


for xml_name in os.listdir(path_lab):
    # xml_name = 'Camera12_20231001_30.xml'

    # 1. 读取XML文档
    tree = ET.parse(path_lab + xml_name)
    root = tree.getroot()

    # 存储 字典
    dict_class = {}
    for i in list_class:
        dict_class[i] = []
        # 创建单个文件夹
        folder_name = save_xml_path + "/" + i + "/"
        if not os.path.exists(folder_name):
            os.mkdir(folder_name)
        #     print(f"文件夹 '{folder_name}' 创建成功。")
        # else:
        #     print(f"文件夹 '{folder_name}' 已存在。")

    # for i, j in dict_class.items():
    #     print(i, j)

    size_find_0 = root.find("size")
    size_w = size_find_0.find("width")
    size_h = size_find_0.find("height")
    size_d = size_find_0.find("depth")
    #
    # print(size_w.text)
    # print(size_h.text)
    # print(size_d.text)

    # 分离文件名 与 文件后缀
    name_lab, xml_ = os.path.splitext(xml_name)

    folder_jpg = root.find("folder")
    # print(folder_jpg.text)

    path_jpg = root.find("path")
    # print(path_jpg.text)

    filename_jpg = root.find("filename")
    # print(filename_jpg.text)

    # 2. 查找 object 全部
    objects = root.findall('object')
    for object_find_0 in objects:
        # print('Tag:', child.tag)
        # print('Text:', child.text)
        # print('Attributes:', child.attrib)

        class_name = object_find_0.find("name")
        class_bndbox = object_find_0.find("bndbox")
        class_bndbox_xmin = class_bndbox.find("xmin")
        class_bndbox_ymin = class_bndbox.find("ymin")
        class_bndbox_xmax = class_bndbox.find("xmax")
        class_bndbox_ymax = class_bndbox.find("ymax")

        # print(class_name.text)
        # print(class_bndbox_xmin.text)
        # print(class_bndbox_ymin.text)
        # print(class_bndbox_xmax.text)
        # print(class_bndbox_ymax.text)

        dict_class[class_name.text].append(
            (class_name.text,
             class_bndbox_xmin.text,
             class_bndbox_ymin.text,
             class_bndbox_xmax.text,
             class_bndbox_ymax.text,
             )
        )

    for ob_class, ob_list in dict_class.items():

        # 创建根元素
        root = ET.Element("annotation")

        folder_save = ET.SubElement(root, "folder")
        folder_save.text = folder_jpg.text

        filename_jpg_save = ET.SubElement(root, "filename")
        filename_jpg_save.text = filename_jpg.text

        path_save_xml = ET.SubElement(root, "path")
        path_save_xml.text = path_jpg.text

        # 创建子元素
        size_save = ET.SubElement(root, "size")
        # 创建二级子元素 只需输入参数不同即可
        size_w_save = ET.SubElement(size_save, "width")
        size_w_save.text = size_w.text
        size_h_save = ET.SubElement(size_save, "height")
        size_h_save.text = size_h.text
        size_d_save = ET.SubElement(size_save, "depth")
        size_d_save.text = size_d.text

        for ob_list_i in ob_list:
            object_save = ET.SubElement(root, "object")
            name_save = ET.SubElement(object_save, "name")
            name_save.text = str(ob_list_i[0])

            bndbox_save = ET.SubElement(object_save, "bndbox")
            xmin_save = ET.SubElement(bndbox_save, "xmin")
            xmin_save.text = str(ob_list_i[1])

            ymin_save = ET.SubElement(bndbox_save, "ymin")
            ymin_save.text = str(ob_list_i[2])

            xmax_save = ET.SubElement(bndbox_save, "xmax")
            xmax_save.text = str(ob_list_i[3])

            ymax_save = ET.SubElement(bndbox_save, "ymax")
            ymax_save.text = str(ob_list_i[4])

        # 写入文件
        if len(dict_class[ob_class]) != 0:
            if ob_class in list_class:
                path_save_i = save_xml_path + "/" + ob_class + "/" + xml_name
                print(path_save_i)
                # tree.write(path_save_i, encoding="utf-8", xml_declaration=True)
                with open(path_save_i, 'w', encoding='utf-8') as xml_file:
                    # 将 XML 元素树转换为字节串,编码为 utf-8
                    rough_string = ET.tostring(root, 'utf-8')
                    # 使用 minidom 模块解析生成的字节串,得到一个可操作的 XML 对象
                    reparsed = minidom.parseString(rough_string)
                    # 将重新解析后的 XML 对象转换为格式打印(pretty-print)的字符串形式,
                    # 其中 indent="  "表示使用两个空格作为缩进
                    string_ = reparsed.toprettyxml(indent="  ")
                    xml_file.write(string_)

合并所有类别xml

参数:全部类别、读取图片路径、单类别拆分的总路径、合并保存的路径、

python 复制代码
import xml.etree.ElementTree as ET
import os
from xml.dom import minidom

list_class = ['helmet', 'human_backward', 'human_forward',
              'closed_door', 'opened_door', 'covered_door',
              'head_with_hat', 'head_no_hat']
# 图片路径
path_img = "E:/zyk_lab/炸药库ataset240820/ce/img/"
# 拆分的总路径
path_lab = "E:/zyk_lab/炸药库ataset240820/ce/ce/"
# 合并后保存的路径
path_lab_save = "E:/zyk_lab/炸药库ataset240820/ce/lab_ce/"

for img_name in os.listdir(path_img):
    # img_name = "Camera12_20231001_31"
    img_name = os.path.splitext(img_name)[0]

    dict_class = {}
    for i in list_class:
        dict_class[i] = []

    dict_class["width"] = 0
    dict_class["height"] = 0
    dict_class["depth"] = 0
    dict_class["folder"] = "null"
    dict_class["path"] = "null"
    dict_class["filename"] = "null"

    for file_1 in os.listdir(path_lab):
        path_i = os.path.join(path_lab, file_1)
        for xml_name in os.listdir(path_i):
            if img_name == os.path.splitext(xml_name)[0]:

                # 1. 读取XML文档
                xml_path = os.path.join(path_i, xml_name)
                tree = ET.parse(xml_path)
                root = tree.getroot()

                size_find_0 = root.find("size")
                size_w = size_find_0.find("width").text
                size_h = size_find_0.find("height").text
                size_d = size_find_0.find("depth").text

                folder_jpg = root.find("folder").text
                # print(folder_jpg.text)
                try:
                    path_jpg = root.find("path").text
                    # print(path_jpg.text)
                except BaseException:
                    path_jpg = "path_jpg"

                filename_jpg = root.find("filename").text
                # print(filename_jpg.text)

                dict_class["width"] = size_w
                dict_class["height"] = size_h
                dict_class["depth"] = size_d
                dict_class["folder"] = folder_jpg
                dict_class["path"] = path_jpg
                dict_class["filename"] = filename_jpg

                # 2. 查找 object 全部
                try:
                    objects = root.findall('object')
                    for object_find_0 in objects:
                        # print('Tag:', child.tag)
                        # print('Text:', child.text)
                        # print('Attributes:', child.attrib)

                        class_name = object_find_0.find("name")
                        class_bndbox = object_find_0.find("bndbox")
                        class_bndbox_xmin = class_bndbox.find("xmin")
                        class_bndbox_ymin = class_bndbox.find("ymin")
                        class_bndbox_xmax = class_bndbox.find("xmax")
                        class_bndbox_ymax = class_bndbox.find("ymax")

                        # print(class_name.text)
                        # print(class_bndbox_xmin.text)
                        # print(class_bndbox_ymin.text)
                        # print(class_bndbox_xmax.text)
                        # print(class_bndbox_ymax.text)

                        dict_class[class_name.text].append(
                            (class_name.text,
                             class_bndbox_xmin.text,
                             class_bndbox_ymin.text,
                             class_bndbox_xmax.text,
                             class_bndbox_ymax.text,
                             )
                        )
                except BaseException:
                    print("读取 object 失败")

    # 保存
    # 创建根元素
    root = ET.Element("annotation")

    folder_save = ET.SubElement(root, "folder")
    if dict_class["folder"] != "null":
        folder_save.text = dict_class["folder"]

    filename_jpg_save = ET.SubElement(root, "filename")
    if dict_class["filename"] != "null":
        filename_jpg_save.text = dict_class["filename"]

    path_save_xml = ET.SubElement(root, "path")
    if dict_class["path"] != "null":
        path_save_xml.text = dict_class["path"]

    # 创建子元素
    size_save = ET.SubElement(root, "size")
    # 创建二级子元素 只需输入参数不同即可
    size_w_save = ET.SubElement(size_save, "width")
    if dict_class["width"] != "null":
        size_w_save.text = dict_class["width"]
    size_h_save = ET.SubElement(size_save, "height")
    if dict_class["height"] != "null":
        size_h_save.text = dict_class["height"]
    size_d_save = ET.SubElement(size_save, "depth")
    if dict_class["depth"] != "null":
        size_d_save.text = dict_class["depth"]

    for ob_class, ob_list in dict_class.items():
        print(ob_class, ob_list)
        if ob_class in ["folder", "filename", "path", "size", "width", "height", "depth"]:
            continue
        for ob_list_i in ob_list:
            # print(ob_list_i)
            object_save = ET.SubElement(root, "object")
            name_save = ET.SubElement(object_save, "name")
            name_save.text = str(ob_list_i[0])

            bndbox_save = ET.SubElement(object_save, "bndbox")
            xmin_save = ET.SubElement(bndbox_save, "xmin")
            xmin_save.text = str(ob_list_i[1])

            ymin_save = ET.SubElement(bndbox_save, "ymin")
            ymin_save.text = str(ob_list_i[2])

            xmax_save = ET.SubElement(bndbox_save, "xmax")
            xmax_save.text = str(ob_list_i[3])

            ymax_save = ET.SubElement(bndbox_save, "ymax")
            ymax_save.text = str(ob_list_i[4])

    # 写入文件
    # if len(dict_class[ob_class]) != 0:
    #     if ob_class in list_class:
    path_save_i = path_lab_save + img_name + ".xml"
    print(path_save_i)

    # tree.write(path_save_i, encoding="utf-8", xml_declaration=True)
    with open(path_save_i, 'w', encoding='utf-8') as xml_file:
        # 将 XML 元素树转换为字节串,编码为 utf-8
        rough_string = ET.tostring(root, 'utf-8')
        # 使用 minidom 模块解析生成的字节串,得到一个可操作的 XML 对象
        reparsed = minidom.parseString(rough_string)
        # 将重新解析后的 XML 对象转换为格式打印(pretty-print)的字符串形式,
        # 其中 indent="  "表示使用两个空格作为缩进
        string_ = reparsed.toprettyxml(indent="  ")
        xml_file.write(string_)

合并指定几个类别

注意:合并后检测完毕后,需要到原本拆分成单类别目录,将检测前的xml文件删除,把检查后的xml放入其中一个类别文件夹即可。

参数:全部类别、需要合并的类别、读取图像路径、拆分总目录、保存xml路径、

cpp 复制代码
import xml.etree.ElementTree as ET
import os
from xml.dom import minidom

# 全部类别
# list_class = ['class1', 'class2', 'class3']
list_class = [
              'class1', 'class2'
              ]

# 图像路径
path_lab = "E:/zyk_lab/ataset240820/ce/img/"

# 拆分后的总路径
xml_path_all = "E:/zyk_lab/ataset240820/ce/ce/"
# xml 保存的路径
xml_save_path = "E:/zyk_lab/ataset240820/ce/xml_ce/"


for xml_name in os.listdir(path_lab):
    # xml_name = 'Camera12_20231001_30.xml'

    # 存储 字典
    dict_class = {}
    for i in list_class:
        dict_class[i] = []

    dict_class["width"] = 0
    dict_class["height"] = 0
    dict_class["depth"] = 0
    dict_class["folder"] = "null"
    dict_class["path"] = "null"
    dict_class["filename"] = "null"

    # 读取 拆分后的 单类别路径
    for list_class_i in list_class:
        class_path = os.path.join(xml_path_all, list_class_i)
        # 拆分.jpg后缀
        xml_name = os.path.splitext(xml_name)[0] + ".xml"
        path_xml_name = os.path.join(class_path, xml_name)
        # 检测文件是否存在
        if os.path.exists(path_xml_name):

            # 1. 读取XML文档
            tree = ET.parse(path_xml_name)
            root = tree.getroot()

            size_find_0 = root.find("size")
            size_w = size_find_0.find("width")
            size_h = size_find_0.find("height")
            size_d = size_find_0.find("depth")

            # print(size_w.text)
            # print(size_h.text)
            # print(size_d.text)


            folder_jpg = root.find("folder")
            # print(folder_jpg.text)
            try:
                path_jpg = root.find("path").text
                # print(path_jpg.text)
            except BaseException:
                path_jpg = "path_jpg"

            filename_jpg = root.find("filename")
            # print(filename_jpg.text)

            dict_class["width"] = size_w.text
            dict_class["height"] = size_h.text
            dict_class["depth"] = size_d.text
            dict_class["folder"] = folder_jpg.text
            dict_class["path"] = path_jpg
            dict_class["filename"] = filename_jpg.text

            # 2. 查找 object 全部
            try:
                objects = root.findall('object')
                for object_find_0 in objects:
                    # print('Tag:', child.tag)
                    # print('Text:', child.text)
                    # print('Attributes:', child.attrib)

                    class_name = object_find_0.find("name")
                    class_bndbox = object_find_0.find("bndbox")
                    class_bndbox_xmin = class_bndbox.find("xmin")
                    class_bndbox_ymin = class_bndbox.find("ymin")
                    class_bndbox_xmax = class_bndbox.find("xmax")
                    class_bndbox_ymax = class_bndbox.find("ymax")

                    # print(class_name.text)
                    # print(class_bndbox_xmin.text)
                    # print(class_bndbox_ymin.text)
                    # print(class_bndbox_xmax.text)
                    # print(class_bndbox_ymax.text)

                    if class_name.text in list_class:

                        dict_class[class_name.text].append(
                            (class_name.text,
                             class_bndbox_xmin.text,
                             class_bndbox_ymin.text,
                             class_bndbox_xmax.text,
                             class_bndbox_ymax.text,
                             )
                        )
            except BaseException:
                print("读取 object 失败")

            # for i, j in dict_class.items():
            #     print(i, j)
            # 创建根元素
            root = ET.Element("annotation")

            folder_save = ET.SubElement(root, "folder")
            folder_save.text = dict_class["folder"]

            filename_jpg_save = ET.SubElement(root, "filename")
            filename_jpg_save.text = dict_class["filename"]

            path_save_xml = ET.SubElement(root, "path")
            path_save_xml.text = dict_class["filename"]

            # 创建子元素
            size_save = ET.SubElement(root, "size")
            # 创建二级子元素 只需输入参数不同即可
            size_w_save = ET.SubElement(size_save, "width")
            size_w_save.text = dict_class["width"]
            size_h_save = ET.SubElement(size_save, "height")
            size_h_save.text = dict_class["height"]
            size_d_save = ET.SubElement(size_save, "depth")
            size_d_save.text = dict_class["depth"]

            for class_i in list_class:
                for ob_list_i in dict_class[class_i]:
                    object_save = ET.SubElement(root, "object")
                    name_save = ET.SubElement(object_save, "name")
                    name_save.text = str(ob_list_i[0])

                    bndbox_save = ET.SubElement(object_save, "bndbox")
                    xmin_save = ET.SubElement(bndbox_save, "xmin")
                    xmin_save.text = str(ob_list_i[1])

                    ymin_save = ET.SubElement(bndbox_save, "ymin")
                    ymin_save.text = str(ob_list_i[2])

                    xmax_save = ET.SubElement(bndbox_save, "xmax")
                    xmax_save.text = str(ob_list_i[3])

                    ymax_save = ET.SubElement(bndbox_save, "ymax")
                    ymax_save.text = str(ob_list_i[4])

                    # 写入文件
                    path_save_i = os.path.join(xml_save_path, xml_name)
                    print(path_save_i)
                    # tree.write(path_save_i, encoding="utf-8", xml_declaration=True)
                    with open(path_save_i, 'w', encoding='utf-8') as xml_file:
                        # 将 XML 元素树转换为字节串,编码为 utf-8
                        rough_string = ET.tostring(root, 'utf-8')
                        # 使用 minidom 模块解析生成的字节串,得到一个可操作的 XML 对象
                        reparsed = minidom.parseString(rough_string)
                        # 将重新解析后的 XML 对象转换为格式打印(pretty-print)的字符串形式,
                        # 其中 indent="  "表示使用两个空格作为缩进
                        string_ = reparsed.toprettyxml(indent="  ")
                        xml_file.write(string_)

增删类别

注意:不会新建xml

在xml文件内,想在固定位置绘制一个矩形框作为标签。

在xml文件内,指定左上角与右下角坐标检查某类别的标注矩形是否有重合,有则删除。

python 复制代码
import xml.etree.ElementTree as ET
import os
from xml.dom import minidom


# 矩形重合检测

def is_rectangles_overlap(x1, y1, x2, y2, x1_, y1_, x2_, y2_):
    stat = False

    num_x2 = x2_ - x1_
    for i in range(num_x2):
        x1_i = x1_ + i
        # print(x1_i)

        if (x1 < x1_i < x2) and (y1 < y1_ < y2):
            # print("x上面线内的点有重合")
            stat = True
            break
        if (x1 < x1_i < x2) and (y1 < y2_ < y2):
            # print("x下面线内的点有重合")
            stat = True
            break

    num_y1 = y2_ - y1_
    for j in range(num_y1):
        y1_i = y1_ + j
        if (x1 < x1_ < x2) and (y1 < y1_i < y2):
            # print("y左边线内有重合")
            stat = True
            break
        if (x1 < x2_ < x2) and (y1 < y1_i < y2):
            # print("y右边线内有重合")
            stat = True
            break

    # 矩形完全重合
    if x1 == x1_ and y1 == y1_ and x2 == x2_ and y2 == y2_:
        stat = True

    return stat


def save_xml_add(name, xmin, ymin, xmax, ymax):
    object_save = ET.SubElement(root, "object")
    name_save = ET.SubElement(object_save, "name")
    name_save.text = str(name)

    bndbox_save = ET.SubElement(object_save, "bndbox")
    xmin_save = ET.SubElement(bndbox_save, "xmin")
    xmin_save.text = str(xmin)

    ymin_save = ET.SubElement(bndbox_save, "ymin")
    ymin_save.text = str(ymin)

    xmax_save = ET.SubElement(bndbox_save, "xmax")
    xmax_save.text = str(xmax)

    ymax_save = ET.SubElement(bndbox_save, "ymax")
    ymax_save.text = str(ymax)


# 添加标注的信息
# 根据分辨率的不同设置不同类别和坐标
dict_add = {}
# 704 只是为了提醒自己在多大分辨率的图上绘制
# dict_add = {704: [("covered_door", 184, 33, 331, 330), ]
#             }


# 删除标注信息
# dict_del = [("closed_door", 184, 33, 331, 330)]
# dle 不管什么类别 只要重合就删除
dict_del = [("dle", 184, 33, 331, 330)]
# dict_del = []

# 读取路径 拆分后单类别的路径
path_i = r"E:\zyk_lab\炸药库ataset240820\ce\ce\closed_door"

# 保存路径 自己新建路径
path_lab_save = r"E:\zyk_lab\炸药库ataset240820\ce\xml_ce"

# xml_name = "Camera12_20231001_30.xml"
for xml_name in os.listdir(path_i):
    # 存储字典 全部类别
    list_class = ['helmet', 'human_backward', 'human_forward',
                  'closed_door', 'opened_door', 'covered_door',
                  'head_with_hat', 'head_no_hat']
    dict_class = {}
    for i in list_class:
        dict_class[i] = []

    # 1. 读取XML文档
    xml_path = os.path.join(path_i, xml_name)
    print(xml_path)
    tree = ET.parse(xml_path)
    root = tree.getroot()

    size_find_0 = root.find("size")
    size_w = size_find_0.find("width").text
    size_h = size_find_0.find("height").text
    size_d = size_find_0.find("depth").text

    folder_jpg = root.find("folder").text
    # print(folder_jpg.text)
    path_jpg = root.find("path").text
    # print(path_jpg.text)
    filename_jpg = root.find("filename").text
    # print(filename_jpg.text)

    dict_class["width"] = size_w
    dict_class["height"] = size_h
    dict_class["depth"] = size_d
    dict_class["folder"] = folder_jpg
    dict_class["path"] = path_jpg
    dict_class["filename"] = filename_jpg

    # 2. 查找 object 全部
    objects = root.findall('object')
    for object_find_0 in objects:

        # print('Tag:', child.tag)
        # print('Text:', child.text)
        # print('Attributes:', child.attrib)

        class_name = object_find_0.find("name")
        class_bndbox = object_find_0.find("bndbox")
        class_bndbox_xmin = class_bndbox.find("xmin")
        class_bndbox_ymin = class_bndbox.find("ymin")
        class_bndbox_xmax = class_bndbox.find("xmax")
        class_bndbox_ymax = class_bndbox.find("ymax")

        # print(class_name.text)
        # print(class_bndbox_xmin.text)
        # print(class_bndbox_ymin.text)
        # print(class_bndbox_xmax.text)
        # print(class_bndbox_ymax.text)

        con_1 = False

        if dict_del:
            for list_con in dict_del:

                x1_, y1_, x2_, y2_ = int(class_bndbox_xmin.text), int(class_bndbox_ymin.text), int(
                    class_bndbox_xmax.text), int(class_bndbox_ymax.text)
                # print(f"{list_con[1]}, {list_con[2]}, {list_con[3]}, {list_con[4]},{x1_}, {y1_}, {x2_}, {y2_}")

                # 制定类别
                if class_name == list_con[0] or list_con[0] == "dle":
                    # 框若有重合 则跳过
                    if is_rectangles_overlap(list_con[1], list_con[2], list_con[3], list_con[4], x1_, y1_, x2_, y2_):
                        print("有重合")
                        print(f"{list_con[1]}, {list_con[2]}, {list_con[3]}, {list_con[4]},{x1_}, {y1_}, {x2_}, {y2_}")
                        con_1 = True
                    else:
                        print("无重合")
                        print(f"{list_con[1]}, {list_con[2]}, {list_con[3]}, {list_con[4]},{x1_}, {y1_}, {x2_}, {y2_}")
        # print(con_1)
        if con_1:
            continue

        dict_class[class_name.text].append(
            (class_name.text,
             class_bndbox_xmin.text,
             class_bndbox_ymin.text,
             class_bndbox_xmax.text,
             class_bndbox_ymax.text,
             )
        )

    # 保存
    # 创建根元素
    root = ET.Element("annotation")

    folder_save = ET.SubElement(root, "folder")
    if dict_class["folder"] != "null":
        folder_save.text = dict_class["folder"]

    filename_jpg_save = ET.SubElement(root, "filename")
    if dict_class["filename"] != "null":
        filename_jpg_save.text = dict_class["filename"]

    path_save_xml = ET.SubElement(root, "path")
    if dict_class["path"] != "null":
        path_save_xml.text = dict_class["path"]

    # 创建子元素
    size_save = ET.SubElement(root, "size")
    # 创建二级子元素 只需输入参数不同即可
    size_w_save = ET.SubElement(size_save, "width")
    if dict_class["width"] != "null":
        size_w_save.text = dict_class["width"]
    size_h_save = ET.SubElement(size_save, "height")
    if dict_class["height"] != "null":
        size_h_save.text = dict_class["height"]
    size_d_save = ET.SubElement(size_save, "depth")
    if dict_class["depth"] != "null":
        size_d_save.text = dict_class["depth"]

    for ob_class, ob_list in dict_class.items():
        # print(ob_class, ob_list)
        if ob_class in ["folder", "filename", "path", "size", "width", "height", "depth"]:
            continue
        for ob_list_i in ob_list:
            # print(ob_list_i)
            save_xml_add(ob_list_i[0], ob_list_i[1], ob_list_i[2], ob_list_i[3], ob_list_i[4])

    # 读取字典内容 添加新的标注
    if dict_add:
        for img_size, list_class_xyxy in dict_add.items():
            for class_xyxy_i in list_class_xyxy:
                save_xml_add(class_xyxy_i[0], class_xyxy_i[1], class_xyxy_i[2], class_xyxy_i[3], class_xyxy_i[4])

    # 写入文件
    # if len(dict_class[ob_class]) != 0:
    #     if ob_class in list_class:
    path_save_i = os.path.join(path_lab_save, xml_name)
    # print(path_save_i)

    # tree.write(path_save_i, encoding="utf-8", xml_declaration=True)
    with open(path_save_i, 'w', encoding='utf-8') as xml_file:
        # 将 XML 元素树转换为字节串,编码为 utf-8
        rough_string = ET.tostring(root, 'utf-8')
        # 使用 minidom 模块解析生成的字节串,得到一个可操作的 XML 对象
        reparsed = minidom.parseString(rough_string)
        # 将重新解析后的 XML 对象转换为格式打印(pretty-print)的字符串形式,
        # 其中 indent="  "表示使用两个空格作为缩进
        string_ = reparsed.toprettyxml(indent="  ")
        xml_file.write(string_)

固定矩形位置生成xml

参数:

cpp 复制代码
import os
import xml.etree.ElementTree as ET
from xml.dom import minidom


def save_xml_add(root, lei_bie_name, xmin, ymin, xmax, ymax):
    object_elem1 = ET.SubElement(root, "object")
    ET.SubElement(object_elem1, "name").text = lei_bie_name
    bndbox1 = ET.SubElement(object_elem1, "bndbox")
    ET.SubElement(bndbox1, "xmin").text = str(xmin)
    ET.SubElement(bndbox1, "ymin").text = str(ymin)
    ET.SubElement(bndbox1, "xmax").text = str(xmax)
    ET.SubElement(bndbox1, "ymax").text = str(ymax)


def create_xml_annotation(image_path, coord1, xml_save_path, size_3, lei_bie_name):
    root = ET.Element("annotation")

    folder = os.path.dirname(image_path)
    ET.SubElement(root, "folder").text = folder

    filename = os.path.basename(image_path)
    ET.SubElement(root, "filename").text = filename

    size = ET.SubElement(root, "size")
    ET.SubElement(size, "width").text = size_3[0]
    ET.SubElement(size, "height").text = size_3[1]
    ET.SubElement(size, "depth").text = size_3[2]

    for i in coord1:
        save_xml_add(root, lei_bie_name, i[0], i[1], i[2], i[3])

    with open(xml_save_path, 'w', encoding='utf-8') as xml_file:
        # 将 XML 元素树转换为字节串,编码为 utf-8
        rough_string = ET.tostring(root, 'utf-8')
        # 使用 minidom 模块解析生成的字节串,得到一个可操作的 XML 对象
        reparsed = minidom.parseString(rough_string)
        # 将重新解析后的 XML 对象转换为格式打印(pretty-print)的字符串形式,
        # 其中 indent="  "表示使用两个空格作为缩进
        string_ = reparsed.toprettyxml(indent="  ")
        xml_file.write(string_)
    print("已保存:", xml_save_path)


def xml_save(folder_path, xml_save_path, coord1, size_list, lei_bie_name):
    for filename in os.listdir(folder_path):
        if filename.endswith(".jpg"):
            image_path = os.path.join(folder_path, filename)
            xml_save_path_name = os.path.join(xml_save_path, os.path.splitext(filename)[0] + ".xml")

            create_xml_annotation(image_path, coord1, xml_save_path_name, size_list, lei_bie_name)


# 图片路径
folder_path = r"E:\zyk_lab\ataset240820\ce\img"
# 保存xml的路径
xml_save_path = r"E:\zyk_lab\ataset240820\ce\xml_ce"
# 绘制坐标
coord1 = [(221, 47, 330, 325)]
# 绘制类别
lei_bie_name = "covered_door"
# 图片分辨率信息
size_list = ["704", "576", "3"]

xml_save(folder_path, xml_save_path, coord1, size_list, lei_bie_name)
相关推荐
禁默几秒前
深入浅出:AWT的基本组件及其应用
java·开发语言·界面编程
Cachel wood7 分钟前
python round四舍五入和decimal库精确四舍五入
java·linux·前端·数据库·vue.js·python·前端框架
Code哈哈笑9 分钟前
【Java 学习】深度剖析Java多态:从向上转型到向下转型,解锁动态绑定的奥秘,让代码更优雅灵活
java·开发语言·学习
終不似少年遊*12 分钟前
pyecharts
python·信息可视化·数据分析·学习笔记·pyecharts·使用技巧
程序猿进阶13 分钟前
深入解析 Spring WebFlux:原理与应用
java·开发语言·后端·spring·面试·架构·springboot
Python之栈14 分钟前
【无标题】
数据库·python·mysql
qq_4336184415 分钟前
shell 编程(二)
开发语言·bash·shell
charlie11451419129 分钟前
C++ STL CookBook
开发语言·c++·stl·c++20
袁袁袁袁满30 分钟前
100天精通Python(爬虫篇)——第113天:‌爬虫基础模块之urllib详细教程大全
开发语言·爬虫·python·网络爬虫·爬虫实战·urllib·urllib模块教程
ELI_He99936 分钟前
PHP中替换某个包或某个类
开发语言·php