155 lines
5.6 KiB
Python
155 lines
5.6 KiB
Python
|
import os
|
|||
|
|
import xml.etree.ElementTree as ET
|
|||
|
|
from xml.dom import minidom
|
|||
|
|
from PIL import Image
|
|||
|
|
|
|||
|
|
def convert_txt_to_xml(txt_path, image_path, xml_path):
|
|||
|
|
# 打开图像，获取尺寸及通道数
|
|||
|
|
with Image.open(image_path) as img:
|
|||
|
|
img_width, img_height = img.size
|
|||
|
|
mode = img.mode
|
|||
|
|
if mode == "RGB":
|
|||
|
|
depth = 3
|
|||
|
|
elif mode == "L":
|
|||
|
|
depth = 1
|
|||
|
|
else:
|
|||
|
|
depth = 3 # 默认3
|
|||
|
|
|
|||
|
|
# 读取TXT文件中的标注信息
|
|||
|
|
with open(txt_path, "r") as f:
|
|||
|
|
lines = f.readlines()
|
|||
|
|
|
|||
|
|
# 类别编号到标签名称的映射
|
|||
|
|
id_to_label = {
|
|||
|
|
0: "light",
|
|||
|
|
1: "red",
|
|||
|
|
2: "yellow",
|
|||
|
|
3: "green"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 构建XML结构
|
|||
|
|
annotation = ET.Element("annotation", verified="yes")
|
|||
|
|
|
|||
|
|
folder_elem = ET.SubElement(annotation, "folder")
|
|||
|
|
folder_elem.text = os.path.basename(os.path.dirname(image_path))
|
|||
|
|
|
|||
|
|
filename_elem = ET.SubElement(annotation, "filename")
|
|||
|
|
image_filename = os.path.basename(image_path)
|
|||
|
|
filename_elem.text = image_filename
|
|||
|
|
|
|||
|
|
path_elem = ET.SubElement(annotation, "path")
|
|||
|
|
path_elem.text = os.path.abspath(image_path)
|
|||
|
|
|
|||
|
|
source_elem = ET.SubElement(annotation, "source")
|
|||
|
|
database_elem = ET.SubElement(source_elem, "database")
|
|||
|
|
database_elem.text = "Unknown"
|
|||
|
|
|
|||
|
|
size_elem = ET.SubElement(annotation, "size")
|
|||
|
|
width_elem = ET.SubElement(size_elem, "width")
|
|||
|
|
width_elem.text = str(img_width)
|
|||
|
|
height_elem = ET.SubElement(size_elem, "height")
|
|||
|
|
height_elem.text = str(img_height)
|
|||
|
|
depth_elem = ET.SubElement(size_elem, "depth")
|
|||
|
|
depth_elem.text = str(depth)
|
|||
|
|
|
|||
|
|
segmented_elem = ET.SubElement(annotation, "segmented")
|
|||
|
|
segmented_elem.text = "0"
|
|||
|
|
|
|||
|
|
# 逐行处理TXT标注
|
|||
|
|
for line in lines:
|
|||
|
|
parts = line.strip().split()
|
|||
|
|
if len(parts) != 5:
|
|||
|
|
print(f"跳过格式不正确的行: {line.strip()}")
|
|||
|
|
continue
|
|||
|
|
class_id_str, x_center_norm_str, y_center_norm_str, width_norm_str, height_norm_str = parts
|
|||
|
|
try:
|
|||
|
|
class_id = int(class_id_str)
|
|||
|
|
x_center_norm = float(x_center_norm_str)
|
|||
|
|
y_center_norm = float(y_center_norm_str)
|
|||
|
|
width_norm = float(width_norm_str)
|
|||
|
|
height_norm = float(height_norm_str)
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"转换数据出错 {line.strip()}: {e}")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
if class_id not in id_to_label:
|
|||
|
|
print(f"未知标签编号 {class_id}，跳过该行")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
label = id_to_label[class_id]
|
|||
|
|
|
|||
|
|
# 计算绝对坐标
|
|||
|
|
x_center = x_center_norm * img_width
|
|||
|
|
y_center = y_center_norm * img_height
|
|||
|
|
bbox_width = width_norm * img_width
|
|||
|
|
bbox_height = height_norm * img_height
|
|||
|
|
|
|||
|
|
xmin = int(x_center - bbox_width / 2)
|
|||
|
|
ymin = int(y_center - bbox_height / 2)
|
|||
|
|
xmax = int(x_center + bbox_width / 2)
|
|||
|
|
ymax = int(y_center + bbox_height / 2)
|
|||
|
|
|
|||
|
|
# 构建object节点
|
|||
|
|
object_elem = ET.SubElement(annotation, "object")
|
|||
|
|
name_elem = ET.SubElement(object_elem, "name")
|
|||
|
|
name_elem.text = label
|
|||
|
|
pose_elem = ET.SubElement(object_elem, "pose")
|
|||
|
|
pose_elem.text = "Unspecified"
|
|||
|
|
truncated_elem = ET.SubElement(object_elem, "truncated")
|
|||
|
|
truncated_elem.text = "0"
|
|||
|
|
difficult_elem = ET.SubElement(object_elem, "difficult")
|
|||
|
|
difficult_elem.text = "0"
|
|||
|
|
|
|||
|
|
bndbox_elem = ET.SubElement(object_elem, "bndbox")
|
|||
|
|
xmin_elem = ET.SubElement(bndbox_elem, "xmin")
|
|||
|
|
xmin_elem.text = str(xmin)
|
|||
|
|
ymin_elem = ET.SubElement(bndbox_elem, "ymin")
|
|||
|
|
ymin_elem.text = str(ymin)
|
|||
|
|
xmax_elem = ET.SubElement(bndbox_elem, "xmax")
|
|||
|
|
xmax_elem.text = str(xmax)
|
|||
|
|
ymax_elem = ET.SubElement(bndbox_elem, "ymax")
|
|||
|
|
ymax_elem.text = str(ymax)
|
|||
|
|
|
|||
|
|
# 美化输出的XML字符串
|
|||
|
|
xml_str = ET.tostring(annotation, encoding="utf-8")
|
|||
|
|
parsed_str = minidom.parseString(xml_str)
|
|||
|
|
pretty_xml_str = parsed_str.toprettyxml(indent=" ")
|
|||
|
|
|
|||
|
|
with open(xml_path, "w", encoding="utf-8") as f:
|
|||
|
|
f.write(pretty_xml_str)
|
|||
|
|
|
|||
|
|
def batch_convert_txt_to_xml(txt_dir, image_dir, output_dir, image_ext=".jpg"):
|
|||
|
|
"""
|
|||
|
|
批量将TXT转换为XML：
|
|||
|
|
txt_dir: TXT文件所在目录
|
|||
|
|
image_dir: 对应图像所在目录（图片文件名与TXT同名，后缀可指定）
|
|||
|
|
output_dir: XML文件保存目录
|
|||
|
|
image_ext: 图像文件扩展名（默认 .jpg）
|
|||
|
|
"""
|
|||
|
|
if not os.path.exists(output_dir):
|
|||
|
|
os.makedirs(output_dir)
|
|||
|
|
|
|||
|
|
for file in os.listdir(txt_dir):
|
|||
|
|
if file.lower().endswith(".txt"):
|
|||
|
|
txt_path = os.path.join(txt_dir, file)
|
|||
|
|
base_name = os.path.splitext(file)[0]
|
|||
|
|
image_filename = base_name + image_ext
|
|||
|
|
image_path = os.path.join(image_dir, image_filename)
|
|||
|
|
if not os.path.exists(image_path):
|
|||
|
|
print(f"图像文件不存在: {image_path}，跳过 {txt_path}")
|
|||
|
|
continue
|
|||
|
|
xml_filename = base_name + ".xml"
|
|||
|
|
xml_path = os.path.join(output_dir, xml_filename)
|
|||
|
|
try:
|
|||
|
|
convert_txt_to_xml(txt_path, image_path, xml_path)
|
|||
|
|
print(f"成功转换: {txt_path} --> {xml_path}")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"处理 {txt_path} 时出错: {e}")
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
# 请修改以下目录为实际路径
|
|||
|
|
txt_directory = "path_to_txt_directory" # TXT文件所在目录
|
|||
|
|
image_directory = "path_to_image_directory" # 图像文件所在目录
|
|||
|
|
output_directory = "path_to_output_xml_directory" # 输出XML文件保存目录
|
|||
|
|
batch_convert_txt_to_xml(txt_directory, image_directory, output_directory, image_ext=".jpg")
|