How do I load and store an mp4 video file?

[anna-charlotte]

Overview

I want to have two functions:

• one to load a video array from an .mp4 file
• to store a video array to an .mp4 file

I tried implementing those as follows:

from ctypes import Union
from pathlib import Path
from typing import BinaryIO

import av
import numpy as np


def load(file_path, **kwargs) -> np.ndarray:
    frames = []
    with av.open(file_path, **kwargs) as container:
        for i, frame in enumerate(container.decode(video=0)):
            img = frame.to_image()
            frames.append(img)

    return np.moveaxis(np.stack(frames), -3, -2)


# following this approach: https://github.com/PyAV-Org/PyAV/issues/498
def save_to_file(
    np_tensor: np.ndarray,
    file_path: str,
    frame_rate: int = 30,
    codec: str = 'mpeg4',
) -> None:

    video_tensor = np.moveaxis(np.clip(np_tensor, 0, 255), -3, -2).astype('uint8')

    with av.open(file_path, mode='w') as container:
        if video_tensor.ndim == 3:
            video_tensor = np.expand_dims(video_tensor, axis=0)

        stream = container.add_stream(codec, rate=frame_rate)
        stream.width = video_tensor.shape[-2]
        stream.height = video_tensor.shape[-3]
        stream.pix_fmt = 'yuv420p'

        for vid in video_tensor:
            frame = av.VideoFrame.from_ndarray(vid, format="rgb24")
            frame.pts = None

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

        for packet in stream.encode(None):
            container.mux(packet)


tmp_file = str(Path.cwd() / 'tmp.mp4')

video_tensor = np.array(
    [[
        [[1, 2, 3], [4, 5, 6]],
        [[7, 8, 9], [10, 11, 12]],
        [[13, 14, 15], [16, 17, 18]],
        [[19, 20, 21], [22, 23, 24]],
    ]]
)


save_to_file(np_tensor=video_tensor, file_path=tmp_file)

video_from_file = load(tmp_file)
assert video_tensor.shape == video_from_file.shape
print(f"tensor = {video_tensor}")
print(f"video_from_file = {video_from_file}")
assert np.allclose(video_tensor, video_from_file)

Expected behavior

I have an mp4 file, that I want to load to video_1. When storing the video_1 array to an mp4 file, and then load the video from that new file to a second array, video_2, I expect: np.allclose(video_1, video_2)

Actual behavior

I am getting an Assertion error:

tensor = [[[[ 1  2  3]
   [ 4  5  6]]

  [[ 7  8  9]
   [10 11 12]]

  [[13 14 15]
   [16 17 18]]

  [[19 20 21]
   [22 23 24]]]]
video_from_file = [[[[ 3  4  7]
   [ 3  4  7]]

  [[ 8  9 12]
   [ 8  9 12]]

  [[15 16 19]
   [15 16 19]]

  [[20 21 24]
   [20 21 24]]]]
Traceback (most recent call last):
  File "/Users/charlottegerhaher/Library/Application Support/JetBrains/PyCharm2022.2/scratches/scratch_26.py", line 68, in <module>
    assert np.allclose(video_tensor, video_from_file)
AssertionError

Traceback:

Traceback (most recent call last):
  File "/Users/charlottegerhaher/Library/Application Support/JetBrains/PyCharm2022.2/scratches/scratch_26.py", line 68, in <module>
    assert np.allclose(video_tensor, video_from_file)
AssertionError

Research

I have done the following:

• Checked the PyAV documentation
• Searched on Google
• Searched on Stack Overflow
• Looked through old GitHub issues
• Asked on PyAV Gitter
• ... and waited 72 hours for a response.