import cv2
import numpy as np


def run_optical_flow(video_path):
    # 1. 打开视频
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"错误: 无法打开视频 {video_path}")
        return

    # 2. 读取第一帧并转换为灰度
    ret, first_frame = cap.read()
    if not ret:
        print("错误: 视频无法读取")
        return
    prev_gray = cv2.cvtColor(first_frame, cv2.COLOR_BGR2GRAY)

    # 3. 设置可视化参数
    step = 16  # 步长：不需要每个像素都画箭头，每隔 16 个像素画一个
    scale = 1  # 缩放因子：控制箭头的长度，如果箭头太短看不清，把这个改大 (例如 2 或 3)
    color = (0, 255, 0)  # 绿色箭头

    print("按 'q' 退出...")

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

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

        # 4. 计算稠密光流 (Farneback 算法)
        # flow 是一个 (h, w, 2) 的数组，flow[..., 0] 是水平移动(dx)，flow[..., 1] 是垂直移动(dy)
        flow = cv2.calcOpticalFlowFarneback(prev_gray, gray, None,
                                            pyr_scale=0.5, levels=3, winsize=15,
                                            iterations=3, poly_n=5, poly_sigma=1.2, flags=0)

        # 5. 可视化：在原图上画箭头
        # 创建网格点 (只在 step 的倍数位置画)
        h, w = gray.shape
        y, x = np.mgrid[step / 2:h:step, step / 2:w:step].reshape(2, -1).astype(int)

        # 获取这些网格点对应的 dx, dy
        fx, fy = flow[y, x].T

        # 创建画线的图层
        vis_frame = frame.copy()

        # 绘制线条
        lines = np.vstack([x, y, x + fx * scale, y + fy * scale]).T.reshape(-1, 2, 2)
        lines = np.int32(lines + 0.5)

        # 遍历绘制箭头 (cv2.arrowedLine)
        for (x1, y1), (x2, y2) in lines:
            # 只有当运动幅度超过一定阈值才画，减少噪点 (对应 MATLAB 的 NoiseThreshold)
            if np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2) > 1:
                cv2.arrowedLine(vis_frame, (x1, y1), (x2, y2), color, 1, tipLength=0.3)

        # 6. 显示结果
        cv2.imshow('Optical Flow (Dense)', vis_frame)

        # 更新上一帧
        prev_gray = gray

        # 按 'q' 退出
        if cv2.waitKey(20) & 0xFF == ord('q'):  # waitKey(20) 控制播放速度
            break

    cap.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    # 请确保目录下有 thumb.avi，或者改成 0 使用摄像头测试
    video_file = "thumb.avi"

    # 如果没有视频文件，想用摄像头测试，取消下面这行的注释：
    # video_file = 0

    run_optical_flow(video_file)