首先标注一部分图片，进行训练，生成模型，标注文件为xml方便后面统一做处理。
1、标注数据（文件为xml, 转为txt用于训练，保留xml标签文件）
2、模型训练（训练配置、训练代码、）
3、使用训练好的模型进行预标注 （生成标注文件 xml）
4、检测标注文件工具：
单类别拆分、
合并所有类别xml、
合并指定几个类别、
固定矩形位置增加类别与固定位置绘制矩形检测是否重合删除类别（增删类别）、
固定矩形位置生成xml、
1、标注数据
注意：标注使用voc
这样标注文件是xml文件。
生成文件结构：放入当前目录下运行
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 文件夹中
可指定训练与测试的百分比
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 数据集文件结构：当前文件结构可以直接用于训练
训练配置：
配置图片的路径
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']
训练代码：
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的路径、需要标注的类别、需要标注的文件后缀、模型的路径+名称、图像的分辨率信息、
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保存路径、
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
参数：全部类别、读取图片路径、单类别拆分的总路径、合并保存的路径、
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路径、
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文件内，指定左上角与右下角坐标检查某类别的标注矩形是否有重合，有则删除。
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
参数：
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)
