livekit-utils crate + VideoRenderer proto

examples/simple_room/src/video_renderer.rs

@@ -0,0 +1,172 @@
+use livekit::webrtc::media_stream::VideoTrack;
+use livekit::webrtc::video_frame_buffer::PlanarYuv8Buffer;
+use livekit::webrtc::video_frame_buffer::PlanarYuvBuffer;
+use livekit::webrtc::video_frame_buffer::VideoFrameBufferTrait;
+use livekit::webrtc::yuv_helper;
+use std::convert::TryInto;
+use std::num::NonZeroU32;
+use std::{
+    ops::DerefMut,
+    sync::{Arc, Mutex},
+};
+
+pub struct VideoRenderer {
+    internal: Arc<Mutex<RendererInternal>>,
+    rtc_track: Arc<VideoTrack>,
+}
+
+struct RendererInternal {
+    render_state: egui_wgpu::RenderState,
+    width: u32,
+    height: u32,
+    rgba_data: Vec<u8>,
+    texture: Option<wgpu::Texture>,
+    texture_view: Option<wgpu::TextureView>,
+    egui_texture: Option<egui::TextureId>,
+}
+
+impl RendererInternal {
+    fn ensure_texture_size(&mut self, width: u32, height: u32) {
+        if self.width == width && self.height == height {
+            return;
+        }
+
+        self.width = width;
+        self.height = height;
+        self.rgba_data.resize((width * height * 4) as usize, 0);
+
+        self.texture = Some(
+            self.render_state
+                .device
+                .create_texture(&wgpu::TextureDescriptor {
+                    label: Some("lk-videotexture"),
+                    usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
+                    dimension: wgpu::TextureDimension::D2,
+                    size: wgpu::Extent3d {
+                        width,
+                        height,
+                        ..Default::default()
+                    },
+                    sample_count: 1,
+                    mip_level_count: 1,
+                    format: wgpu::TextureFormat::Rgba8UnormSrgb,
+                }),
+        );
+
+        self.texture_view = Some(self.texture.as_mut().unwrap().create_view(
+            &wgpu::TextureViewDescriptor {
+                label: Some("lk-videotexture-view"),
+                format: Some(wgpu::TextureFormat::Rgba8UnormSrgb),
+                dimension: Some(wgpu::TextureViewDimension::D2),
+                mip_level_count: NonZeroU32::new(1),
+                array_layer_count: NonZeroU32::new(1),
+                ..Default::default()
+            },
+        ));
+
+        if let Some(texture_id) = self.egui_texture {
+            // Update the existing texture
+            self.render_state
+                .renderer
+                .write()
+                .update_egui_texture_from_wgpu_texture(
+                    &*self.render_state.device,
+                    self.texture_view.as_ref().unwrap(),
+                    wgpu::FilterMode::Linear,
+                    texture_id,
+                );
+        } else {
+            self.egui_texture = Some(self.render_state.renderer.write().register_native_texture(
+                &*self.render_state.device,
+                self.texture_view.as_ref().unwrap(),
+                wgpu::FilterMode::Linear,
+            ));
+        }
+    }
+}
+
+impl VideoRenderer {
+    pub fn new(render_state: egui_wgpu::RenderState, rtc_track: Arc<VideoTrack>) -> Self {
+        let internal = Arc::new(Mutex::new(RendererInternal {
+            render_state,
+            width: 0,
+            height: 0,
+            rgba_data: Vec::default(),
+            texture: None,
+            texture_view: None,
+            egui_texture: None,
+        }));
+
+        rtc_track.on_frame({
+            let internal = internal.clone();
+
+            Box::new(move |_frame, buffer| {
+                let mut internal = internal.lock().unwrap();
+                let buffer = buffer.to_i420();
+
+                let width: u32 = buffer.width().try_into().unwrap();
+                let height: u32 = buffer.height().try_into().unwrap();
+
+                internal.ensure_texture_size(width, height);
+
+                let rgba_ptr = internal.rgba_data.deref_mut();
+                let rgba_stride = buffer.width() * 4;
+
+                yuv_helper::i420_to_abgr(
+                    buffer.data_y(),
+                    buffer.stride_y(),
+                    buffer.data_u(),
+                    buffer.stride_u(),
+                    buffer.data_v(),
+                    buffer.stride_v(),
+                    rgba_ptr,
+                    rgba_stride,
+                    buffer.width(),
+                    buffer.height(),
+                );
+
+                let copy_desc = wgpu::ImageCopyTexture {
+                    texture: internal.texture.as_ref().unwrap(),
+                    mip_level: 0,
+                    origin: wgpu::Origin3d::default(),
+                    aspect: wgpu::TextureAspect::default(),
+                };
+
+                let copy_layout = wgpu::ImageDataLayout {
+                    bytes_per_row: Some(NonZeroU32::new(width * 4).unwrap()),
+                    ..Default::default()
+                };
+
+                let copy_size = wgpu::Extent3d {
+                    width,
+                    height,
+                    ..Default::default()
+                };
+
+                internal.render_state.queue.write_texture(
+                    copy_desc,
+                    &internal.rgba_data,
+                    copy_layout,
+                    copy_size,
+                );
+
+                println!("wrote");
+            })
+        });
+
+        Self {
+            rtc_track,
+            internal,
+        }
+    }
+
+    pub fn texture_id(&self) -> Option<egui::TextureId> {
+        self.internal.lock().unwrap().egui_texture.clone()
+    }
+}
+
+impl Drop for VideoRenderer {
+    fn drop(&mut self) {
+        self.rtc_track.on_frame(Box::new(|_, _| {}));
+    }
+}