使用opencv调用TV_L1算法提取光流

提取视频的光流并使用灰度图可视化

python 复制代码

import cv2
import numpy as np
import pickle
import os

def compute_tvl1_optical_flow(video_path):
    # 创建视频捕获对象
    cap = cv2.VideoCapture(video_path)
    ret, frame1 = cap.read()
    if not ret:
        print("Failed to read video")
        return None  # 返回None，表示无法读取视频

    # 将第一帧转换为灰度
    prvs = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)

    # 创建TVL1光流对象
    optical_flow = cv2.optflow.DualTVL1OpticalFlow_create()

    # 初始化用于存储光流数据的列表
    flows = []

    # 读取视频帧并计算光流
    while True:
        ret, frame2 = cap.read()
        if not ret:
            break

        # 将当前帧转换为灰度
        next = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)

        # 计算光流
        flow = optical_flow.calc(prvs, next, None)

        # 将光流数据添加到列表中
        flows.append(flow)

        # 更新前一帧
        prvs = next

    cap.release()

    # 返回光流数据列表
    return flows

def save_optical_flows(flows, filepath):
    with open(filepath, 'wb') as file:
        pickle.dump(flows, file)

def draw_gray(flow, scale=1, gamma=0.1):
    # 计算每个像素点的速度大小
    magnitude = np.sqrt(flow[:, :, 0] ** 2 + flow[:, :, 1] ** 2)

    # 应用gamma校正来增强低亮度区域的可视化
    magnitude = np.power(magnitude, gamma)

    # 将速度大小归一化到0-255的范围
    normalized_mag = cv2.normalize(magnitude, None, 0, 255, cv2.NORM_MINMAX)

    # 将浮点数转换为8-bit整数，并调整亮度
    gray_img = np.uint8(normalized_mag * scale)

    return gray_img

# 使用函数
video_path = '/home/yunchuan/video_features-master/sample/v_ZNVhz7ctTq0.mp4'
output_dir = '/home/yunchuan/video_features-master/sample/flow_images/'
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

flows = compute_tvl1_optical_flow(video_path)

if flows is not None:
    print("Successfully computed optical flows.")
    save_optical_flows(flows, '/home/yunchuan/video_features-master/sample/flows.pkl')  # 保存光流数据
    for i, flow in enumerate(flows):
        img = draw_gray(flow)
        flow_image_path = os.path.join(output_dir, f'gray_flow_{i}.png')
        cv2.imwrite(flow_image_path, img)  # 保存光流图像
    cv2.destroyAllWindows()
else:
    print("Failed to compute optical flows.")

计算两张图片之间的光流并使用灰度图可视化

python 复制代码

import cv2
import numpy as np
import os

def calculate_optical_flow(image1_path, image2_path):
    # 读取两张图片
    frame1 = cv2.imread(image1_path)
    frame2 = cv2.imread(image2_path)

    if frame1 is None or frame2 is None:
        print("Error loading images!")
        return None

    # 检查并调整两张图片的尺寸，确保它们大小相同
    h1, w1 = frame1.shape[:2]
    h2, w2 = frame2.shape[:2]
    h = min(h1, h2)
    w = min(w1, w2)
    h=180
    w=280

    # 将图片大小调整为最小的共同尺寸
    frame1 = cv2.resize(frame1, (w, h))
    frame2 = cv2.resize(frame2, (w, h))

    # 将图片转换为灰度
    prvs = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
    next = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)

    # 创建TVL1光流对象
    optical_flow = cv2.optflow.DualTVL1OpticalFlow_create()

    # 计算两张图片之间的光流
    flow = optical_flow.calc(prvs, next, None)
    return flow




def draw_gray(flow, scale=1, gamma=0.1):
    # 计算每个像素点的速度大小
    magnitude = np.sqrt(flow[:, :, 0] ** 2 + flow[:, :, 1] ** 2)

    # 应用gamma校正来增强低亮度区域的可视化
    magnitude = np.power(magnitude, gamma)

    # 将速度大小归一化到0-255的范围
    normalized_mag = cv2.normalize(magnitude, None, 0, 255, cv2.NORM_MINMAX)

    # 将浮点数转换为8-bit整数，并调整亮度
    gray_img = np.uint8(normalized_mag * scale)

    return gray_img


# 指定两张图片的路径
image1_path = '/home/yunchuan/video_features-master/sample_images/image1.jpg'
image2_path = '/home/yunchuan/video_features-master/sample_images/image2.jpg'

# 输出路径
output_dir = '/home/yunchuan/video_features-master/sample_images/'
if not os.path.exists(output_dir):
    os.makedirs(output_dir)


# 计算光流
flow = calculate_optical_flow(image1_path, image2_path)
if flow is not None:
    # 可视化光流为灰度图
    gray_flow_image = draw_gray(flow)
    gray_flow_image_path = os.path.join(output_dir, 'optical_flow_visualization.png')
    cv2.imwrite(gray_flow_image_path, gray_flow_image)
    print(f"Optical flow visualization saved to {gray_flow_image_path}")

    # 保存原始光流数据
    raw_flow_path = os.path.join(output_dir, 'raw_optical_flow.npy')
    np.save(raw_flow_path, flow)
    print(f"Raw optical flow data saved to {raw_flow_path}")
else:
    print("Failed to compute optical flow.")