fix(rtc): use the right plane layout for I420A, I422 and NV12 frames

livekit-rtc/livekit/rtc/video_frame.py

@@ -280,7 +280,8 @@ def _get_plane_length(type: proto_video.VideoBufferType.ValueType, width: int, h
         return width * height + chroma_width * chroma_height * 2
     elif type == proto_video.VideoBufferType.I420A:
         chroma_width = (width + 1) // 2
-        return width * height * 2 + chroma_width * chroma_width * 2
+        chroma_height = (height + 1) // 2
+        return width * height * 2 + chroma_width * chroma_height * 2
     elif type == proto_video.VideoBufferType.I422:
         chroma_width = (width + 1) // 2
         return width * height + chroma_width * height * 2
@@ -293,7 +294,7 @@ def _get_plane_length(type: proto_video.VideoBufferType.ValueType, width: int, h
     elif type == proto_video.VideoBufferType.NV12:
         chroma_width = (width + 1) // 2
         chroma_height = (height + 1) // 2
-        return width * height + chroma_width * chroma_width * 2
+        return width * height + chroma_width * chroma_height * 2
 
     raise Exception(f"unsupported video buffer type: {type}")
 
@@ -320,7 +321,7 @@ def _get_plane_infos(
         chroma_width = (width + 1) // 2
         y = _component_info(addr, width, width * height)
         u = _component_info(y.data_ptr + y.size, chroma_width, chroma_width * height)
-        v = _component_info(u.data_ptr + u.size + u.size, chroma_width, chroma_width * height)
+        v = _component_info(u.data_ptr + u.size, chroma_width, chroma_width * height)
         return [y, u, v]
     elif type == proto_video.VideoBufferType.I444:
         y = _component_info(addr, width, width * height)

tests/rtc/test_video_frame.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import ctypes
+
+import pytest
+
+from livekit import rtc
+from livekit.rtc import video_frame
+from livekit.rtc._proto import video_frame_pb2 as proto_video
+
+T = proto_video.VideoBufferType
+
+
+def _chroma(n: int) -> int:
+    return (n + 1) // 2
+
+
+# Plane sizes in bytes, in memory order, for a tightly packed frame of each format.
+def _expected_planes(t: int, w: int, h: int) -> list[int]:
+    cw, ch = _chroma(w), _chroma(h)
+    return {
+        T.I420: [w * h, cw * ch, cw * ch],
+        T.I420A: [w * h, cw * ch, cw * ch, w * h],
+        T.I422: [w * h, cw * h, cw * h],
+        T.I444: [w * h, w * h, w * h],
+        T.I010: [w * h * 2, cw * ch * 2, cw * ch * 2],
+        T.NV12: [w * h, cw * ch * 2],
+    }[t]
+
+
+@pytest.mark.parametrize("name", ["I420", "I420A", "I422", "I444", "I010", "NV12"])
+@pytest.mark.parametrize(("w", "h"), [(640, 480), (480, 640), (5, 3)])
+def test_planes_are_packed_and_match_the_frame_length(name: str, w: int, h: int) -> None:
+    t = getattr(T, name)
+    planes = video_frame._get_plane_infos(0, t, w, h)
+
+    assert [p.size for p in planes] == _expected_planes(t, w, h)
+    offset = 0
+    for plane in planes:
+        assert plane.data_ptr == offset
+        offset += plane.size
+    assert video_frame._get_plane_length(t, w, h) == offset
+
+
+@pytest.mark.parametrize("name", ["I420A", "NV12"])
+def test_frame_from_the_ffi_copies_exactly_the_frame(name: str) -> None:
+    # Frames received from a VideoStream are copied out of FFI memory with this length.
+    t, w, h = getattr(T, name), 640, 480
+    expected = sum(_expected_planes(t, w, h))
+    buf = (ctypes.c_uint8 * (expected * 2))()  # room past the frame, so main over-reads safely
+    owned = proto_video.OwnedVideoBuffer()
+    owned.info.type, owned.info.width, owned.info.height = t, w, h
+    owned.info.data_ptr = ctypes.addressof(buf)
+
+    frame = rtc.VideoFrame._from_owned_info(owned)
+
+    assert len(frame.data) == expected
+
+
+def test_get_plane_returns_each_i422_plane() -> None:
+    w, h = 4, 2
+    y, u, v = b"\x01" * (w * h), b"\x02" * (_chroma(w) * h), b"\x03" * (_chroma(w) * h)
+    frame = rtc.VideoFrame(w, h, T.I422, y + u + v)
+
+    assert [bytes(frame.get_plane(i)) for i in range(3)] == [y, u, v]