import os
import shutil
from pathlib import Path

def copy_png_images(source_folder, destination_folder):
    """
    从源文件夹递归搜索所有 .png 文件并复制到目标文件夹。

    参数:
    source_folder (str): 源文件夹路径
    destination_folder (str): 目标文件夹路径
    """
    # 将字符串路径转换为 Path 对象，更方便操作
    src_path = Path(source_folder)
    dest_path = Path(destination_folder)

    # 检查源文件夹是否存在
    if not src_path.exists():
        print(f"错误：源文件夹 '{source_folder}' 不存在。")
        return

    if not src_path.is_dir():
        print(f"错误：'{source_folder}' 不是一个有效的文件夹。")
        return

    # 如果目标文件夹不存在，则创建它
    dest_path.mkdir(parents=True, exist_ok=True)

    # 用于统计复制的文件数量
    copied_count = 0

    # 使用 rglob 递归搜索所有 .png 文件（不区分大小写）
    # rglob('*.[pP][nN][gG]') 可以匹配 .png, .PNG, .Png 等
    png_files = src_path.rglob('*.[jJ][pP][gG]')

    for png_file in png_files:
        try:
            # 计算目标文件的完整路径
            # 保持源文件夹的相对结构，但只保留文件名（可选：如果想保持结构，去掉 .name）
            # 这里我们只复制文件名到目标文件夹，避免路径过长或结构复杂
            dest_file = dest_path / png_file.name

            # 处理重名文件：在文件名后添加序号
            counter = 1
            original_dest_file = dest_file
            while dest_file.exists():
                dest_file = original_dest_file.parent / f"{original_dest_file.stem}_{counter}{original_dest_file.suffix}"
                counter += 1

            # 执行复制
            shutil.copy2(png_file, dest_file)  # copy2 会保留文件的元数据（如修改时间）
            print(f"已复制: {png_file} -> {dest_file}")
            copied_count += 1

        except Exception as e:
            print(f"复制文件时出错: {png_file} - {e}")

    print(f"\n复制完成！共复制了 {copied_count} 个 PNG 文件到 '{destination_folder}'。")

# ------------------ 主程序 ------------------
if __name__ == "__main__":
    # ====== 请在这里修改源文件夹和目标文件夹的路径 ======
    source_folder = r"/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/3"  # 替换为你的源文件夹路径
    destination_folder = r"/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/yolodataset/point"               # 替换为你的目标文件夹路径
    # ====================================================

    copy_png_images(source_folder, destination_folder)