StreamWriter 進階用法

作者: Moto Hira

本教學課程展示如何使用 torchaudio.io.StreamWriter 播放音訊和視訊。

注意

本教學課程使用硬體裝置，因此無法跨不同的作業系統移植。

本教學課程在 MacBook Pro (M1, 2020) 上撰寫和測試。

注意

本教學課程需要 FFmpeg 函式庫。有關詳細資訊，請參閱 FFmpeg 相依性。

警告

TorchAudio 會動態載入系統上安裝的相容 FFmpeg 函式庫。支援的格式類型 (媒體格式、編碼器、編碼器選項等) 取決於這些函式庫。

若要檢查可用的裝置、多工器和編碼器，您可以使用以下命令

ffmpeg -muxers
ffmpeg -encoders
ffmpeg -devices
ffmpeg -protocols

準備

import torch
import torchaudio

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

from torchaudio.io import StreamWriter

from torchaudio.utils import download_asset

AUDIO_PATH = download_asset("tutorial-assets/Lab41-SRI-VOiCES-src-sp0307-ch127535-sg0042.wav")
VIDEO_PATH = download_asset(
    "tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4_small.mp4"
)

裝置可用性

StreamWriter 利用 FFmpeg 的 IO 抽象化，並將資料寫入媒體裝置，例如喇叭和 GUI。

若要寫入裝置，請將 format 選項提供給 StreamWriter 的建構函式。

不同的作業系統會有不同的裝置選項，其可用性取決於 FFmpeg 的實際安裝情況。

若要檢查哪些裝置可用，您可以使用 ffmpeg -devices 命令。

「audiotoolbox」(喇叭) 和「sdl」(視訊 GUI) 可用。

$ ffmpeg -devices
...
Devices:
 D. = Demuxing supported
 .E = Muxing supported
 --
  E audiotoolbox    AudioToolbox output device
 D  avfoundation    AVFoundation input device
 D  lavfi           Libavfilter virtual input device
  E opengl          OpenGL output
  E sdl,sdl2        SDL2 output device

有關哪些裝置在哪些作業系統上可用的詳細資訊，請查看官方 FFmpeg 文件。https://ffmpeg.dev.org.tw/ffmpeg-devices.html

播放音訊

透過提供 format="audiotoolbox" 選項，StreamWriter 會將資料寫入喇叭裝置。

# Prepare sample audio
waveform, sample_rate = torchaudio.load(AUDIO_PATH, channels_first=False, normalize=False)
num_frames, num_channels = waveform.shape

# Configure StreamWriter to write to speaker device
s = StreamWriter(dst="-", format="audiotoolbox")
s.add_audio_stream(sample_rate, num_channels, format="s16")

# Write audio to the device
with s.open():
    for i in range(0, num_frames, 256):
        s.write_audio_chunk(0, waveform[i : i + 256])

注意

寫入「audiotoolbox」是封鎖操作，但它不會等待音訊播放。在播放音訊時，必須保持裝置開啟。

以下程式碼會在音訊寫入後立即關閉裝置，並在播放完成之前關閉。新增 time.sleep() 將有助於保持裝置開啟，直到播放完成。

with s.open():
    s.write_audio_chunk(0, waveform)

播放視訊

若要播放影片，您可以使用 format="sdl" 或 format="opengl"。同樣地，您需要啟用相應整合功能的 FFmpeg 版本。可以使用 ffmpeg -devices 來檢查可用的裝置。

在此，我們使用 SDL 裝置 (https://ffmpeg.dev.org.tw/ffmpeg-devices.html#sdl)。

# note:
#  SDL device does not support specifying frame rate, and it has to
#  match the refresh rate of display.
frame_rate = 120
width, height = 640, 360

由於我們定義了一個輔助函式，將影片載入委派給背景執行緒並給予區塊。

running = True


def video_streamer(path, frames_per_chunk):
    import queue
    import threading

    from torchaudio.io import StreamReader

    q = queue.Queue()

    # Streaming process that runs in background thread
    def _streamer():
        streamer = StreamReader(path)
        streamer.add_basic_video_stream(
            frames_per_chunk, format="rgb24", frame_rate=frame_rate, width=width, height=height
        )
        for (chunk_,) in streamer.stream():
            q.put(chunk_)
            if not running:
                break

    # Start the background thread and fetch chunks
    t = threading.Thread(target=_streamer)
    t.start()
    while running:
        try:
            yield q.get()
        except queue.Empty:
            break
    t.join()

現在我們開始串流。按下“Q”將停止影片。

注意

對 SDL 裝置呼叫 write_video_chunk 會阻塞，直到 SDL 完成影片播放。

# Set output device to SDL
s = StreamWriter("-", format="sdl")

# Configure video stream (RGB24)
s.add_video_stream(frame_rate, width, height, format="rgb24", encoder_format="rgb24")

# Play the video
with s.open():
    for chunk in video_streamer(VIDEO_PATH, frames_per_chunk=256):
        try:
            s.write_video_chunk(0, chunk)
        except RuntimeError:
            running = False
            break

[程式碼]

串流影片

到目前為止，我們已經了解如何寫入硬體裝置。還有一些替代的影片串流方法。

RTMP（即時訊息協定）

使用 RMTP，您可以將媒體（影片和/或音訊）串流到單一用戶端。這不需要硬體裝置，但需要一個獨立的播放器。

要使用 RMTP，請在 StreamWriter 建構函式的 dst 參數中指定協定和路由，然後在開啟目的地時傳遞 {"listen": "1"} 選項。

StreamWriter 將監聽該連接埠，並等待用戶端請求影片。對 open 的呼叫會被阻塞，直到收到請求。

s = StreamWriter(dst="rtmp://127.0.0.1:1935/live/app", format="flv")
s.add_audio_stream(sample_rate=sample_rate, num_channels=num_channels, encoder="aac")
s.add_video_stream(frame_rate=frame_rate, width=width, height=height)

with s.open(option={"listen": "1"}):
    for video_chunk, audio_chunk in generator():
        s.write_audio_chunk(0, audio_chunk)
        s.write_video_chunk(1, video_chunk)

[程式碼]

UDP（使用者資料包協定）

使用 UDP，您可以將媒體（影片和/或音訊）串流到 socket。這不需要硬體裝置，但需要一個獨立的播放器。

與 RTMP 不同，串流和客戶端進程是斷開連接的。串流進程不知道客戶端進程。

s = StreamWriter(dst="udp://127.0.0.1:48550", format="mpegts")
s.add_audio_stream(sample_rate=sample_rate, num_channels=num_channels, encoder="aac")
s.add_video_stream(frame_rate=frame_rate, width=width, height=height)

with s.open():
    for video_chunk, audio_chunk in generator():
        s.write_audio_chunk(0, audio_chunk)
        s.write_video_chunk(1, video_chunk)

[程式碼]

標籤：torchaudio.io