working video sources

livekit · Oct 27, 2023 · b2dc10b · b2dc10b
1 parent bb7bc6f
commit b2dc10b
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 14 deletions.
diff --git a/examples/publish_hue.py b/examples/publish_hue.py
@@ -3,6 +3,8 @@
 import logging
 from signal import SIGINT, SIGTERM
 
+import cv2
+
 import numpy as np
 from livekit import rtc
 
@@ -11,10 +13,9 @@
 
 
 async def draw_color_cycle(source: rtc.VideoSource):
-    argb_frame = rtc.ArgbFrame(
+    argb_frame = rtc.ArgbFrame.create(
         rtc.VideoFormatType.FORMAT_ARGB, 1280, 720)
-
-    arr = np.ctypeslib.as_array(argb_frame.data)
+    arr = np.frombuffer(argb_frame.data, dtype=np.uint8)
 
     framerate = 1 / 30
     hue = 0.0
@@ -33,6 +34,7 @@ async def draw_color_cycle(source: rtc.VideoSource):
 
         frame = rtc.VideoFrame(
             0, rtc.VideoRotation.VIDEO_ROTATION_0, argb_frame.to_i420())
+
         source.capture_frame(frame)
         hue = (hue + framerate / 3) % 1.0
 

diff --git a/livekit-rtc/livekit/rtc/video_frame.py b/livekit-rtc/livekit/rtc/video_frame.py
@@ -132,7 +132,6 @@ def to_argb(self, dst: 'ArgbFrame') -> None:
         self.to_i420().to_argb(dst)
 
 
-
 class PlanarYuvBuffer(VideoFrameBuffer, ABC):
     def __init__(self,
                  data: bytearray,
@@ -153,12 +152,14 @@ def __init__(self,
 
     def _proto_info(self) -> proto_video_frame.VideoFrameBufferInfo:
         info = proto_video_frame.VideoFrameBufferInfo()
-        info.width = self._width
-        info.height = self._height
-        info.buffer_type = self._buffer_type
-        info.yuv.stride_y = self._stride_y
-        info.yuv.stride_u = self._stride_u
-        info.yuv.stride_v = self._stride_v
+        info.width = self.width
+        info.height = self.height
+        info.yuv.chroma_width = self.chroma_width
+        info.yuv.chroma_height = self.chroma_height
+        info.buffer_type = self.type
+        info.yuv.stride_y = self.stride_y
+        info.yuv.stride_u = self.stride_u
+        info.yuv.stride_v = self.stride_v
         return info
 
     @property
@@ -193,8 +194,8 @@ def __init__(self,
                  stride_v: int,
                  chroma_width: int,
                  chroma_height: int) -> None:
-        super().__init__(data, width, height, buffer_type, stride_u,
-                         stride_y, stride_v, chroma_width, chroma_height)
+        super().__init__(data, width, height, buffer_type, stride_y,
+                         stride_u, stride_v, chroma_width, chroma_height)
 
     def _proto_info(self) -> proto_video_frame.VideoFrameBufferInfo:
         info = super()._proto_info()
@@ -282,6 +283,8 @@ def _proto_info(self) -> proto_video_frame.VideoFrameBufferInfo:
         info = proto_video_frame.VideoFrameBufferInfo()
         info.width = self._width
         info.height = self._height
+        info.bi_yuv.chroma_width = self.chroma_width
+        info.bi_yuv.chroma_height = self.chroma_height
         info.buffer_type = self._buffer_type
         info.bi_yuv.stride_y = self._stride_y
         info.bi_yuv.stride_uv = self._stride_uv
@@ -575,8 +578,12 @@ def __init__(self,
         self._height = height
         self._stride = stride
 
+    @staticmethod
+    def create(format: VideoFormatType.ValueType, width: int, height: int) -> 'ArgbFrame':
+        data = bytearray(width * height * ctypes.sizeof(ctypes.c_uint32))
+        return ArgbFrame(data, format, width, height)
+
     def to_i420(self) -> I420Buffer:
-        # TODO(theomonnom): avoid unnecessary buffer allocation
         req = proto_ffi.FfiRequest()
         req.to_i420.argb.format = self.format
         req.to_i420.argb.width = self.width

diff --git a/livekit-rtc/livekit/rtc/video_source.py b/livekit-rtc/livekit/rtc/video_source.py
@@ -31,7 +31,7 @@ def __init__(self) -> None:
     def capture_frame(self, frame: VideoFrame) -> None:
         req = proto_ffi.FfiRequest()
         req.capture_video_frame.source_handle = self._ffi_handle.handle
-        req.capture_video_frame.buffer_handle = frame.buffer._ffi_handle.handle
+        req.capture_video_frame.info.CopyFrom(frame.buffer._proto_info())
         req.capture_video_frame.frame.rotation = frame.rotation
         req.capture_video_frame.frame.timestamp_us = frame.timestamp_us
         ffi_client.request(req)