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

目标

在本示例中，我们将逐步创建使用 readers.video 操作符的管线。该管线将返回一批帧序列。这些序列是任意数量的帧（图像）。不同之处在于图像的维度是 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 视频并根据 Create Common 许可分发的 mp4 容器。让我们将其拆分为 10 秒的片段，以检查 readers.Video 如何处理多个视频文件。这可以使用 ffmpeg 独立工具轻松完成。

然后我们可以设置将在管线中使用的参数。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 存储库 下载数据的位置。请确保检出正确的发布标签。

运行管线

然后我们可以定义一个最小的管线，它将直接输出 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])