Voidljc

Your commit message

aa24fe8 over 1 year ago

3.89 kB

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)

	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)