我要评论按xml中提取的标签顺序转索引
import os
import random
import time
from pathlib import path
import shutil
import tkinter as tk
from tkinter import filedialog
from loguru import logger
import xml.etree.elementtree as et
class analysisxml(object):
'''清洗xml'''
def __init__(self):
root = tk.tk()
root.withdraw()
root.attributes('-topmost', 1)
self.directory = filedialog.askdirectory() # 打开目录选择器
root.destroy()
logger.warning(f'路径选择:【{self.directory}】')
def xml_img_split(self):
'''分割图片和xml'''
logger.info(f'---------------------------分割图片和xml------------------------')
self.images_path = path(self.directory).parent.joinpath('images')
self.xml_labels_path = path(self.directory).parent.joinpath('xml_labels')
self.images_path.mkdir(parents=true, exist_ok=true)
self.xml_labels_path.mkdir(parents=true, exist_ok=true)
for i in path(self.directory).iterdir():
if i.suffix == '.xml':
new_path = self.xml_labels_path.joinpath(i.name)
logger.debug(f'移动:【{i}】 -> 【{new_path}】')
shutil.copy(str(i), str(new_path))
if i.suffix in ('.jpg', '.png'):
new_path = self.images_path.joinpath(i.name)
logger.debug(f'移动:【{i}】 -> 【{new_path}】')
shutil.copy(str(i), str(new_path))
def xml_to_txt(self):
'''xml转txt'''
logger.info(f'----------------------------正在将xml转为txt-----------------------')
self.txt_labels = self.xml_labels_path.joinpath('labels') # 替换为实际的输出txt文件夹路径
os.makedirs(self.txt_labels, exist_ok=true)
names_set = set()
for filename in os.listdir(self.xml_labels_path):
if filename.endswith('.xml'):
tree = et.parse(os.path.join(self.xml_labels_path, filename))
root = tree.getroot()
for obj in root.findall('object'):
name = obj.find('name').text
names_set.add(name)
# 输出所有的name
categories = []
for name in names_set:
categories.append(name)
logger.success(f'标注的内容names:【{categories}】')
category_to_index = {category: index for index, category in enumerate(categories)}
# 遍历输入文件夹中的所有xml文件
for filename in os.listdir(self.xml_labels_path):
if filename.endswith('.xml'):
xml_path = os.path.join(self.xml_labels_path, filename)
logger.warning(f'正在处理:【{xml_path}】')
# 解析xml文件
tree = et.parse(xml_path)
root = tree.getroot()
# 提取图像的尺寸
size = root.find('size')
width = int(size.find('width').text)
height = int(size.find('height').text)
# 存储name和对应的归一化坐标
objects = []
# 遍历xml中的object标签
for obj in root.findall('object'):
name = obj.find('name').text
if name in category_to_index:
category_index = category_to_index[name]
else:
continue # 如果name不在指定类别中，跳过该object
bndbox = obj.find('bndbox')
xmin = int(bndbox.find('xmin').text)
ymin = int(bndbox.find('ymin').text)
xmax = int(bndbox.find('xmax').text)
ymax = int(bndbox.find('ymax').text)
# 转换为中心点坐标和宽高
x_center = (xmin + xmax) / 2.0
y_center = (ymin + ymax) / 2.0
w = xmax - xmin
h = ymax - ymin
# 归一化
x = x_center / width
y = y_center / height
w = w / width
h = h / height
objects.append(f"{category_index} {x:.6f} {y:.6f} {w:.6f} {h:.6f}")
# 输出结果到对应的txt文件
txt_filename = os.path.splitext(filename)[0] + '.txt'
txt_path = os.path.join(self.txt_labels, txt_filename)
with open(txt_path, 'w') as f:
for obj in objects:
f.write(obj + '\n')
def to_dataset(self, test_ratio):
'''整理为dataset'''
output_folder = os.path.join(os.path.dirname(self.directory), 'datasets')
input_image_folder = self.images_path
input_label_folder = self.txt_labels
train_images_folder = os.path.join(output_folder, 'train', 'images')
train_labels_folder = os.path.join(output_folder, 'train', 'labels')
val_images_folder = os.path.join(output_folder, 'val', 'images')
val_labels_folder = os.path.join(output_folder, 'val', 'labels')
os.makedirs(train_images_folder, exist_ok=true)
os.makedirs(train_labels_folder, exist_ok=true)
os.makedirs(val_images_folder, exist_ok=true)
os.makedirs(val_labels_folder, exist_ok=true)
# 获取所有图像文件列表
images = [f for f in os.listdir(input_image_folder) if f.endswith('.jpg') or f.endswith('.png')]
# 随机打乱图像文件列表
random.shuffle(images)
# 计算验证集的数量
val_size = int(len(images) * test_ratio)
# 划分验证集和训练集
val_images = images[:val_size]
train_images = images[val_size:]
# 复制验证集图像和标签
for image in val_images:
label = os.path.splitext(image)[0] + '.txt'
if os.path.exists(os.path.join(input_label_folder, label)):
shutil.copy(os.path.join(input_image_folder, image), os.path.join(val_images_folder, image))
shutil.copy(os.path.join(input_label_folder, label), os.path.join(val_labels_folder, label))
logger.debug(
f'【{os.path.join(input_image_folder, image)}】 --> 【{os.path.join(val_images_folder, image)}】')
logger.success(
f'【{os.path.join(input_label_folder, label)}】 --> 【{os.path.join(val_labels_folder, label)}】')
else:
logger.error(f"warning: label file {label} not found for image {image}")
# 复制训练集图像和标签
for image in train_images:
label = os.path.splitext(image)[0] + '.txt'
if os.path.exists(os.path.join(input_label_folder, label)):
shutil.copy(os.path.join(input_image_folder, image), os.path.join(train_images_folder, image))
shutil.copy(os.path.join(input_label_folder, label), os.path.join(train_labels_folder, label))
logger.debug(
f'【{os.path.join(input_image_folder, image)}】 --> 【{os.path.join(train_images_folder, image)}】')
logger.success(
f'【{os.path.join(input_label_folder, label)}】 --> 【{os.path.join(train_labels_folder, label)}】')
else:
logger.error(f"warning: label file {label} not found for image {image}")
def start(self):
'''启动'''
time.sleep(1)
self.xml_img_split()
time.sleep(1)
self.xml_to_txt()
time.sleep(1)
self.to_dataset(0.2)
if __name__ == '__main__':
base_dir = os.path.dirname(__file__)
log_path = os.path.join(base_dir, 'log.log')
if os.path.exists(log_path):
os.unlink(log_path)
logger.add(log_path)
print('...第一层文件夹')
print(' -->第二层文件夹↓')
print(' -->[xml和img混合文件夹]')
print('\n')
status = input('请确认xml和图片在同一个文件夹(99:确认)(任意值:取消):')
if status in (99, '99'):
a = analysisxml()
a.start()
logger.success('系统完成')
for i in (3, 2, 1):
time.sleep(1)
logger.success(f'{i}/秒')
else:
logger.error('系统退出!')
for i in (3, 2, 1):
time.sleep(1)
logger.error(f'{i}/秒')
根据自定义标签转格式(推荐)
import os
import random
import time
from pathlib import path
import shutil
import tkinter as tk
from tkinter import filedialog, simpledialog
from loguru import logger
import xml.etree.elementtree as et
class analysisxml(object):
'''清洗xml并按自定义标签列表转换为yolo格式'''
def __init__(self):
root = tk.tk()
root.withdraw()
root.attributes('-topmost', 1)
self.directory = filedialog.askdirectory() # 打开目录选择器
# 新增：输入自定义标签列表（用英文逗号分隔，如 cat,dog,car）
self.custom_categories = simpledialog.askstring(
"请输入",
'自定义标签列表,(空格逗号分隔); 输入示例: cat dog car person',
parent=root
)
root.destroy()
# 处理输入的标签列表
if self.custom_categories:
# 去空格、拆分、去重
self.categories = [cat.strip() for cat in self.custom_categories.split() if cat.strip()]
logger.warning(f'自定义标签列表:【{self.categories}】')
else:
logger.error('未输入标签列表，程序退出！')
exit(1)
logger.warning(f'路径选择:【{self.directory}】')
def xml_img_split(self):
'''分割图片和xml'''
logger.info(f'---------------------------分割图片和xml------------------------')
self.images_path = path(self.directory).parent.joinpath('images')
self.xml_labels_path = path(self.directory).parent.joinpath('xml_labels')
self.images_path.mkdir(parents=true, exist_ok=true)
self.xml_labels_path.mkdir(parents=true, exist_ok=true)
for i in path(self.directory).iterdir():
if i.suffix == '.xml':
new_path = self.xml_labels_path.joinpath(i.name)
logger.debug(f'移动:【{i}】 -> 【{new_path}】')
shutil.copy(str(i), str(new_path))
if i.suffix in ('.jpg', '.png'):
new_path = self.images_path.joinpath(i.name)
logger.debug(f'移动:【{i}】 -> 【{new_path}】')
shutil.copy(str(i), str(new_path))
def xml_to_txt(self):
'''xml转txt（按自定义标签列表映射索引）'''
logger.info(f'----------------------------正在将xml转为txt-----------------------')
self.txt_labels = self.xml_labels_path.joinpath('labels') # 替换为实际的输出txt文件夹路径
os.makedirs(self.txt_labels, exist_ok=true)
# 核心修改：用自定义标签列表生成索引映射，而非自动无序生成
category_to_index = {category: index for index, category in enumerate(self.categories)}
logger.success(f'标签->索引映射:【{category_to_index}】')
# 遍历输入文件夹中的所有xml文件
for filename in os.listdir(self.xml_labels_path):
if filename.endswith('.xml'):
xml_path = os.path.join(self.xml_labels_path, filename)
logger.warning(f'正在处理:【{xml_path}】')
# 解析xml文件
tree = et.parse(xml_path)
root = tree.getroot()
# 提取图像的尺寸
size = root.find('size')
width = int(size.find('width').text)
height = int(size.find('height').text)
# 存储name和对应的归一化坐标
objects = []
# 遍历xml中的object标签
for obj in root.findall('object'):
name = obj.find('name').text
# 检查标签是否在自定义列表中，不在则跳过并警告
if name not in category_to_index:
logger.warning(f'xml【{filename}】中发现未定义标签【{name}】，已跳过！')
continue
category_index = category_to_index[name]
bndbox = obj.find('bndbox')
xmin = int(bndbox.find('xmin').text)
ymin = int(bndbox.find('ymin').text)
xmax = int(bndbox.find('xmax').text)
ymax = int(bndbox.find('ymax').text)
# 转换为中心点坐标和宽高
x_center = (xmin + xmax) / 2.0
y_center = (ymin + ymax) / 2.0
w = xmax - xmin
h = ymax - ymin
# 归一化
x = x_center / width
y = y_center / height
w = w / width
h = h / height
objects.append(f"{category_index} {x:.6f} {y:.6f} {w:.6f} {h:.6f}")
# 输出结果到对应的txt文件
txt_filename = os.path.splitext(filename)[0] + '.txt'
txt_path = os.path.join(self.txt_labels, txt_filename)
with open(txt_path, 'w') as f:
for obj in objects:
f.write(obj + '\n')
def to_dataset(self, test_ratio):
'''整理为dataset'''
output_folder = os.path.join(os.path.dirname(self.directory), 'datasets')
input_image_folder = self.images_path
input_label_folder = self.txt_labels
train_images_folder = os.path.join(output_folder, 'train', 'images')
train_labels_folder = os.path.join(output_folder, 'train', 'labels')
val_images_folder = os.path.join(output_folder, 'val', 'images')
val_labels_folder = os.path.join(output_folder, 'val', 'labels')
os.makedirs(train_images_folder, exist_ok=true)
os.makedirs(train_labels_folder, exist_ok=true)
os.makedirs(val_images_folder, exist_ok=true)
os.makedirs(val_labels_folder, exist_ok=true)
# 获取所有图像文件列表
images = [f for f in os.listdir(input_image_folder) if f.endswith('.jpg') or f.endswith('.png')]
# 随机打乱图像文件列表
random.shuffle(images)
# 计算验证集的数量
val_size = int(len(images) * test_ratio)
# 划分验证集和训练集
val_images = images[:val_size]
train_images = images[val_size:]
# 复制验证集图像和标签
for image in val_images:
label = os.path.splitext(image)[0] + '.txt'
if os.path.exists(os.path.join(input_label_folder, label)):
shutil.copy(os.path.join(input_image_folder, image), os.path.join(val_images_folder, image))
shutil.copy(os.path.join(input_label_folder, label), os.path.join(val_labels_folder, label))
logger.debug(
f'【{os.path.join(input_image_folder, image)}】 --> 【{os.path.join(val_images_folder, image)}】')
logger.success(
f'【{os.path.join(input_label_folder, label)}】 --> 【{os.path.join(val_labels_folder, label)}】')
else:
logger.error(f"warning: label file {label} not found for image {image}")
# 复制训练集图像和标签
for image in train_images:
label = os.path.splitext(image)[0] + '.txt'
if os.path.exists(os.path.join(input_label_folder, label)):
shutil.copy(os.path.join(input_image_folder, image), os.path.join(train_images_folder, image))
shutil.copy(os.path.join(input_label_folder, label), os.path.join(train_labels_folder, label))
logger.debug(
f'【{os.path.join(input_image_folder, image)}】 --> 【{os.path.join(train_images_folder, image)}】')
logger.success(
f'【{os.path.join(input_label_folder, label)}】 --> 【{os.path.join(train_labels_folder, label)}】')
else:
logger.error(f"warning: label file {label} not found for image {image}")
def start(self):
'''启动'''
time.sleep(1)
self.xml_img_split()
time.sleep(1)
self.xml_to_txt()
time.sleep(1)
self.to_dataset(0.2)
if __name__ == '__main__':
base_dir = os.path.dirname(__file__)
log_path = os.path.join(base_dir, 'log.log')
if os.path.exists(log_path):
os.unlink(log_path)
logger.add(log_path)
print('...第一层文件夹')
print(' -->第二层文件夹↓')
print(' -->[xml和img混合文件夹]')
print('\n')
status = input('请确认xml和图片在同一个文件夹(99:确认)(任意值:取消):')
if status in (99, '99'):
a = analysisxml()
a.start()
logger.success('系统完成')
for i in (3, 2, 1):
time.sleep(1)
logger.success(f'{i}/秒')
else:
logger.error('系统退出!')
for i in (3, 2, 1):
time.sleep(1)
logger.error(f'{i}/秒')
回显yolo标注文件
import cv2
import os
def draw_yolo_annotation(image_path, txt_path, save_path=none):
"""
在图片上绘制yolo标注框和类别索引
:param image_path: 原图路径（如 test.jpg）
:param txt_path: 对应的yolo标注txt文件路径（如 test.txt）
:param save_path: 绘制后的图片保存路径，none则直接显示
"""
# 1. 读取图片
img = cv2.imread(image_path)
if img is none:
print(f"错误：无法读取图片 {image_path}")
return
h, w = img.shape[:2] # 获取图片高、宽
# 2. 读取并解析yolo标注文件
if not os.path.exists(txt_path):
print(f"警告：标注文件 {txt_path} 不存在，直接返回原图")
if save_path:
cv2.imwrite(save_path, img)
else:
cv2.imshow("no annotation", img)
cv2.waitkey(0)
cv2.destroyallwindows()
return
with open(txt_path, 'r', encoding='utf-8') as f:
lines = f.readlines()
# 3. 遍历每个标注框，转换坐标并绘制
for line in lines:
line = line.strip()
if not line:
continue
# 解析标注行：类别索引、cx、cy、bw、bh
parts = line.split()
if len(parts) < 5:
print(f"无效标注行：{line}")
continue
class_idx = int(parts[0])
cx = float(parts[1]) * w # 归一化→像素坐标
cy = float(parts[2]) * h
bw = float(parts[3]) * w
bh = float(parts[4]) * h
# 计算框的左上角、右下角坐标（yolo中心点→opencv矩形框）
x1 = int(cx - bw / 2)
y1 = int(cy - bh / 2)
x2 = int(cx + bw / 2)
y2 = int(cy + bh / 2)
# 4. 绘制矩形框+类别索引文本
# 框的样式：绿色（bgr）、线宽2
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
# 文本背景（黑色半透明），避免文字被遮挡
text = f"class: {class_idx}"
text_size = cv2.gettextsize(text, cv2.font_hershey_simplex, 0.6, 2)[0]
text_x = x1
text_y = y1 - 10 if y1 - 10 > 10 else y1 + 20 # 避免文字超出图片
cv2.rectangle(img, (text_x, text_y - text_size[1] - 5),
(text_x + text_size[0] + 5, text_y + 5), (0, 0, 0), -1)
# 绘制文本：白色、字体大小0.6、线宽2
cv2.puttext(img, text, (text_x + 2, text_y),
cv2.font_hershey_simplex, 0.6, (255, 255, 255), 2)
# 5. 保存/显示结果
if save_path:
# 确保保存目录存在
os.makedirs(os.path.dirname(save_path), exist_ok=true)
cv2.imwrite(save_path, img)
print(f"标注可视化完成，保存至：{save_path}")
else:
img = cv2.resize(img, (1500, 800))
cv2.imshow("yolo annotation", img)
cv2.waitkey(0) # 按任意键关闭窗口
cv2.destroyallwindows()
# ===================== 测试调用 =====================
if __name__ == "__main__":
# 替换为你的图片和标注文件路径
image_path = r"c:\users\123\desktop\22\控制室东_2026-01-16 15-22-23.249.jpg" # 原图路径
txt_path = r"c:\users\123\desktop\22\控制室东_2026-01-16 15-22-23.249.txt" # 对应的yolo标注txt
# save_path = "test_anno.jpg" # 绘制后的保存路径（可选）
# 方式1：直接显示标注后的图片
draw_yolo_annotation(image_path, txt_path)
# 方式2：保存标注后的图片（不显示）
# draw_yolo_annotation(image_path, txt_path, save_path)
总结
以上为个人经验，希望能给大家一个参考，也希望大家多多支持代码网。
发表评论