File size: 3,886 Bytes
aa24fe8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | import os
import xml.etree.ElementTree as ET
from shutil import move
from PIL import Image
def yolo_to_voc(yolo_label_folder, images_folder, voc_output_folder):
# 确保输出文件夹存在
if not os.path.exists(voc_output_folder):
os.makedirs(voc_output_folder)
# 获取标签文件和图片文件的列表
yolo_labels = [f for f in os.listdir(yolo_label_folder) if f.endswith('.txt')]
image_files = [f for f in os.listdir(images_folder) if f.endswith(('.jpg', '.png', '.jpeg'))]
# 创建一个集合存储图片文件名（不带扩展名）
image_files_set = {os.path.splitext(f)[0] for f in image_files}
for label_file in yolo_labels:
label_name = os.path.splitext(label_file)[0] # 获取没有扩展名的标签文件名
label_path = os.path.join(yolo_label_folder, label_file)
# 检查标签文件是否有对应的图片文件
if label_name not in image_files_set:
os.remove(label_path) # 删除没有对应图片的标签文件
print(f'Removed label file: {label_file}')
continue # 跳过这个标签文件
# 转换YOLO格式标签为VOC格式
voc_xml_path = os.path.join(voc_output_folder, label_name + '.xml')
with open(label_path, 'r') as f:
lines = f.readlines()
# 获取对应图片的尺寸
image_file_path = os.path.join(images_folder, label_name + '.jpg') # 假设图片是.jpg格式
if not os.path.exists(image_file_path): # 如果是PNG或其他格式，也可以扩展检查
image_file_path = os.path.join(images_folder, label_name + '.png')
image = Image.open(image_file_path)
width, height = image.size
# 创建VOC XML文件
annotation = ET.Element('annotation')
ET.SubElement(annotation, 'folder').text = images_folder
ET.SubElement(annotation, 'filename').text = label_name + '.jpg'
ET.SubElement(annotation, 'path').text = image_file_path
source = ET.SubElement(annotation, 'source')
ET.SubElement(source, 'database').text = 'Unknown'
size = ET.SubElement(annotation, 'size')
ET.SubElement(size, 'width').text = str(width)
ET.SubElement(size, 'height').text = str(height)
ET.SubElement(size, 'depth').text = '3' # 假设图片是RGB三通道
ET.SubElement(annotation, 'segmented').text = '0'
# 对应每个标签进行转换
for line in lines:
parts = line.strip().split()
class_id = int(parts[0]) # 类别ID
xmin = float(parts[1]) * width
ymin = float(parts[2]) * height
xmax = float(parts[3]) * width
ymax = float(parts[4]) * height
obj = ET.SubElement(annotation, 'object')
ET.SubElement(obj, 'name').text = str(class_id)
ET.SubElement(obj, 'pose').text = 'Unspecified'
ET.SubElement(obj, 'truncated').text = '0'
ET.SubElement(obj, 'difficult').text = '0'
bndbox = ET.SubElement(obj, 'bndbox')
ET.SubElement(bndbox, 'xmin').text = str(int(xmin))
ET.SubElement(bndbox, 'ymin').text = str(int(ymin))
ET.SubElement(bndbox, 'xmax').text = str(int(xmax))
ET.SubElement(bndbox, 'ymax').text = str(int(ymax))
# 保存VOC格式的XML文件
tree = ET.ElementTree(annotation)
tree.write(voc_xml_path)
print(f'Converted {label_file} to VOC format and saved as {voc_xml_path}')
# 调用函数，输入文件夹路径
yolo_label_folder = '/home/lab/LJ/wampee/WampeeDataSets/train/labels'
images_folder = '/home/lab/LJ/wampee/WampeeDataSets/train/images'
voc_output_folder = '/home/lab/LJ/wampee/Wampee_datasets_Voc'
yolo_to_voc(yolo_label_folder, images_folder, voc_output_folder)
|