#!/usr/bin/env python3
"""
Seamless图片收集程序

这个程序的功能是：
1. 从images文件夹的所有子文件夹中收集以seamless开头的图片
2. 从exps文件夹的子文件夹中的seamless文件夹收集图片
3. 按模型名称（如bag-000）整理到seamless_collection文件夹中

使用方法：
1. 修改下面的路径配置
2. 运行程序：python seamless_collector.py
"""

import os
import shutil
from pathlib import Path


def find_seamless_images_in_folder(folder_path):
    """在指定文件夹中查找所有以seamless开头的图片"""
    seamless_images = []
    if not os.path.exists(folder_path):
        return seamless_images

    for item in os.listdir(folder_path):
        item_path = os.path.join(folder_path, item)
        if os.path.isfile(item_path) and item.lower().startswith('seamless'):
            # 检查是否为图片文件
            if item.lower().endswith(('.png', '.jpg', '.jpeg')):
                seamless_images.append(item_path)

    return seamless_images


def collect_seamless_images(images_dir, exps_dir, output_base_dir):
    """收集所有seamless图片"""
    print("开始收集seamless图片...")

    # 创建输出目录
    seamless_output_dir = os.path.join(output_base_dir, "seamless_collection")
    os.makedirs(seamless_output_dir, exist_ok=True)

    # 1. 从images文件夹收集seamless图片
    print("从images文件夹收集seamless图片...")
    images_seamless_count = 0

    if os.path.exists(images_dir):
        for subfolder in os.listdir(images_dir):
            subfolder_path = os.path.join(images_dir, subfolder)
            if os.path.isdir(subfolder_path):
                print(f"  检查文件夹: {subfolder}")
                seamless_images = find_seamless_images_in_folder(subfolder_path)

                if seamless_images:
                    # 创建对应的输出文件夹
                    output_subfolder = os.path.join(seamless_output_dir, subfolder)
                    os.makedirs(output_subfolder, exist_ok=True)

                    for seamless_image in seamless_images:
                        filename = os.path.basename(seamless_image)
                        output_path = os.path.join(output_subfolder, filename)
                        shutil.copy2(seamless_image, output_path)
                        print(f"    复制: {filename}")
                        images_seamless_count += 1
                else:
                    print(f"    未找到seamless图片")
    else:
        print(f"  警告：images文件夹不存在: {images_dir}")

    # 2. 从exps文件夹收集seamless图片
    print("从exps文件夹收集seamless图片...")
    exps_seamless_count = 0

    if os.path.exists(exps_dir):
        for subfolder in os.listdir(exps_dir):
            subfolder_path = os.path.join(exps_dir, subfolder)
            if os.path.isdir(subfolder_path):
                # 检查是否有seamless子文件夹
                seamless_subfolder = os.path.join(subfolder_path, "seamless")
                if os.path.exists(seamless_subfolder):
                    print(f"  检查文件夹: {subfolder}/seamless")
                    seamless_images = find_seamless_images_in_folder(seamless_subfolder)

                    if seamless_images:
                        # 提取基础名称（如bag-000）
                        base_name = subfolder.split('-')[0] if '-' in subfolder else subfolder

                        # 创建对应的输出文件夹
                        output_subfolder = os.path.join(seamless_output_dir, base_name)
                        os.makedirs(output_subfolder, exist_ok=True)

                        for seamless_image in seamless_images:
                            filename = os.path.basename(seamless_image)
                            # 添加前缀以区分来源
                            new_filename = f"exps_{filename}"
                            output_path = os.path.join(output_subfolder, new_filename)
                            shutil.copy2(seamless_image, output_path)
                            print(f"    复制: {new_filename}")
                            exps_seamless_count += 1
                    else:
                        print(f"    未找到seamless图片")
    else:
        print(f"  警告：exps文件夹不存在: {exps_dir}")

    print(f"\nseamless图片收集完成！")
    print(f"从images文件夹收集: {images_seamless_count} 个文件")
    print(f"从exps文件夹收集: {exps_seamless_count} 个文件")
    print(f"总收集: {images_seamless_count + exps_seamless_count} 个文件")
    print(f"输出目录: {seamless_output_dir}")

    return seamless_output_dir


def main():
    # ==================== 路径配置 ====================
    # 请根据您的实际情况修改这些路径

    # images文件夹：包含bag-000等子文件夹，每个子文件夹可能包含seamless图片
    images_dir = "/home/swu/szp/MaterialTransfer/assets/images"

    # exps文件夹：包含bag-000-0等子文件夹，每个子文件夹可能包含seamless文件夹
    exps_dir = "/home/swu/szp/MaterialTransfer/data_backup/exps"

    # 输出文件夹：收集的seamless图片将保存在这里
    output_dir = "/home/swu/szp/MaterialTransfer/data_backup/seamless_all"

    # ==================== 路径配置结束 ====================

    print("Seamless图片收集程序")
    print("=" * 50)
    print(f"Images文件夹: {images_dir}")
    print(f"Exps文件夹: {exps_dir}")
    print(f"输出文件夹: {output_dir}")
    print("=" * 50)

    # 检查路径是否存在
    if not os.path.exists(images_dir):
        print(f"错误：images文件夹不存在: {images_dir}")
        return

    if not os.path.exists(exps_dir):
        print(f"错误：exps文件夹不存在: {exps_dir}")
        return

    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    # 开始收集seamless图片
    collect_seamless_images(images_dir, exps_dir, output_dir)


if __name__ == "__main__":
    main()