import os
import cv2
import numpy as np

def process_single_video(video_path):
    """处理单个视频文件，计算帧平均并保存为PNG"""
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"无法打开视频文件: {video_path}")
        return False
    
    frame_count = 0
    accumulated_image = None
    
    print(f"正在处理视频: {os.path.basename(video_path)}...", end=' ', flush=True)
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        # 转换为浮点型以便累加
        frame_float = frame.astype(np.float64)
        
        if accumulated_image is None:
            accumulated_image = np.zeros_like(frame_float)
        
        accumulated_image += frame_float
        frame_count += 1
    
    cap.release()
    
    if frame_count == 0:
        print(f"视频中没有有效帧: {video_path}")
        return False
    
    # 计算平均值并转换回8位图像
    averaged_image = (accumulated_image / frame_count).astype(np.uint8)
    
    # 生成输出文件名 (PNG格式)
    base_name = os.path.splitext(video_path)[0]
    output_filename = f"{base_name}_average.png"
    
    # 保存为PNG文件
    if not cv2.imwrite(output_filename, averaged_image):
        print(f"保存失败: {output_filename}")
        return False
    
    print(f"完成 -> {os.path.basename(output_filename)} (共{frame_count}帧)")
    return True

def process_all_avi_files():
    """处理当前目录下所有AVI文件"""
    # 获取当前目录下所有AVI文件
    video_files = sorted([f for f in os.listdir('.') 
                         if f.lower().endswith('.avi')])
    
    if not video_files:
        print("当前目录下未找到AVI视频文件")
        return
    
    print(f"找到 {len(video_files)} 个AVI文件:")
    processed_count = 0
    
    for video_file in video_files:
        if process_single_video(video_file):
            processed_count += 1
    
    print(f"\n处理完成，成功处理 {processed_count}/{len(video_files)} 个文件")

if __name__ == "__main__":
    print("===== 视频帧平均处理程序 =====")
    print("将当前目录下所有AVI文件的帧平均并保存为PNG")
    print("----------------------------")
    process_all_avi_files()
    input("\n按Enter键退出...")