import cv2
import os
import re
import argparse
from pathlib import Path

def natural_sort_key(filename):
    """提取文件名中的数字进行自然排序"""
    numbers = re.findall(r'\d+', filename)
    return [int(num) for num in numbers] if numbers else [0]

def create_video_from_images(input_folder, output_path="output_video.mp4", fps=30):
    """
    将文件夹中的图片按索引顺序合成视频
    
    参数:
    - input_folder: 输入图片文件夹路径
    - output_path: 输出视频文件路径
    - fps: 视频帧率
    """
    
    # 支持的图片格式
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
    
    # 获取所有图片文件
    image_files = []
    for file in os.listdir(input_folder):
        if Path(file).suffix.lower() in image_extensions:
            image_files.append(file)
    
    if not image_files:
        print("未找到支持的图片文件！")
        return False
    
    # 按照文件名中的数字索引排序
    image_files.sort(key=natural_sort_key)
    print(image_files)
    
    print(f"找到 {len(image_files)} 张图片")
    print(f"第一张图片: {image_files[0]}")
    print(f"最后一张图片: {image_files[-1]}")
    
    # 读取第一张图片获取尺寸
    first_image_path = os.path.join(input_folder, image_files[0])
    first_image = cv2.imread(first_image_path)
    
    if first_image is None:
        print(f"无法读取图片: {first_image_path}")
        return False
    
    height, width, layers = first_image.shape
    print(f"视频尺寸: {width}x{height}")
    
    # 创建视频编写器
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    video_writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    if not video_writer.isOpened():
        print("无法创建视频文件！")
        return False
    
    # 逐帧写入视频
    for i, image_file in enumerate(image_files):
        image_path = os.path.join(input_folder, image_file)
        image = cv2.imread(image_path)
        
        if image is None:
            print(f"跳过无效图片: {image_file}")
            continue
        
        # 确保图片尺寸一致
        if image.shape[:2] != (height, width):
            image = cv2.resize(image, (width, height))
        
        video_writer.write(image)
        
        # 显示进度
        if (i + 1) % 10 == 0 or i == len(image_files) - 1:
            print(f"处理进度: {i + 1}/{len(image_files)}")
    
    # 释放资源
    video_writer.release()
    cv2.destroyAllWindows()
    
    print(f"视频生成完成: {output_path}")
    print(f"视频参数: {width}x{height}, {fps}fps, {len(image_files)}帧")
    return True

def main():
    parser = argparse.ArgumentParser(description='将图片序列合成为视频')
    parser.add_argument('--input', '-i', default='.', help='输入图片文件夹路径 (默认: 当前目录)')
    parser.add_argument('--output', '-o', default='output_video.mp4', help='输出视频文件路径 (默认: output_video.mp4)')
    parser.add_argument('--fps', '-f', type=int, default=30, help='视频帧率 (默认: 30)')
    
    args = parser.parse_args()
    
    # 检查输入文件夹是否存在
    if not os.path.exists(args.input):
        print(f"错误: 输入文件夹不存在: {args.input}")
        return
    
    print(f"输入文件夹: {args.input}")
    print(f"输出文件: {args.output}")
    print(f"帧率: {args.fps}")
    print("=" * 50)
    
    # 创建视频
    success = create_video_from_images(args.input, args.output, args.fps)
    
    if success:
        print("视频生成成功！")
    else:
        print("视频生成失败！")

if __name__ == "__main__":
    main()