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_xml_dimensions(data_dir):
    print(f"\n---------检查数据集 {data_dir} 目录下的所有xml文件是否正确------------")
    xml_directory = data_dir + ANNOTATIONS_PATH

    # 遍历指定目录下的所有文件
    for filename in os.listdir(xml_directory):
        if filename.endswith('.xml'):  # 确保只处理XML文件
            filepath = os.path.join(xml_directory, filename)

            # 解析XML文件
            tree = ET.parse(filepath)
            root = tree.getroot()

            # 查找标注中的size元素
            size_element = root.find('size')
            if size_element is not None:
                # 获取width和height属性
                width = int(size_element.find('width').text)
                height = int(size_element.find('height').text)

                # 检查width和height是否为0
                if width == 0 or height == 0:
                    print(f"File: {filepath}, Width: {width}, Height: {height}")
                    remove_xml_image_by_name(data_dir, os.path.splitext(filename)[0])


if __name__ == '__main__':
    current_path = os.getcwd()

    # 指定要检查的目录
    check_xml_dimensions(current_path)