StreamWriter 基本用法

本教程展示了如何使用 torchaudio.io.StreamWriter 将音频/视频数据编码并保存为各种格式/目标。

本教程需要 FFmpeg 库。有关详细信息，请参阅 FFmpeg 依赖。

TorchAudio 动态加载系统中已安装的兼容 FFmpeg 库。支持的格式类型（媒体格式、编码器、编码器选项等）取决于这些库。

要检查可用的复用器和编码器，您可以使用以下命令
ffmpeg -muxers
ffmpeg -encoders

准备工作

import torch
import torchaudio

print(torch.__version__)
print(torchaudio.__version__)

from torchaudio.io import StreamWriter

print("FFmpeg library versions")
for k, v in torchaudio.utils.ffmpeg_utils.get_versions().items():
    print(f"  {k}: {v}")

2.6.0
2.6.0
FFmpeg library versions
  libavcodec: (60, 3, 100)
  libavdevice: (60, 1, 100)
  libavfilter: (9, 3, 100)
  libavformat: (60, 3, 100)
  libavutil: (58, 2, 100)

import io
import os
import tempfile

from IPython.display import Audio, Video

from torchaudio.utils import download_asset

SAMPLE_PATH = download_asset("tutorial-assets/Lab41-SRI-VOiCES-src-sp0307-ch127535-sg0042.wav")
WAVEFORM, SAMPLE_RATE = torchaudio.load(SAMPLE_PATH, channels_first=False)
NUM_FRAMES, NUM_CHANNELS = WAVEFORM.shape

_BASE_DIR = tempfile.TemporaryDirectory()


def get_path(filename):
    return os.path.join(_BASE_DIR.name, filename)

基本用法

要将 Tensor 数据保存为媒体格式并使用 StreamWriter，需要执行以下三个步骤

指定输出
配置流
写入数据

以下代码展示了如何将音频数据保存为 WAV 文件。

# 1. Define the destination. (local file in this case)
path = get_path("test.wav")
s = StreamWriter(path)

# 2. Configure the stream. (8kHz, Stereo WAV)
s.add_audio_stream(
    sample_rate=SAMPLE_RATE,
    num_channels=NUM_CHANNELS,
)

# 3. Write the data
with s.open():
    s.write_audio_chunk(0, WAVEFORM)

Audio(path)

现在我们将更详细地探讨每个步骤。

写入目标

StreamWriter 支持不同类型的写入目标

本地文件
类文件对象
流媒体协议（如 RTMP 和 UDP）
媒体设备（扬声器和视频播放器）†

† 对于媒体设备，请参考 StreamWriter 高级用法。

本地文件

StreamWriter 支持将媒体保存到本地文件。

StreamWriter(dst="audio.wav")

StreamWriter(dst="audio.mp3")

这同样适用于静态图像和视频。

StreamWriter(dst="image.jpeg")

StreamWriter(dst="video.mpeg")

类文件对象

您也可以传递一个类文件对象。该对象必须实现符合 io.RawIOBase.write 的 write 方法。

# Open the local file as fileobj
with open("audio.wav", "wb") as dst:
    StreamWriter(dst=dst)

# In-memory encoding
buffer = io.BytesIO()
StreamWriter(dst=buffer)

流媒体协议

您可以使用流媒体协议来传输媒体内容。

# Real-Time Messaging Protocol
StreamWriter(dst="rtmp://localhost:1234/live/app", format="flv")

# UDP
StreamWriter(dst="udp://localhost:48550", format="mpegts")

配置输出流

一旦指定了目标，下一步就是配置流。对于典型的音频和静态图像情况，只需要一个流，但对于带音频的视频，至少需要配置两个流（一个用于音频，另一个用于视频）。

音频流

可以通过 add_audio_stream() 方法添加音频流。

对于写入常规音频文件，至少需要 sample_rate 和 num_channels 参数。

s = StreamWriter("audio.wav")
s.add_audio_stream(sample_rate=8000, num_channels=2)

默认情况下，音频流期望输入波形张量的类型为 torch.float32。在上述情况下，数据将被编码为 WAV 格式的默认编码格式，即 16 位有符号整数线性 PCM。StreamWriter 会在内部进行样本格式的转换。

如果编码器支持多种样本格式，并且您希望更改编码器的样本格式，可以使用 encoder_format 选项。

在以下示例中，StreamWriter 期望输入波形张量的数据类型为 torch.float32，但在编码时会将其转换为 16 位有符号整数。

s = StreamWriter("audio.mp3")
s.add_audio_stream(
    ...,
    encoder="libmp3lame",   # "libmp3lame" is often the default encoder for mp3,
                            # but specifying it manually, for the sake of illustration.

    encoder_format="s16p",  # "libmp3lame" encoder supports the following sample format.
                            #  - "s16p" (16-bit signed integer)
                            #  - "s32p" (32-bit signed integer)
                            #  - "fltp" (32-bit floating point)
)

如果您的波形 Tensor 的数据类型不是 torch.float32，您可以提供 format 选项来更改预期的数据类型。

以下示例配置 StreamWriter 以期望 torch.int16 类型的 Tensor。

# Audio data passed to StreamWriter must be torch.int16
s.add_audio_stream(..., format="s16")

下图展示了 format 和 encoder_format 选项在音频流中的工作原理。

视频流

要添加静态图像或视频流，您可以使用 add_video_stream() 方法。

至少需要提供 frame_rate、height 和 width 参数。

s = StreamWriter("video.mp4")
s.add_video_stream(frame_rate=10, height=96, width=128)

对于静态图像，请使用 frame_rate=1。

s = StreamWriter("image.png")
s.add_video_stream(frame_rate=1, ...)

与音频流类似，您可以提供 format 和 encoder_format 选项来控制输入数据的格式和编码。

以下示例以 YUV422 格式编码视频数据。

s = StreamWriter("video.mov")
s.add_video_stream(
    ...,
    encoder="libx264",  # libx264 supports different YUV formats, such as
                        # yuv420p yuvj420p yuv422p yuvj422p yuv444p yuvj444p nv12 nv16 nv21

    encoder_format="yuv422p",  # StreamWriter will convert the input data to YUV422 internally
)

YUV 格式通常用于视频编码。许多 YUV 格式的色度通道平面大小与亮度通道不同，这使得直接将其表示为 torch.Tensor 类型变得困难。因此，StreamWriter 会自动将输入的视频 Tensor 转换为目标格式。

StreamWriter 期望输入的图像 Tensor 是 4 维的（时间、通道、高度、宽度）并且是 torch.uint8 类型。

默认的颜色通道是 RGB，即对应红色、绿色和蓝色的三个颜色通道。如果您的输入具有不同的颜色通道，例如 BGR 或 YUV，您可以通过 format 选项来指定。

以下示例指定了 BGR 格式。

s.add_video_stream(..., format="bgr24")
                   # Image data passed to StreamWriter must have
                   # three color channels representing Blue Green Red.
                   #
                   # The shape of the input tensor has to be
                   # (time, channel==3, height, width)

下图展示了 format 和 encoder_format 选项在视频流中的工作原理。

写入数据

配置好流之后，下一步是打开输出位置并开始写入数据。

使用 open() 方法打开目标位置，然后使用 write_audio_chunk() 和/或 write_video_chunk() 写入数据。

音频张量的形状应为 (时间, 通道)，视频/图像张量的形状应为 (时间, 通道, 高度, 宽度)。

通道、高度和宽度必须与相应流的配置匹配，该配置通过 "format" 选项指定。

表示静态图像的张量在时间维度上必须只有一帧，但音频和视频张量在时间维度上可以有任意数量的帧。

以下代码片段展示了这一点；

示例) 音频

# Configure stream
s = StreamWriter(dst=get_path("audio.wav"))
s.add_audio_stream(sample_rate=SAMPLE_RATE, num_channels=NUM_CHANNELS)

# Write data
with s.open():
    s.write_audio_chunk(0, WAVEFORM)

示例) 图像

# Image config
height = 96
width = 128

# Configure stream
s = StreamWriter(dst=get_path("image.png"))
s.add_video_stream(frame_rate=1, height=height, width=width, format="rgb24")

# Generate image
chunk = torch.randint(256, (1, 3, height, width), dtype=torch.uint8)

# Write data
with s.open():
    s.write_video_chunk(0, chunk)

示例) 无音频的视频

# Video config
frame_rate = 30
height = 96
width = 128

# Configure stream
s = StreamWriter(dst=get_path("video.mp4"))
s.add_video_stream(frame_rate=frame_rate, height=height, width=width, format="rgb24")

# Generate video chunk (3 seconds)
time = int(frame_rate * 3)
chunk = torch.randint(256, (time, 3, height, width), dtype=torch.uint8)

# Write data
with s.open():
    s.write_video_chunk(0, chunk)

示例) 带音频的视频

要为视频写入音频，必须分别配置流。

# Configure stream
s = StreamWriter(dst=get_path("video.mp4"))
s.add_audio_stream(sample_rate=SAMPLE_RATE, num_channels=NUM_CHANNELS)
s.add_video_stream(frame_rate=frame_rate, height=height, width=width, format="rgb24")

# Generate audio/video chunk (3 seconds)
time = int(SAMPLE_RATE * 3)
audio_chunk = torch.randn((time, NUM_CHANNELS))
time = int(frame_rate * 3)
video_chunk = torch.randint(256, (time, 3, height, width), dtype=torch.uint8)

# Write data
with s.open():
    s.write_audio_chunk(0, audio_chunk)
    s.write_video_chunk(1, video_chunk)

逐块写入数据

在写入数据时，可以沿时间维度拆分数据，并按较小的块进行写入。

# Write data in one-go
dst1 = io.BytesIO()
s = StreamWriter(dst=dst1, format="mp3")
s.add_audio_stream(SAMPLE_RATE, NUM_CHANNELS)
with s.open():
    s.write_audio_chunk(0, WAVEFORM)

# Write data in smaller chunks
dst2 = io.BytesIO()
s = StreamWriter(dst=dst2, format="mp3")
s.add_audio_stream(SAMPLE_RATE, NUM_CHANNELS)
with s.open():
    for start in range(0, NUM_FRAMES, SAMPLE_RATE):
        end = start + SAMPLE_RATE
        s.write_audio_chunk(0, WAVEFORM[start:end, ...])

# Check that the contents are same
dst1.seek(0)
bytes1 = dst1.read()

print(f"bytes1: {len(bytes1)}")
print(f"{bytes1[:10]}...{bytes1[-10:]}\n")

dst2.seek(0)
bytes2 = dst2.read()

print(f"bytes2: {len(bytes2)}")
print(f"{bytes2[:10]}...{bytes2[-10:]}\n")

assert bytes1 == bytes2

import matplotlib.pyplot as plt

bytes1: 10700
b'ID3\x04\x00\x00\x00\x00\x00"'...b'\xaa\xaa\xaa\xaa\xaa\xaa\xaa\xaa\xaa\xaa'

bytes2: 10700
b'ID3\x04\x00\x00\x00\x00\x00"'...b'\xaa\xaa\xaa\xaa\xaa\xaa\xaa\xaa\xaa\xaa'

示例 - 频谱可视化器

在本节中，我们使用 StreamWriter 创建音频的频谱可视化并将其保存为视频文件。

为了创建频谱可视化，我们使用 torchaudio.transforms.Spectrogram 来获取音频的频谱表示，使用 matplotlib 生成其可视化的栅格图像，然后使用 StreamWriter 将其转换为带有原始音频的视频。

import torchaudio.transforms as T

准备数据

首先，我们准备频谱图数据。我们使用 Spectrogram。

我们调整 hop_length，使得频谱图的一帧对应视频的一帧。

frame_rate = 20
n_fft = 4000

trans = T.Spectrogram(
    n_fft=n_fft,
    hop_length=SAMPLE_RATE // frame_rate,  # One FFT per one video frame
    normalized=True,
    power=1,
)
specs = trans(WAVEFORM.T)[0].T

生成的频谱图如下所示。

spec_db = T.AmplitudeToDB(stype="magnitude", top_db=80)(specs.T)
_ = plt.imshow(spec_db, aspect="auto", origin="lower")

streamwriter basic tutorial

准备画布

我们使用 matplotlib 来可视化每帧的频谱图。我们创建了一个辅助函数，用于绘制频谱图数据并生成该图的栅格图像。

fig, ax = plt.subplots(figsize=[3.2, 2.4])
ax.set_position([0, 0, 1, 1])
ax.set_facecolor("black")
ncols, nrows = fig.canvas.get_width_height()


def _plot(data):
    ax.clear()
    x = list(range(len(data)))
    R, G, B = 238 / 255, 76 / 255, 44 / 255
    for coeff, alpha in [(0.8, 0.7), (1, 1)]:
        d = data**coeff
        ax.fill_between(x, d, -d, color=[R, G, B, alpha])
    xlim = n_fft // 2 + 1
    ax.set_xlim([-1, n_fft // 2 + 1])
    ax.set_ylim([-1, 1])
    ax.text(
        xlim,
        0.95,
        f"Created with TorchAudio\n{torchaudio.__version__}",
        color="white",
        ha="right",
        va="top",
        backgroundcolor="black",
    )
    fig.canvas.draw()
    frame = torch.frombuffer(fig.canvas.tostring_rgb(), dtype=torch.uint8)
    return frame.reshape(nrows, ncols, 3).permute(2, 0, 1)


# sphinx_gallery_defer_figures

写入视频

最后，我们使用 StreamWriter 并写入视频。我们一次处理一秒的音频和视频帧。

s = StreamWriter(get_path("example.mp4"))
s.add_audio_stream(sample_rate=SAMPLE_RATE, num_channels=NUM_CHANNELS)
s.add_video_stream(frame_rate=frame_rate, height=nrows, width=ncols)

with s.open():
    i = 0
    # Process by second
    for t in range(0, NUM_FRAMES, SAMPLE_RATE):
        # Write audio chunk
        s.write_audio_chunk(0, WAVEFORM[t : t + SAMPLE_RATE, :])

        # write 1 second of video chunk
        frames = [_plot(spec) for spec in specs[i : i + frame_rate]]
        if frames:
            s.write_video_chunk(1, torch.stack(frames))
        i += frame_rate

plt.close(fig)

/pytorch/audio/examples/tutorials/streamwriter_basic_tutorial.py:566: MatplotlibDeprecationWarning: The tostring_rgb function was deprecated in Matplotlib 3.8 and will be removed two minor releases later. Use buffer_rgba instead.
  frame = torch.frombuffer(fig.canvas.tostring_rgb(), dtype=torch.uint8)
/pytorch/audio/examples/tutorials/streamwriter_basic_tutorial.py:566: UserWarning: The given buffer is not writable, and PyTorch does not support non-writable tensors. This means you can write to the underlying (supposedly non-writable) buffer using the tensor. You may want to copy the buffer to protect its data or make it writable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at /opt/conda/conda-bld/pytorch_1727971112454/work/torch/csrc/utils/tensor_new.cpp:1560.)
  frame = torch.frombuffer(fig.canvas.tostring_rgb(), dtype=torch.uint8)

结果

结果如下所示。

Video(get_path("example.mp4"), embed=True)

仔细观看视频可以发现，“s”音（如curiosity、besides、this）在较高频率侧（视频的右侧）分配了更多的能量。

标签: torchaudio.io

下载 Python 源代码: streamwriter_basic_tutorial.py

下载 Jupyter 笔记本: streamwriter_basic_tutorial.ipynb