Pytorch – 使用Pyav解码视频文件并将视频帧转换为Pytorch tensor作为网络模型输入数据

1 视频文件作为网络模型的输入数据

越来越多的神经网络模型开始以视频作为训练数据，比如基于视频数据的行为识别等等，这就需要我们将视频转换为可适用的张量，本文将以pytorch为例，展示一下视频数据转换为pytorch tensor的过程。

2 使用Pyav解码视频文件并将视频帧转换为Pytorch tensor

2.1 安装ffmepg以及pyav

不赘述。

2.2 视频转换为tensor

import torch
import numpy as np
import os
import av

def get_tensor_from_video(video_path,is_multi_thread_decode = True):
    """
    :param video_path: 视频文件地址
    :param is_multi_thread_decode: 是否多线程解码文件
    :return: pytorch tensor
    """
    if not os.access(video_path, os.F_OK):
        print('测试文件不存在')
        return

    container = av.open(video_path)
    if is_multi_thread_decode:
        container.streams.video[0].thread_type = "AUTO"

    container.seek(0, any_frame=False, backward=True, stream=container.streams.video[0])

    frames = []
    for frame in container.decode(video=0):
        frames.append(frame)
    container.close()

    result_frams = None

    # 从视频帧转换为ndarray
    result_frames = [frame.to_rgb().to_ndarray() for frame in frames]
    # 转换成tensor
    result_frames = torch.as_tensor(np.stack(result_frames))

    # 注意：此时result_frames组成的维度为[视频帧数量，宽，高，通道数]
    return result_frames

if __name__ == '__main__':
    tensor = get_tensor_from_video(r'example.avi')
    print(tensor.shape)

详细的代码如上，请注意上述函数get_tensor_from_video返回的tensor的维度为[视频帧数量，宽，高，通道数]，而有的神经网络模型要求输入[通道数，视频帧数量，宽，高]，我们可以使用以下函数进行转换。

2.3 tensor维度交换

#  T H W C -> C T H W
result_frames = result_frames.permute(3, 0, 1, 2)