import cv2
import numpy as np
import sys

def convert_to_red_pseudo_color(input_path, output_path):
    """
    将自然光视频转换为红色伪彩效果
    """
    # 打开视频
    cap = cv2.VideoCapture(input_path)
    if not cap.isOpened():
        print("❌ 无法打开视频文件，请检查路径是否正确")
        return

    # 获取视频参数
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    print(f"✅ 视频已加载，总帧数: {total_frames}, 分辨率: {width}x{height}, 帧率: {fps}")

    # 创建视频写入对象
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    frame_count = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # 转为灰度图
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # 纯红色伪彩
        red_pseudo = np.zeros_like(frame)
        red_pseudo[:, :, 2] = gray  # BGR 红色通道

        # 增强效果
        red_pseudo = cv2.convertScaleAbs(red_pseudo, alpha=1.6, beta=35)

        out.write(red_pseudo)
        frame_count += 1

        if frame_count % 100 == 0:
            print(f"处理进度: {frame_count}/{total_frames} ({frame_count/total_frames*100:.1f}%)")

    cap.release()
    out.release()
    print("\n🎉 转换完成！")

if __name__ == "__main__":
    # 支持命令行传参
    if len(sys.argv) != 3:
        print("使用方法：python 脚本名.py 输入视频路径 输出视频路径")
    else:
        input_video = sys.argv[1]
        output_video = sys.argv[2]
        convert_to_red_pseudo_color(input_video, output_video)