简单的视频管道从多个文件读取

目标

在本示例中，我们将介绍如何使用 readers.video 操作符创建 pipeline。该 pipeline 将返回一批帧序列。这些序列是任意数量的帧（图像）。不同之处在于图像的维度为 HWC，而序列的维度为 FHWC。

有关 readers.video 参数的更多信息，请查看文档。

设置

首先，让我们从导入开始

import os
import numpy as np

from nvidia.dali import pipeline_def
import nvidia.dali.fn as fn
import nvidia.dali.types as types

我们需要一些视频容器来处理。我们可以使用 Sintel 预告片，这是一个包含 h264 视频的 mp4 容器，并在 Create Common 许可下分发。让我们将其拆分为 10 秒的片段，以检查 readers.Video 如何处理多个视频文件。这可以使用 ffmpeg 独立工具轻松完成。

然后我们可以设置将在 pipeline 中使用的参数。count 参数将定义每个序列样本中我们想要的帧数。

我们可以将 video_directory 替换为任何其他包含 FFmpeg 识别的视频容器文件的目录。

batch_size = 2
sequence_length = 8
initial_prefetch_size = 16
video_directory = os.path.join(
    os.environ["DALI_EXTRA_PATH"], "db", "video", "sintel", "video_files"
)
video_files = [video_directory + "/" + f for f in os.listdir(video_directory)]
n_iter = 6

注意：DALI_EXTRA_PATH 环境变量应指向从 DALI extra 存储库 下载数据的位置。请确保检出正确的发布标签。

运行 Pipeline

然后我们可以定义一个最小的 pipeline，它将直接输出 readers.video 输出

@pipeline_def
def video_pipe(filenames):
    videos = fn.readers.video(
        device="gpu",
        filenames=filenames,
        sequence_length=sequence_length,
        shard_id=0,
        num_shards=1,
        random_shuffle=True,
        initial_fill=initial_prefetch_size,
    )
    return videos

注意：这里一个重要的点是调整 initial_fill，它对应于读取器预取缓冲区的初始大小。由于此缓冲区将填充 initial_fill 序列，因此帧的总数可能非常大！因此，请相应地设置它，以避免在训练期间 OOM。

让我们尝试在设备 0 上构建并运行一个 video_pipe，它将在每次迭代时输出 batch_size 个 count 帧的序列。

pipe = video_pipe(
    batch_size=batch_size,
    num_threads=2,
    device_id=0,
    filenames=video_files,
    seed=123456,
)
pipe.build()
for i in range(n_iter):
    pipe_out = pipe.run()
    sequences_out = pipe_out[0].as_cpu().as_array()
    print(sequences_out.shape)

(2, 8, 720, 1280, 3)
(2, 8, 720, 1280, 3)
(2, 8, 720, 1280, 3)
(2, 8, 720, 1280, 3)
(2, 8, 720, 1280, 3)
(2, 8, 720, 1280, 3)

可视化结果

之前的迭代似乎产生了预期形状的批次。我们将使用 matplotlib 来显示我们在最后一个批次中获得的帧。

%matplotlib inline
from matplotlib import pyplot as plt
import matplotlib.gridspec as gridspec

def show_sequence(sequence):
    columns = 4
    rows = (sequence_length + 1) // (columns)
    fig = plt.figure(figsize=(32, (16 // columns) * rows))
    gs = gridspec.GridSpec(rows, columns)
    for j in range(rows * columns):
        plt.subplot(gs[j])
        plt.axis("off")
        plt.imshow(sequence[j])

pipe_out = pipe.run()
sequences_out = pipe_out[0].as_cpu().as_array()
show_sequence(sequences_out[0])