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)