82 lines
2.7 KiB
Python
82 lines
2.7 KiB
Python
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) |