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

    for root, dirs, files in os.walk(images_directory):
        for file in files:
            if file.endswith('.jpg') or file.endswith('.jpeg'):
                file_path = os.path.join(root, file)
                try:
                    img = Image.open(file_path)  # Open the image file
                    img.verify()  # Verify that it's a valid image
                except (IOError, SyntaxError) as e:
                    print(f'发现损坏的图片: {file_path}')
                    remove_xml_image_by_name(data_dir, os.path.splitext(file)[0])

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

    # 检查图片文件
    check_images(current_path)