Merge pull request #127 from CapSoftware/raw-pixel-buffers

raw pixel buffers

Author: Brendonovich

src/capturer/engine/mac/apple_sys.rs

@@ -1,3 +1,5 @@
+#![allow(non_upper_case_globals)]
+
 pub use screencapturekit_sys::os_types::base::*;
 
 #[repr(C)]

src/capturer/engine/mac/mod.rs

@@ -1,6 +1,7 @@
 use std::cmp;
 use std::sync::mpsc;
 
+use pixelformat::get_pts_in_nanoseconds;
 use screencapturekit::{
     cm_sample_buffer::CMSampleBuffer,
     sc_content_filter::{InitParams, SCContentFilter},
@@ -16,13 +17,16 @@ use screencapturekit_sys::os_types::geometry::{CGPoint, CGRect, CGSize};
 
 use crate::frame::{Frame, FrameType};
 use crate::targets::Target;
-use crate::{capturer::Resolution, targets};
 use crate::{
-    capturer::{Area, Options, Point, Size},
+    capturer::{Area, Options, Point, Resolution, Size},
     frame::BGRAFrame,
+    targets,
 };
 
+use super::ChannelItem;
+
 mod apple_sys;
+mod pixel_buffer;
 mod pixelformat;
 
 struct ErrorHandler;
@@ -33,69 +37,22 @@ impl StreamErrorHandler for ErrorHandler {
 }
 
 pub struct Capturer {
-    pub tx: mpsc::Sender<Frame>,
-    pub output_type: FrameType,
+    pub tx: mpsc::Sender<ChannelItem>,
 }
 
 impl Capturer {
-    pub fn new(tx: mpsc::Sender<Frame>, output_type: FrameType) -> Self {
-        Capturer { tx, output_type }
+    pub fn new(tx: mpsc::Sender<ChannelItem>) -> Self {
+        Capturer { tx }
     }
 }
 
 impl StreamOutput for Capturer {
     fn did_output_sample_buffer(&self, sample: CMSampleBuffer, of_type: SCStreamOutputType) {
-        match of_type {
-            SCStreamOutputType::Screen => {
-                let frame_status = &sample.frame_status;
-
-                match frame_status {
-                    SCFrameStatus::Complete | SCFrameStatus::Started => unsafe {
-                        let frame = match self.output_type {
-                            FrameType::YUVFrame => {
-                                let yuvframe = pixelformat::create_yuv_frame(sample).unwrap();
-                                Frame::YUVFrame(yuvframe)
-                            }
-                            FrameType::RGB => {
-                                let rgbframe = pixelformat::create_rgb_frame(sample).unwrap();
-                                Frame::RGB(rgbframe)
-                            }
-                            FrameType::BGR0 => {
-                                let bgrframe = pixelformat::create_bgr_frame(sample).unwrap();
-                                Frame::BGR0(bgrframe)
-                            }
-                            FrameType::BGRAFrame => {
-                                let bgraframe = pixelformat::create_bgra_frame(sample).unwrap();
-                                Frame::BGRA(bgraframe)
-                            }
-                        };
-                        self.tx.send(frame).unwrap_or(());
-                    },
-                    SCFrameStatus::Idle => {
-                        // Quick hack - just send an empty frame, and the caller can figure out how to handle it
-                        match self.output_type {
-                            FrameType::BGRAFrame => {
-                                let display_time = pixelformat::get_pts_in_nanoseconds(&sample);
-                                let frame = BGRAFrame {
-                                    display_time,
-                                    width: 0,
-                                    height: 0,
-                                    data: vec![],
-                                };
-                                self.tx.send(Frame::BGRA(frame)).unwrap_or(());
-                            }
-                            _ => {}
-                        }
-                    }
-                    _ => {}
-                }
-            }
-            _ => {}
-        }
+        self.tx.send((sample, of_type)).unwrap_or(());
     }
 }
 
-pub fn create_capturer(options: &Options, tx: mpsc::Sender<Frame>) -> SCStream {
+pub fn create_capturer(options: &Options, tx: mpsc::Sender<ChannelItem>) -> SCStream {
     // If no target is specified, capture the main display
     let target = options
         .target
@@ -133,14 +90,13 @@ pub fn create_capturer(options: &Options, tx: mpsc::Sender<Frame>) -> SCStream {
                         .into_iter()
                         .filter(|window| {
                             excluded_targets
-                                .into_iter()
-                                .find(|excluded_target| match excluded_target {
+                                .iter()
+                                .any(|excluded_target| match excluded_target {
                                     Target::Window(excluded_window) => {
                                         excluded_window.id == window.window_id
                                     }
                                     _ => false,
                                 })
-                                .is_some()
                         })
                         .collect();
 
@@ -190,10 +146,7 @@ pub fn create_capturer(options: &Options, tx: mpsc::Sender<Frame>) -> SCStream {
     };
 
     let mut stream = SCStream::new(filter, stream_config, ErrorHandler);
-    stream.add_output(
-        Capturer::new(tx, options.output_type),
-        SCStreamOutputType::Screen,
-    );
+    stream.add_output(Capturer::new(tx), SCStreamOutputType::Screen);
 
     stream
 }
@@ -251,8 +204,8 @@ pub fn get_crop_area(options: &Options) -> Area {
                     y: val.origin.y,
                 },
                 size: Size {
-                    width: input_width as f64,
-                    height: input_height as f64,
+                    width: input_width,
+                    height: input_height,
                 },
             }
         })
@@ -264,3 +217,50 @@ pub fn get_crop_area(options: &Options) -> Area {
             },
         })
 }
+
+pub fn process_sample_buffer(
+    sample: CMSampleBuffer,
+    of_type: SCStreamOutputType,
+    output_type: FrameType,
+) -> Option<Frame> {
+    if let SCStreamOutputType::Screen = of_type {
+        let frame_status = &sample.frame_status;
+
+        match frame_status {
+            SCFrameStatus::Complete | SCFrameStatus::Started => unsafe {
+                return Some(match output_type {
+                    FrameType::YUVFrame => {
+                        let yuvframe = pixelformat::create_yuv_frame(sample).unwrap();
+                        Frame::YUVFrame(yuvframe)
+                    }
+                    FrameType::RGB => {
+                        let rgbframe = pixelformat::create_rgb_frame(sample).unwrap();
+                        Frame::RGB(rgbframe)
+                    }
+                    FrameType::BGR0 => {
+                        let bgrframe = pixelformat::create_bgr_frame(sample).unwrap();
+                        Frame::BGR0(bgrframe)
+                    }
+                    FrameType::BGRAFrame => {
+                        let bgraframe = pixelformat::create_bgra_frame(sample).unwrap();
+                        Frame::BGRA(bgraframe)
+                    }
+                });
+            },
+            SCFrameStatus::Idle => {
+                // Quick hack - just send an empty frame, and the caller can figure out how to handle it
+                if let FrameType::BGRAFrame = output_type {
+                    return Some(Frame::BGRA(BGRAFrame {
+                        display_time: get_pts_in_nanoseconds(&sample),
+                        width: 0,
+                        height: 0,
+                        data: vec![],
+                    }));
+                }
+            }
+            _ => {}
+        }
+    }
+
+    None
+}

src/capturer/engine/mac/pixel_buffer.rs

@@ -0,0 +1,191 @@
+use core::slice;
+use core_video_sys::{
+    CVPixelBufferGetBaseAddress, CVPixelBufferGetBaseAddressOfPlane, CVPixelBufferGetBytesPerRow,
+    CVPixelBufferGetBytesPerRowOfPlane, CVPixelBufferGetHeight, CVPixelBufferGetHeightOfPlane,
+    CVPixelBufferGetPlaneCount, CVPixelBufferGetWidth, CVPixelBufferGetWidthOfPlane,
+    CVPixelBufferLockBaseAddress, CVPixelBufferRef, CVPixelBufferUnlockBaseAddress,
+};
+use screencapturekit::{cm_sample_buffer::CMSampleBuffer, sc_types::SCFrameStatus};
+use screencapturekit_sys::cm_sample_buffer_ref::CMSampleBufferGetImageBuffer;
+use std::{ops::Deref, sync::mpsc};
+
+use crate::capturer::{engine::ChannelItem, RawCapturer};
+
+pub struct PixelBuffer {
+    display_time: u64,
+    width: usize,
+    height: usize,
+    bytes_per_row: usize,
+    buffer: CMSampleBuffer,
+}
+
+impl PixelBuffer {
+    pub fn display_time(&self) -> u64 {
+        self.display_time
+    }
+
+    pub fn width(&self) -> usize {
+        self.width
+    }
+
+    pub fn height(&self) -> usize {
+        self.height
+    }
+
+    pub fn buffer(&self) -> &CMSampleBuffer {
+        &self.buffer
+    }
+
+    pub fn bytes_per_row(&self) -> usize {
+        self.bytes_per_row
+    }
+
+    pub fn data(&self) -> PixelBufferData {
+        unsafe {
+            let pixel_buffer = sample_buffer_to_pixel_buffer(&self.buffer);
+
+            CVPixelBufferLockBaseAddress(pixel_buffer, 0);
+
+            let base_address = CVPixelBufferGetBaseAddress(pixel_buffer);
+
+            PixelBufferData {
+                buffer: pixel_buffer,
+                data: slice::from_raw_parts(
+                    base_address as *mut _,
+                    self.bytes_per_row * self.height,
+                ),
+            }
+        }
+    }
+
+    pub fn planes(&self) -> Vec<Plane> {
+        unsafe {
+            let pixel_buffer = sample_buffer_to_pixel_buffer(&self.buffer);
+            let count = CVPixelBufferGetPlaneCount(pixel_buffer);
+
+            CVPixelBufferLockBaseAddress(pixel_buffer, 0);
+
+            (0..count)
+                .map(|i| Plane {
+                    buffer: pixel_buffer,
+                    width: CVPixelBufferGetWidthOfPlane(pixel_buffer, i),
+                    height: CVPixelBufferGetHeightOfPlane(pixel_buffer, i),
+                    bytes_per_row: CVPixelBufferGetBytesPerRowOfPlane(pixel_buffer, i),
+                    index: i,
+                })
+                .collect()
+        }
+    }
+
+    pub(crate) fn new(item: ChannelItem) -> Option<Self> {
+        unsafe {
+            let display_time = pixel_buffer_display_time(&item.0);
+            let pixel_buffer = sample_buffer_to_pixel_buffer(&item.0);
+            let (width, height) = pixel_buffer_bounds(pixel_buffer);
+
+            match item.0.frame_status {
+                SCFrameStatus::Complete | SCFrameStatus::Started | SCFrameStatus::Idle => {
+                    Some(Self {
+                        display_time,
+                        width,
+                        height,
+                        bytes_per_row: pixel_buffer_bytes_per_row(pixel_buffer),
+                        buffer: item.0,
+                    })
+                }
+                _ => None,
+            }
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct Plane {
+    buffer: CVPixelBufferRef,
+    index: usize,
+    width: usize,
+    height: usize,
+    bytes_per_row: usize,
+}
+
+impl Plane {
+    pub fn width(&self) -> usize {
+        self.width
+    }
+
+    pub fn height(&self) -> usize {
+        self.height
+    }
+
+    pub fn bytes_per_row(&self) -> usize {
+        self.bytes_per_row
+    }
+
+    pub fn data(&self) -> PixelBufferData {
+        unsafe {
+            CVPixelBufferLockBaseAddress(self.buffer, 0);
+
+            let base_address = CVPixelBufferGetBaseAddressOfPlane(self.buffer, self.index);
+
+            PixelBufferData {
+                buffer: self.buffer,
+                data: slice::from_raw_parts(
+                    base_address as *mut _,
+                    self.bytes_per_row * self.height,
+                ),
+            }
+        }
+    }
+}
+
+pub struct PixelBufferData<'a> {
+    buffer: CVPixelBufferRef,
+    data: &'a [u8],
+}
+
+impl<'a> Deref for PixelBufferData<'a> {
+    type Target = [u8];
+
+    fn deref(&self) -> &'a Self::Target {
+        self.data
+    }
+}
+
+impl<'a> Drop for PixelBufferData<'a> {
+    fn drop(&mut self) {
+        unsafe { CVPixelBufferUnlockBaseAddress(self.buffer, 0) };
+    }
+}
+
+impl RawCapturer<'_> {
+    #[cfg(target_os = "macos")]
+    pub fn get_next_pixel_buffer(&self) -> Result<PixelBuffer, mpsc::RecvError> {
+        loop {
+            if let Some(frame) = PixelBuffer::new(self.capturer.rx.recv()?) {
+                return Ok(frame);
+            }
+        }
+    }
+}
+
+pub unsafe fn sample_buffer_to_pixel_buffer(sample_buffer: &CMSampleBuffer) -> CVPixelBufferRef {
+    let buffer_ref = &(*sample_buffer.sys_ref);
+    
+
+    CMSampleBufferGetImageBuffer(buffer_ref) as CVPixelBufferRef
+}
+
+pub unsafe fn pixel_buffer_bounds(pixel_buffer: CVPixelBufferRef) -> (usize, usize) {
+    let width = CVPixelBufferGetWidth(pixel_buffer);
+    let height = CVPixelBufferGetHeight(pixel_buffer);
+
+    (width, height)
+}
+
+pub unsafe fn pixel_buffer_bytes_per_row(pixel_buffer: CVPixelBufferRef) -> usize {
+    CVPixelBufferGetBytesPerRow(pixel_buffer)
+}
+
+pub unsafe fn pixel_buffer_display_time(sample_buffer: &CMSampleBuffer) -> u64 {
+    sample_buffer.sys_ref.get_presentation_timestamp().value as u64
+}