#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
图片帧重新组织脚本
将extracted_frames_test文件夹中的所有图片按连续序号重命名并合并到一个目录
"""

import os
import shutil
from pathlib import Path
import re

def get_frame_number(filename):
    """从文件名中提取帧序号"""
    match = re.search(r'frame_(\d+)', filename)
    return int(match.group(1)) if match else 0

def collect_all_frames(source_dir):
    """收集所有子文件夹中的图片文件"""
    all_frames = []
    source_path = Path(source_dir)
    
    # 获取所有子文件夹并按名称排序
    subdirs = sorted([d for d in source_path.iterdir() if d.is_dir()])
    
    for subdir in subdirs:
        print(f"处理文件夹: {subdir.name}")
        
        # 获取该文件夹中的所有jpg文件
        jpg_files = list(subdir.glob('*.jpg'))
        
        # 按帧序号排序
        jpg_files.sort(key=lambda x: get_frame_number(x.name))
        
        for jpg_file in jpg_files:
            all_frames.append({
                'original_path': jpg_file,
                'folder_name': subdir.name,
                'original_frame_num': get_frame_number(jpg_file.name)
            })
        
        print(f"  找到 {len(jpg_files)} 个图片文件")
    
    return all_frames

def reorganize_frames(source_dir, output_dir):
    """重新组织和重命名图片文件"""
    # 创建输出目录
    output_path = Path(output_dir)
    output_path.mkdir(exist_ok=True)
    
    # 收集所有图片文件
    all_frames = collect_all_frames(source_dir)
    
    print(f"\n总共找到 {len(all_frames)} 个图片文件")
    print("开始重命名和复制...")
    
    # 按新的连续序号重命名并复制
    for new_index, frame_info in enumerate(all_frames):
        # 生成新的文件名
        new_filename = f"frame_{new_index:06d}.jpg"
        new_path = output_path / new_filename
        
        # 复制文件
        shutil.copy2(frame_info['original_path'], new_path)
        
        if new_index % 100 == 0:  # 每100个文件显示一次进度
            print(f"  已处理: {new_index + 1}/{len(all_frames)}")
    
    print(f"\n完成！所有文件已重命名并复制到: {output_dir}")
    print(f"文件命名范围: frame_000000.jpg 到 frame_{len(all_frames)-1:06d}.jpg")
    
    return len(all_frames)

def main():
    """主函数"""
    # 设置路径
    source_dir = r"d:\data\extracted_frames_test"
    output_dir = r"d:\data\reorganized_frames"
    
    print("图片帧重新组织脚本")
    print(f"源目录: {source_dir}")
    print(f"输出目录: {output_dir}")
    print("-" * 50)
    
    # 检查源目录是否存在
    if not os.path.exists(source_dir):
        print(f"错误: 源目录不存在 - {source_dir}")
        return
    
    try:
        # 执行重组织
        total_files = reorganize_frames(source_dir, output_dir)
        print(f"\n成功处理了 {total_files} 个图片文件")
        
    except Exception as e:
        print(f"错误: {e}")
        return

if __name__ == "__main__":
    main()