import os
import shutil


def move_images_to_labels(image_dir, label_dir):
    """
    将标签对应的图片从图片文件夹移动到标签文件夹。

    :param image_dir: 存放图片的文件夹路径
    :param label_dir: 存放标签的文件夹路径，同时也是目标文件夹
    """
    # 支持的图片扩展名集合
    IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff', '.webp'}

    # 遍历标签文件夹中的所有文件
    for label_filename in os.listdir(label_dir):
        if not os.path.isfile(os.path.join(label_dir, label_filename)):
            continue  # 跳过非文件项（如目录）

        # 获取文件的基础名称（不带扩展名）
        base_name, _ = os.path.splitext(label_filename)

        # 在图片文件夹中查找对应的所有可能扩展名的图片文件
        found_image = False
        for ext in IMG_EXTENSIONS:
            image_filename = base_name + ext
            image_path = os.path.join(image_dir, image_filename)

            if os.path.exists(image_path):
                # 构造目标路径
                target_path = os.path.join(label_dir, image_filename)

                # 移动图片文件到标签文件夹
                print(f"移动 {image_path} 到 {target_path}")
                shutil.move(image_path, target_path)

                found_image = True
                break

        if not found_image:
            print(f"未找到与标签 {label_filename} 对应的图片")


if __name__ == "__main__":
    IMAGE_SOURCE_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/obb4/images"  # 替换为你的图片文件夹路径
    LABEL_TARGET_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/obb4/labels"  # 替换为你的标签文件夹路径

    if not os.path.isdir(IMAGE_SOURCE_DIR):
        print(f"图片文件夹不存在: {IMAGE_SOURCE_DIR}")
    elif not os.path.isdir(LABEL_TARGET_DIR):
        print(f"标签文件夹不存在: {LABEL_TARGET_DIR}")
    else:
        move_images_to_labels(IMAGE_SOURCE_DIR, LABEL_TARGET_DIR)