PyAV-Org/User-Help

Hi there! I'm running into a bit of a problem trying to write to a video file both an audio and video stream (from pytorch tensors) ---

This is my code so far:

def write_video(
    filename,
    video_array,
    video_fps,
    video_codec="libx264",
    audio_array=None,
    audio_fps=None,
    audio_codec=None,
    audio_layout=None,
    sample_format=None,
    options=None
):
    video_array = torch.as_tensor(video_array, dtype=torch.uint8).numpy()

    container = av.open(filename, mode="w")

    stream = container.add_stream(video_codec, rate=video_fps)
    stream.width = video_array.shape[2]
    stream.height = video_array.shape[1]
    stream.pix_fmt = "yuv420p" if video_codec != "libx264rgb" else "rgb24"
    stream.options = options or {}

    for img in video_array:
        frame = av.VideoFrame.from_ndarray(img, format="rgb24")
        frame.pict_type = "NONE"
        for packet in stream.encode(frame):
            container.mux(packet)

    # Flush stream
    for packet in stream.encode():
        container.mux(packet)

    if audio_array is not None:
        audio_array = torch.as_tensor(audio_array).numpy()
        a_stream = container.add_stream(audio_codec, rate=audio_fps)
        a_stream.options = {}

        frame = av.AudioFrame.from_ndarray(
            audio_array, format=sample_format, layout=audio_layout
        )

        for packet in a_stream.encode(frame):
            container.mux(packet)

        for packet in a_stream.encode():
            container.mux(packet)

    # Close the file
    container.close()

this is how im calling the function:

video_tensor, audio_tensor, info = torchvision.io.video.read_video("2.mp4", pts_unit="sec")
    write_video(
        "0.mp4",
        video_tensor,
        int(info["video_fps"]),
        "h264",
        audio_tensor,
        info["audio_fps"],
        "aac",
        "stereo",
        "fltp",
    )

However, I keep getting the following error:

Requested input sample rate 0 is invalid
Traceback (most recent call last):
  File "write.py", line 80, in <module>
    "fltp",
  File "write.py", line 59, in write_video
    for packet in a_stream.encode(frame):
  File "av/stream.pyx", line 152, in av.stream.Stream.encode
  File "av/codec/context.pyx", line 466, in av.codec.context.CodecContext.encode
  File "av/audio/codeccontext.pyx", line 40, in av.audio.codeccontext.AudioCodecContext._prepare_frames_for_encode
  File "av/audio/resampler.pyx", line 106, in av.audio.resampler.AudioResampler.resample
  File "av/audio/resampler.pyx", line 103, in av.audio.resampler.AudioResampler.resample
  File "av/error.pyx", line 336, in av.error.err_check
av.error.ValueError: [Errno 22] Invalid argument; last error log: [SWR] Requested input sample rate 0 is invalid
Traceback (most recent call last):
  File "av/container/output.pyx", line 25, in av.container.output.close_output
TypeError: 'NoneType' object is not iterable
Exception ignored in: 'av.container.output.OutputContainer.__dealloc__'
Traceback (most recent call last):
  File "av/container/output.pyx", line 25, in av.container.output.close_output
TypeError: 'NoneType' object is not iterable

I'm not sure why it says that the sample rate is set to zero when I set it as I add streams, so I'm not sure how to move forward. The video is still being produced, but it just has no audio. Would really appreciate any help!