Skip to content

Commit

Permalink
Merge pull request #127 from CapSoftware/raw-pixel-buffers
Browse files Browse the repository at this point in the history
raw pixel buffers
  • Loading branch information
Brendonovich authored Nov 7, 2024
2 parents f63eb09 + 88d0e14 commit b1e7ecb
Show file tree
Hide file tree
Showing 8 changed files with 343 additions and 109 deletions.
2 changes: 2 additions & 0 deletions src/capturer/engine/mac/apple_sys.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
#![allow(non_upper_case_globals)]

pub use screencapturekit_sys::os_types::base::*;

#[repr(C)]
Expand Down
126 changes: 63 additions & 63 deletions src/capturer/engine/mac/mod.rs
Original file line number Diff line number Diff line change
@@ -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},
Expand All @@ -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;
Expand All @@ -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
Expand Down Expand Up @@ -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();

Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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,
},
}
})
Expand All @@ -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
}
191 changes: 191 additions & 0 deletions src/capturer/engine/mac/pixel_buffer.rs
Original file line number Diff line number Diff line change
@@ -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
}
Loading

0 comments on commit b1e7ecb

Please sign in to comment.