display-ambient-light/desktop
1
0
Fork 0
Code Issues 1 Pull Requests Packages Projects 1 Releases Wiki Activity

feat: 使用 ScreenCaptureKit 获取屏幕帧数据。

Browse Source
This commit is contained in:
Ivan Li 2023-06-05 22:34:32 +08:00
parent ed72bdfdb1
commit 268ec1df81
7 changed files with 511 additions and 560 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

499
src-tauri/Cargo.lock generated
View File

File diff suppressed because it is too large Load Diff

1
src-tauri/Cargo.toml
View File

@ -37,6 +37,7 @@ mdns-sd = "0.7.2"
futures = "0.3.28"
ddc-hi = "0.4.1"
coreaudio-rs = "0.11.2"
rust_swift_screencapture = { version = "0.1.1", path = "../../../../demo/rust-swift-screencapture" }
[features]
# this feature is used for production builds or when `devPath` points to the filesystem

97
src-tauri/src/ambient_light/publisher.rs
View File

@ -1,4 +1,4 @@
use std::{collections::HashMap, sync::Arc, time::Duration};
use std::{collections::HashMap, sync::Arc, time::Duration, borrow::Borrow};
use paris::warn;
use tauri::async_runtime::RwLock;
@ -11,8 +11,9 @@ use tokio::{
use crate::{
ambient_light::{config, ConfigManager},
led_color::LedColor,
screenshot::LedSamplePoints,
screenshot_manager::{self, ScreenshotManager}, rpc::UdpRpc,
rpc::UdpRpc,
screenshot::{self, LedSamplePoints},
screenshot_manager::{self, ScreenshotManager},
};
use itertools::Itertools;
@ -48,60 +49,33 @@ impl LedColorsPublisher {
.await
}
fn start_one_display_colors_fetcher(
async fn start_one_display_colors_fetcher(
&self,
display_id: u32,
sample_points: Vec<Vec<LedSamplePoints>>,
sample_points: Vec<LedSamplePoints>,
bound_scale_factor: f32,
mappers: Vec<SamplePointMapper>,
display_colors_tx: broadcast::Sender<(u32, Vec<u8>)>,
) {
let internal_tasks_version = self.inner_tasks_version.clone();
let screenshot_manager = ScreenshotManager::global().await;
let screenshot_rx = screenshot_manager.subscribe_by_display_id(display_id).await;
if let Err(err) = screenshot_rx {
log::error!("{}", err);
return;
}
let mut screenshot_rx = screenshot_rx.unwrap();
tokio::spawn(async move {
let colors = screenshot_manager::get_display_colors(
display_id,
&sample_points,
bound_scale_factor,
);
if let Err(err) = colors {
warn!("Failed to get colors: {}", err);
return;
}
let mut interval = tokio::time::interval(Duration::from_millis(33));
let init_version = internal_tasks_version.read().await.clone();
loop {
interval.tick().await;
tokio::time::sleep(Duration::from_millis(1)).await;
let version = internal_tasks_version.read().await.clone();
if version != init_version {
log::info!(
"inner task version changed, stop. {} != {}",
internal_tasks_version.read().await.clone(),
init_version
);
break;
}
let colors = screenshot_manager::get_display_colors(
display_id,
&sample_points,
bound_scale_factor,
);
if let Err(err) = colors {
warn!("Failed to get colors: {}", err);
sleep(Duration::from_millis(100)).await;
continue;
}
let colors: Vec<crate::led_color::LedColor> = colors.unwrap();
while screenshot_rx.changed().await.is_ok() {
let screenshot = screenshot_rx.borrow().clone();
let colors = screenshot
.get_colors_by_sample_points(&sample_points)
.await;
let colors_copy = colors.clone();
@ -133,6 +107,18 @@ impl LedColorsPublisher {
warn!("Failed to send display_colors: {}", err);
}
};
// Check if the inner task version changed
let version = internal_tasks_version.read().await.clone();
if version != init_version {
log::info!(
"inner task version changed, stop. {} != {}",
internal_tasks_version.read().await.clone(),
init_version
);
break;
}
}
});
}
@ -247,13 +233,15 @@ impl LedColorsPublisher {
let display_id = sample_point_group.display_id;
let sample_points = sample_point_group.points;
let bound_scale_factor = sample_point_group.bound_scale_factor;
publisher.start_one_display_colors_fetcher(
display_id,
sample_points,
bound_scale_factor,
sample_point_group.mappers,
display_colors_tx.clone(),
);
publisher
.start_one_display_colors_fetcher(
display_id,
sample_points,
bound_scale_factor,
sample_point_group.mappers,
display_colors_tx.clone(),
)
.await;
}
let display_ids = configs.sample_point_groups;
@ -402,6 +390,7 @@ impl LedColorsPublisher {
let points: Vec<_> = led_strip_configs
.clone()
.map(|(_, config)| screenshot.get_sample_points(&config))
.flatten()
.collect();
if points.len() == 0 {
@ -451,7 +440,7 @@ pub struct AllColorConfig {
#[derive(Debug, Clone)]
pub struct DisplaySamplePointGroup {
pub display_id: u32,
pub points: Vec<Vec<LedSamplePoints>>,
pub points: Vec<LedSamplePoints>,
pub bound_scale_factor: f32,
pub mappers: Vec<config::SamplePointMapper>,
}

279
src-tauri/src/main.rs
View File

@ -88,7 +88,7 @@ async fn get_led_strips_sample_points(
let screenshot_manager = ScreenshotManager::global().await;
let channels = screenshot_manager.channels.read().await;
if let Some(rx) = channels.get(&config.display_id) {
let rx = rx.clone();
let rx = rx.read().await;
let screenshot = rx.borrow().clone();
let sample_points = screenshot.get_sample_points(&config);
Ok(sample_points)
@ -105,7 +105,7 @@ async fn get_one_edge_colors(
let screenshot_manager = ScreenshotManager::global().await;
let channels = screenshot_manager.channels.read().await;
if let Some(rx) = channels.get(&display_id) {
let rx = rx.clone();
let rx = rx.read().await;
let screenshot = rx.borrow().clone();
let bytes = screenshot.bytes.read().await.to_owned();
let colors =
@ -217,8 +217,12 @@ async fn get_displays() -> Vec<DisplayState> {
async fn main() {
env_logger::init();
let screenshot_manager = ScreenshotManager::global().await;
screenshot_manager.start().unwrap();
tokio::spawn(async move {
let screenshot_manager = ScreenshotManager::global().await;
screenshot_manager.start().await.unwrap_or_else(|e| {
error!("can not start screenshot manager: {}", e);
})
});
let led_color_publisher = ambient_light::LedColorsPublisher::global().await;
led_color_publisher.start();
@ -282,77 +286,86 @@ async fn main() {
let bytes = tokio::task::block_in_place(move || {
tauri::async_runtime::block_on(async move {
let screenshot_manager = ScreenshotManager::global().await;
let channels = screenshot_manager.channels.read().await;
if let Some(rx) = channels.get(&display_id) {
let rx = rx.clone();
let screenshot = rx.borrow().clone();
let bytes = screenshot.bytes.read().await;
let rx: Result<tokio::sync::watch::Receiver<Screenshot>, anyhow::Error> = screenshot_manager.subscribe_by_display_id(display_id).await;
let (scale_factor_x, scale_factor_y, width, height) = if url.query.is_some()
&& url.query.as_ref().unwrap().contains_key("height")
&& url.query.as_ref().unwrap().contains_key("width")
{
let width = url.query.as_ref().unwrap()["width"]
.parse::<u32>()
.map_err(|err| {
warn!("width parse error: {}", err);
err
})?;
let height = url.query.as_ref().unwrap()["height"]
.parse::<u32>()
.map_err(|err| {
warn!("height parse error: {}", err);
err
})?;
(
screenshot.width as f32 / width as f32,
screenshot.height as f32 / height as f32,
width,
height,
)
} else {
log::debug!("scale by scale_factor");
let scale_factor = screenshot.scale_factor;
(
scale_factor,
scale_factor,
(screenshot.width as f32 / scale_factor) as u32,
(screenshot.height as f32 / scale_factor) as u32,
)
};
log::debug!(
"scale by query. width: {}, height: {}, scale_factor: {}, len: {}",
if let Err(err) = rx {
anyhow::bail!("Display#{}: not found. {}", display_id, err);
}
let mut rx = rx.unwrap();
if rx.changed().await.is_err() {
anyhow::bail!("Display#{}: no more screenshot.", display_id);
}
let screenshot = rx.borrow().clone();
let bytes = screenshot.bytes.read().await;
if bytes.len() == 0 {
anyhow::bail!("Display#{}: no screenshot.", display_id);
}
log::debug!("Display#{}: screenshot size: {}", display_id, bytes.len());
let (scale_factor_x, scale_factor_y, width, height) = if url.query.is_some()
&& url.query.as_ref().unwrap().contains_key("height")
&& url.query.as_ref().unwrap().contains_key("width")
{
let width = url.query.as_ref().unwrap()["width"]
.parse::<u32>()
.map_err(|err| {
warn!("width parse error: {}", err);
err
})?;
let height = url.query.as_ref().unwrap()["height"]
.parse::<u32>()
.map_err(|err| {
warn!("height parse error: {}", err);
err
})?;
(
screenshot.width as f32 / width as f32,
screenshot.height as f32 / height as f32,
width,
height,
screenshot.width as f32 / width as f32,
width * height * 4,
);
let bytes_per_row = screenshot.bytes_per_row as f32;
let mut rgba_buffer = vec![0u8; (width * height * 4) as usize];
for y in 0..height {
for x in 0..width {
let offset = ((y as f32) * scale_factor_y).floor() as usize
* bytes_per_row as usize
+ ((x as f32) * scale_factor_x).floor() as usize * 4;
let b = bytes[offset];
let g = bytes[offset + 1];
let r = bytes[offset + 2];
let a = bytes[offset + 3];
let offset_2 = (y * width + x) as usize * 4;
rgba_buffer[offset_2] = r;
rgba_buffer[offset_2 + 1] = g;
rgba_buffer[offset_2 + 2] = b;
rgba_buffer[offset_2 + 3] = a;
}
}
Ok(rgba_buffer.clone())
)
} else {
anyhow::bail!("Display#{}: not found", display_id);
log::debug!("scale by scale_factor");
let scale_factor = screenshot.scale_factor;
(
scale_factor,
scale_factor,
(screenshot.width as f32 / scale_factor) as u32,
(screenshot.height as f32 / scale_factor) as u32,
)
};
log::debug!(
"scale by query. width: {}, height: {}, scale_factor: {}, len: {}",
width,
height,
screenshot.width as f32 / width as f32,
width * height * 4,
);
let bytes_per_row = screenshot.bytes_per_row as f32;
let mut rgba_buffer = vec![0u8; (width * height * 4) as usize];
for y in 0..height {
for x in 0..width {
let offset = ((y as f32) * scale_factor_y).floor() as usize
* bytes_per_row as usize
+ ((x as f32) * scale_factor_x).floor() as usize * 4;
let b = bytes[offset];
let g = bytes[offset + 1];
let r = bytes[offset + 2];
let a = bytes[offset + 3];
let offset_2 = (y * width + x) as usize * 4;
rgba_buffer[offset_2] = r;
rgba_buffer[offset_2 + 1] = g;
rgba_buffer[offset_2 + 2] = b;
rgba_buffer[offset_2 + 3] = a;
}
}
Ok(rgba_buffer.clone())
})
});
@ -388,82 +401,82 @@ async fn main() {
}
});
let app_handle = app.handle().clone();
tokio::spawn(async move {
let publisher = ambient_light::LedColorsPublisher::global().await;
let mut publisher_update_receiver = publisher.clone_sorted_colors_receiver().await;
loop {
if let Err(err) = publisher_update_receiver.changed().await {
error!("publisher update receiver changed error: {}", err);
return;
}
// let app_handle = app.handle().clone();
// tokio::spawn(async move {
// let publisher = ambient_light::LedColorsPublisher::global().await;
// let mut publisher_update_receiver = publisher.clone_sorted_colors_receiver().await;
// loop {
// if let Err(err) = publisher_update_receiver.changed().await {
// error!("publisher update receiver changed error: {}", err);
// return;
// }
let publisher = publisher_update_receiver.borrow().clone();
// let publisher = publisher_update_receiver.borrow().clone();
app_handle
.emit_all("led_sorted_colors_changed", publisher)
.unwrap();
}
});
// app_handle
// .emit_all("led_sorted_colors_changed", publisher)
// .unwrap();
// }
// });
let app_handle = app.handle().clone();
tokio::spawn(async move {
let publisher = ambient_light::LedColorsPublisher::global().await;
let mut publisher_update_receiver = publisher.clone_colors_receiver().await;
loop {
if let Err(err) = publisher_update_receiver.changed().await {
error!("publisher update receiver changed error: {}", err);
return;
}
// let app_handle = app.handle().clone();
// tokio::spawn(async move {
// let publisher = ambient_light::LedColorsPublisher::global().await;
// let mut publisher_update_receiver = publisher.clone_colors_receiver().await;
// loop {
// if let Err(err) = publisher_update_receiver.changed().await {
// error!("publisher update receiver changed error: {}", err);
// return;
// }
let publisher = publisher_update_receiver.borrow().clone();
// let publisher = publisher_update_receiver.borrow().clone();
app_handle
.emit_all("led_colors_changed", publisher)
.unwrap();
}
});
// app_handle
// .emit_all("led_colors_changed", publisher)
// .unwrap();
// }
// });
let app_handle = app.handle().clone();
tokio::spawn(async move {
loop {
match UdpRpc::global().await {
Ok(udp_rpc) => {
let mut receiver = udp_rpc.subscribe_boards_change();
loop {
if let Err(err) = receiver.changed().await {
error!("boards change receiver changed error: {}", err);
return;
}
// let app_handle = app.handle().clone();
// tokio::spawn(async move {
// loop {
// match UdpRpc::global().await {
// Ok(udp_rpc) => {
// let mut receiver = udp_rpc.subscribe_boards_change();
// loop {
// if let Err(err) = receiver.changed().await {
// error!("boards change receiver changed error: {}", err);
// return;
// }
let boards = receiver.borrow().clone();
// let boards = receiver.borrow().clone();
let boards = boards.into_iter().collect::<Vec<_>>();
// let boards = boards.into_iter().collect::<Vec<_>>();
app_handle.emit_all("boards_changed", boards).unwrap();
}
}
Err(err) => {
error!("udp rpc error: {}", err);
return;
}
}
}
});
// app_handle.emit_all("boards_changed", boards).unwrap();
// }
// }
// Err(err) => {
// error!("udp rpc error: {}", err);
// return;
// }
// }
// }
// });
let app_handle = app.handle().clone();
tokio::spawn(async move {
let display_manager = DisplayManager::global().await;
let mut rx =display_manager.subscribe_displays_changed();
// let app_handle = app.handle().clone();
// tokio::spawn(async move {
// let display_manager = DisplayManager::global().await;
// let mut rx = display_manager.subscribe_displays_changed();
while rx.changed().await.is_ok() {
let displays = rx.borrow().clone();
// while rx.changed().await.is_ok() {
// let displays = rx.borrow().clone();
log::info!("displays changed. emit displays_changed event.");
// log::info!("displays changed. emit displays_changed event.");
app_handle.emit_all("displays_changed", displays).unwrap();
}
});
// app_handle.emit_all("displays_changed", displays).unwrap();
// }
// });
Ok(())
})

22
src-tauri/src/screenshot.rs
View File

@ -1,5 +1,6 @@
use std::iter;
use std::fmt::Formatter;
use std::{iter, fmt::Debug};
use std::sync::Arc;
use serde::{Deserialize, Serialize};
@ -7,17 +8,30 @@ use tauri::async_runtime::RwLock;
use crate::{ambient_light::LedStripConfig, led_color::LedColor};
#[derive(Debug, Clone)]
#[derive(Clone)]
pub struct Screenshot {
pub display_id: u32,
pub height: u32,
pub width: u32,
pub bytes_per_row: usize,
pub bytes: Arc<RwLock<Vec<u8>>>,
pub bytes: Arc<RwLock<Arc<Vec<u8>>>>,
pub scale_factor: f32,
pub bound_scale_factor: f32,
}
impl Debug for Screenshot {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Screenshot")
.field("display_id", &self.display_id)
.field("height", &self.height)
.field("width", &self.width)
.field("bytes_per_row", &self.bytes_per_row)
.field("scale_factor", &self.scale_factor)
.field("bound_scale_factor", &self.bound_scale_factor)
.finish()
}
}
static SINGLE_AXIS_POINTS: usize = 5;
impl Screenshot {
@ -26,7 +40,7 @@ impl Screenshot {
height: u32,
width: u32,
bytes_per_row: usize,
bytes: Vec<u8>,
bytes: Arc<Vec<u8>>,
scale_factor: f32,
bound_scale_factor: f32,
) -> Self {

168
src-tauri/src/screenshot_manager.rs
View File

@ -5,47 +5,14 @@ use core_graphics::display::{
};
use core_graphics::geometry::{CGPoint, CGRect, CGSize};
use paris::warn;
use rust_swift_screencapture::display::CGDisplayId;
use tauri::async_runtime::RwLock;
use tokio::sync::{broadcast, watch, OnceCell};
use tokio::time::{self, Duration};
use tokio::sync::{broadcast, watch, Mutex, OnceCell};
use tokio::task::yield_now;
use crate::screenshot::LedSamplePoints;
use crate::{ambient_light::SamplePointMapper, led_color::LedColor, screenshot::Screenshot};
pub fn take_screenshot(display_id: u32, scale_factor: f32) -> anyhow::Result<Screenshot> {
log::debug!("take_screenshot");
let cg_display = CGDisplay::new(display_id);
let cg_image = CGDisplay::screenshot(
cg_display.bounds(),
kCGWindowListOptionOnScreenOnly,
kCGNullWindowID,
kCGWindowImageDefault,
)
.ok_or_else(|| anyhow::anyhow!("Display#{}: take screenshot failed", display_id))?;
let buffer = cg_image.data();
let bytes_per_row = cg_image.bytes_per_row();
let height = cg_image.height();
let width = cg_image.width();
let bytes = buffer.bytes().to_owned();
let cg_display = CGDisplay::new(display_id);
let bound_scale_factor = (cg_display.bounds().size.width / width as f64) as f32;
Ok(Screenshot::new(
display_id,
height as u32,
width as u32,
bytes_per_row,
bytes,
scale_factor,
bound_scale_factor,
))
}
pub fn get_display_colors(
display_id: u32,
sample_points: &Vec<Vec<LedSamplePoints>>,
@ -114,7 +81,7 @@ pub fn get_display_colors(
}
pub struct ScreenshotManager {
pub channels: Arc<RwLock<HashMap<u32, watch::Receiver<Screenshot>>>>,
pub channels: Arc<RwLock<HashMap<u32, Arc::<RwLock<watch::Sender<Screenshot>>>>>>,
merged_screenshot_tx: Arc<RwLock<broadcast::Sender<Screenshot>>>,
}
@ -134,75 +101,70 @@ impl ScreenshotManager {
.await
}
pub fn start(&self) -> anyhow::Result<()> {
pub async fn start(&self) -> anyhow::Result<()> {
let displays = display_info::DisplayInfo::all()?;
for display in displays {
self.start_one(display.id, display.scale_factor)?;
}
Ok(())
}
fn start_one(&self, display_id: u32, scale_factor: f32) -> anyhow::Result<()> {
let channels = self.channels.to_owned();
let merged_screenshot_tx = self.merged_screenshot_tx.clone();
tokio::spawn(async move {
let screenshot = take_screenshot(display_id, scale_factor);
if screenshot.is_err() {
warn!("take_screenshot_loop: {}", screenshot.err().unwrap());
return;
}
let mut interval = time::interval(Duration::from_millis(1000));
let screenshot = screenshot.unwrap();
let (screenshot_tx, screenshot_rx) = watch::channel(screenshot);
{
let channels = channels.clone();
let mut channels = channels.write().await;
channels.insert(display_id, screenshot_rx.clone());
}
let merged_screenshot_tx = merged_screenshot_tx.read().await.clone();
loop {
Self::take_screenshot_loop(
display_id,
scale_factor,
&screenshot_tx,
&merged_screenshot_tx,
)
.await;
interval.tick().await;
tokio::time::sleep(Duration::from_millis(1)).await;
}
let futures = displays.iter().map(|display| async {
self.start_one(display.id, display.scale_factor)
.await
.unwrap_or_else(|err| {
warn!("start_one failed: display_id: {}, err: {}", display.id, err);
});
});
futures::future::join_all(futures).await;
Ok(())
}
async fn take_screenshot_loop(
display_id: u32,
scale_factor: f32,
screenshot_tx: &watch::Sender<Screenshot>,
merged_screenshot_tx: &broadcast::Sender<Screenshot>,
) {
let screenshot = take_screenshot(display_id, scale_factor);
if let Ok(screenshot) = screenshot {
match merged_screenshot_tx.send(screenshot.clone()) {
Ok(_) => {
log::info!(
"take_screenshot_loop: merged_screenshot_tx.send success. display#{}",
display_id
);
}
Err(_) => {
}
async fn start_one(&self, display_id: u32, scale_factor: f32) -> anyhow::Result<()> {
let mut channels = self.channels.write().await;
let merged_screenshot_tx = self.merged_screenshot_tx.clone();
let display = rust_swift_screencapture::display::Display::new(display_id);
display.start_capture().await;
let mut frame_rx = display.subscribe_frame().await;
let (tx, _) = watch::channel(Screenshot::new(
display_id,
0,
0,
0,
Arc::new(vec![]),
scale_factor,
scale_factor,
));
let tx = Arc::new(RwLock::new(tx));
channels.insert(display_id, tx.clone());
drop(channels);
let tx_for_send = tx.read().await;
while frame_rx.changed().await.is_ok() {
let frame = frame_rx.borrow().clone();
let screenshot = Screenshot::new(
display_id,
frame.height as u32,
frame.width as u32,
frame.bytes_per_row as usize,
frame.bytes,
scale_factor,
scale_factor,
);
let merged_screenshot_tx = merged_screenshot_tx.write().await;
if let Err(err) = merged_screenshot_tx.send(screenshot.clone()) {
// log::warn!("merged_screenshot_tx.send failed: {}", err);
}
screenshot_tx.send(screenshot).unwrap();
// log::info!("take_screenshot_loop: send success. display#{}", display_id)
} else {
warn!("take_screenshot_loop: {}", screenshot.err().unwrap());
if let Err(err) = tx_for_send.send(screenshot.clone()) {
log::warn!("display {} screenshot_tx.send failed: {}", display_id, err);
} else {
log::debug!("screenshot: {:?}", screenshot);
}
yield_now().await;
}
Ok(())
}
pub fn get_sorted_colors(colors: &Vec<u8>, mappers: &Vec<SamplePointMapper>) -> Vec<u8> {
@ -257,4 +219,16 @@ impl ScreenshotManager {
pub async fn clone_merged_screenshot_rx(&self) -> broadcast::Receiver<Screenshot> {
self.merged_screenshot_tx.read().await.subscribe()
}
pub async fn subscribe_by_display_id(
&self,
display_id: CGDisplayId,
) -> anyhow::Result<watch::Receiver<Screenshot>> {
let channels = self.channels.read().await;
if let Some(tx) = channels.get(&display_id) {
Ok(tx.read().await.subscribe())
} else {
Err(anyhow::anyhow!("display_id: {} not found", display_id))
}
}
}

5
src-tauri/tauri.conf.json
View File

@ -28,7 +28,10 @@
"icons/icon.ico"
],
"identifier": "cc.ivanli.ambient-light.desktop",
"targets": "all"
"targets": "all",
"macOS": {
"minimumSystemVersion": "13"
}
},
"security": {
"csp": null