import os
from PIL import Image
import xml.etree.ElementTree as ET
import re
import shutil

ANNOTATIONS_PATH = "/Annotations"
IMAGE_PATH = "/images"
LABEL_PATH = "/labels"

def remove_xml_image_by_name(data_dir, filename):
    xml_directory = data_dir + ANNOTATIONS_PATH
    images_directory = data_dir + IMAGE_PATH
    # 删除对应的xml和image
    xml_file_path = os.path.join(xml_directory, f'{filename}.xml')
    if os.path.exists(xml_file_path):  # 检查对应的XML文件是否存在
        print(f'正在删除对应的Annotations文件: {xml_file_path}')
        os.remove(xml_file_path)  # 删除对应的XML文件
    else:
        print(f"警告：对应的图像文件 {xml_file_path} 不存在，无法删除")

    image_file_path = os.path.join(images_directory, f'{filename}.jpg')
    if os.path.exists(image_file_path):  # 检查对应的XML文件是否存在
        print(f'正在删除对应的Image文件: {image_file_path}')
        os.remove(image_file_path)  # 删除对应的XML文件
    else:
        print(f"警告：对应的图像文件 {image_file_path} 不存在，无法删除")

def check_yolo_image_label_match(data_dir, image_extension=('.jpg', '.png', 'jpeg'), text_extension=('.txt')):

    images_directory = data_dir + IMAGE_PATH
    labels_directory = data_dir + LABEL_PATH

    if not os.path.isdir(images_directory) or not os.path.isdir(labels_directory):
        print("请确保指定的Images和labels目录存在。")
        return
    print(f"\n------------检查目录{images_directory} 和 {labels_directory}目录中文件的一致性------------")

    # 获取两个目录下的所有文件名（不包括路径）
    images = set(os.path.splitext(f)[0] for f in os.listdir(images_directory) if os.path.splitext(f)[1].lower() in image_extension)
    labels = set(os.path.splitext(f)[0] for f in os.listdir(labels_directory) if os.path.splitext(f)[1].lower() in text_extension)

    # 查找只有图片没有标签的文件
    missing_labels = images - labels
    if missing_labels:
        print("缺少对应的.txt标签文件的图片:")
        for img in missing_labels:
            print(img + image_extension[0])
            image_file_path = os.path.join(images_directory, img + image_extension[0])
            print(f'正在删除对应的image文件: {image_file_path}')
            os.remove(image_file_path)

    # # 查找只有标签没有图片的文件
    missing_images = labels - images
    if missing_images:
        print("缺少对应的图片文件的.txt标签:")
        for label in missing_images:
            print(label + text_extension)
            txt_file_path = os.path.join(labels_directory, label + text_extension)
            print(f'正在删除对应的txt文件: {txt_file_path}')
            os.remove(txt_file_path)

if __name__ == '__main__':
    current_path = os.getcwd()
    # 确保目录存在
    check_yolo_image_label_match(current_path)