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" futures = "0.3.28"
ddc-hi = "0.4.1" ddc-hi = "0.4.1"
coreaudio-rs = "0.11.2" coreaudio-rs = "0.11.2"
rust_swift_screencapture = { version = "0.1.1", path = "../../../../demo/rust-swift-screencapture" }
[features] [features]
# this feature is used for production builds or when `devPath` points to the filesystem # this feature is used for production builds or when `devPath` points to the filesystem

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

157
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 screenshot_manager = ScreenshotManager::global().await;
let channels = screenshot_manager.channels.read().await; let channels = screenshot_manager.channels.read().await;
if let Some(rx) = channels.get(&config.display_id) { if let Some(rx) = channels.get(&config.display_id) {
let rx = rx.clone(); let rx = rx.read().await;
let screenshot = rx.borrow().clone(); let screenshot = rx.borrow().clone();
let sample_points = screenshot.get_sample_points(&config); let sample_points = screenshot.get_sample_points(&config);
Ok(sample_points) Ok(sample_points)
@ -105,7 +105,7 @@ async fn get_one_edge_colors(
let screenshot_manager = ScreenshotManager::global().await; let screenshot_manager = ScreenshotManager::global().await;
let channels = screenshot_manager.channels.read().await; let channels = screenshot_manager.channels.read().await;
if let Some(rx) = channels.get(&display_id) { if let Some(rx) = channels.get(&display_id) {
let rx = rx.clone(); let rx = rx.read().await;
let screenshot = rx.borrow().clone(); let screenshot = rx.borrow().clone();
let bytes = screenshot.bytes.read().await.to_owned(); let bytes = screenshot.bytes.read().await.to_owned();
let colors = let colors =
@ -217,8 +217,12 @@ async fn get_displays() -> Vec<DisplayState> {
async fn main() { async fn main() {
env_logger::init(); env_logger::init();
tokio::spawn(async move {
let screenshot_manager = ScreenshotManager::global().await; let screenshot_manager = ScreenshotManager::global().await;
screenshot_manager.start().unwrap(); screenshot_manager.start().await.unwrap_or_else(|e| {
error!("can not start screenshot manager: {}", e);
})
});
let led_color_publisher = ambient_light::LedColorsPublisher::global().await; let led_color_publisher = ambient_light::LedColorsPublisher::global().await;
led_color_publisher.start(); led_color_publisher.start();
@ -282,11 +286,23 @@ async fn main() {
let bytes = tokio::task::block_in_place(move || { let bytes = tokio::task::block_in_place(move || {
tauri::async_runtime::block_on(async move { tauri::async_runtime::block_on(async move {
let screenshot_manager = ScreenshotManager::global().await; let screenshot_manager = ScreenshotManager::global().await;
let channels = screenshot_manager.channels.read().await; let rx: Result<tokio::sync::watch::Receiver<Screenshot>, anyhow::Error> = screenshot_manager.subscribe_by_display_id(display_id).await;
if let Some(rx) = channels.get(&display_id) {
let rx = rx.clone(); 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 screenshot = rx.borrow().clone();
let bytes = screenshot.bytes.read().await; 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() 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("height")
@ -350,9 +366,6 @@ async fn main() {
} }
Ok(rgba_buffer.clone()) Ok(rgba_buffer.clone())
} else {
anyhow::bail!("Display#{}: not found", display_id);
}
}) })
}); });
@ -388,82 +401,82 @@ async fn main() {
} }
}); });
let app_handle = app.handle().clone(); // let app_handle = app.handle().clone();
tokio::spawn(async move { // tokio::spawn(async move {
let publisher = ambient_light::LedColorsPublisher::global().await; // let publisher = ambient_light::LedColorsPublisher::global().await;
let mut publisher_update_receiver = publisher.clone_sorted_colors_receiver().await; // let mut publisher_update_receiver = publisher.clone_sorted_colors_receiver().await;
loop { // loop {
if let Err(err) = publisher_update_receiver.changed().await { // if let Err(err) = publisher_update_receiver.changed().await {
error!("publisher update receiver changed error: {}", err); // error!("publisher update receiver changed error: {}", err);
return; // return;
} // }
let publisher = publisher_update_receiver.borrow().clone(); // let publisher = publisher_update_receiver.borrow().clone();
app_handle // app_handle
.emit_all("led_sorted_colors_changed", publisher) // .emit_all("led_sorted_colors_changed", publisher)
.unwrap(); // .unwrap();
} // }
}); // });
let app_handle = app.handle().clone(); // let app_handle = app.handle().clone();
tokio::spawn(async move { // tokio::spawn(async move {
let publisher = ambient_light::LedColorsPublisher::global().await; // let publisher = ambient_light::LedColorsPublisher::global().await;
let mut publisher_update_receiver = publisher.clone_colors_receiver().await; // let mut publisher_update_receiver = publisher.clone_colors_receiver().await;
loop { // loop {
if let Err(err) = publisher_update_receiver.changed().await { // if let Err(err) = publisher_update_receiver.changed().await {
error!("publisher update receiver changed error: {}", err); // error!("publisher update receiver changed error: {}", err);
return; // return;
} // }
let publisher = publisher_update_receiver.borrow().clone(); // let publisher = publisher_update_receiver.borrow().clone();
app_handle // app_handle
.emit_all("led_colors_changed", publisher) // .emit_all("led_colors_changed", publisher)
.unwrap(); // .unwrap();
} // }
}); // });
let app_handle = app.handle().clone(); // let app_handle = app.handle().clone();
tokio::spawn(async move { // tokio::spawn(async move {
loop { // loop {
match UdpRpc::global().await { // match UdpRpc::global().await {
Ok(udp_rpc) => { // Ok(udp_rpc) => {
let mut receiver = udp_rpc.subscribe_boards_change(); // let mut receiver = udp_rpc.subscribe_boards_change();
loop { // loop {
if let Err(err) = receiver.changed().await { // if let Err(err) = receiver.changed().await {
error!("boards change receiver changed error: {}", err); // error!("boards change receiver changed error: {}", err);
return; // 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(); // app_handle.emit_all("boards_changed", boards).unwrap();
} // }
} // }
Err(err) => { // Err(err) => {
error!("udp rpc error: {}", err); // error!("udp rpc error: {}", err);
return; // return;
} // }
} // }
} // }
}); // });
let app_handle = app.handle().clone(); // let app_handle = app.handle().clone();
tokio::spawn(async move { // tokio::spawn(async move {
let display_manager = DisplayManager::global().await; // let display_manager = DisplayManager::global().await;
let mut rx =display_manager.subscribe_displays_changed(); // let mut rx = display_manager.subscribe_displays_changed();
while rx.changed().await.is_ok() { // while rx.changed().await.is_ok() {
let displays = rx.borrow().clone(); // 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(()) 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 std::sync::Arc;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
@ -7,17 +8,30 @@ use tauri::async_runtime::RwLock;
use crate::{ambient_light::LedStripConfig, led_color::LedColor}; use crate::{ambient_light::LedStripConfig, led_color::LedColor};
#[derive(Debug, Clone)] #[derive(Clone)]
pub struct Screenshot { pub struct Screenshot {
pub display_id: u32, pub display_id: u32,
pub height: u32, pub height: u32,
pub width: u32, pub width: u32,
pub bytes_per_row: usize, pub bytes_per_row: usize,
pub bytes: Arc<RwLock<Vec<u8>>>, pub bytes: Arc<RwLock<Arc<Vec<u8>>>>,
pub scale_factor: f32, pub scale_factor: f32,
pub bound_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; static SINGLE_AXIS_POINTS: usize = 5;
impl Screenshot { impl Screenshot {
@ -26,7 +40,7 @@ impl Screenshot {
height: u32, height: u32,
width: u32, width: u32,
bytes_per_row: usize, bytes_per_row: usize,
bytes: Vec<u8>, bytes: Arc<Vec<u8>>,
scale_factor: f32, scale_factor: f32,
bound_scale_factor: f32, bound_scale_factor: f32,
) -> Self { ) -> Self {

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

@ -5,47 +5,14 @@ use core_graphics::display::{
}; };
use core_graphics::geometry::{CGPoint, CGRect, CGSize}; use core_graphics::geometry::{CGPoint, CGRect, CGSize};
use paris::warn; use paris::warn;
use rust_swift_screencapture::display::CGDisplayId;
use tauri::async_runtime::RwLock; use tauri::async_runtime::RwLock;
use tokio::sync::{broadcast, watch, OnceCell}; use tokio::sync::{broadcast, watch, Mutex, OnceCell};
use tokio::time::{self, Duration}; use tokio::task::yield_now;
use crate::screenshot::LedSamplePoints; use crate::screenshot::LedSamplePoints;
use crate::{ambient_light::SamplePointMapper, led_color::LedColor, screenshot::Screenshot}; 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( pub fn get_display_colors(
display_id: u32, display_id: u32,
sample_points: &Vec<Vec<LedSamplePoints>>, sample_points: &Vec<Vec<LedSamplePoints>>,
@ -114,7 +81,7 @@ pub fn get_display_colors(
} }
pub struct ScreenshotManager { 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>>>, merged_screenshot_tx: Arc<RwLock<broadcast::Sender<Screenshot>>>,
} }
@ -134,75 +101,70 @@ impl ScreenshotManager {
.await .await
} }
pub fn start(&self) -> anyhow::Result<()> { pub async fn start(&self) -> anyhow::Result<()> {
let displays = display_info::DisplayInfo::all()?; 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 futures = displays.iter().map(|display| async {
let channels = self.channels.to_owned(); self.start_one(display.id, display.scale_factor)
let merged_screenshot_tx = self.merged_screenshot_tx.clone(); .await
tokio::spawn(async move { .unwrap_or_else(|err| {
let screenshot = take_screenshot(display_id, scale_factor); warn!("start_one failed: display_id: {}, err: {}", display.id, err);
});
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;
}
}); });
futures::future::join_all(futures).await;
Ok(()) Ok(())
} }
async fn take_screenshot_loop( async fn start_one(&self, display_id: u32, scale_factor: f32) -> anyhow::Result<()> {
display_id: u32, let mut channels = self.channels.write().await;
scale_factor: f32, let merged_screenshot_tx = self.merged_screenshot_tx.clone();
screenshot_tx: &watch::Sender<Screenshot>, let display = rust_swift_screencapture::display::Display::new(display_id);
merged_screenshot_tx: &broadcast::Sender<Screenshot>,
) { display.start_capture().await;
let screenshot = take_screenshot(display_id, scale_factor);
if let Ok(screenshot) = screenshot { let mut frame_rx = display.subscribe_frame().await;
match merged_screenshot_tx.send(screenshot.clone()) {
Ok(_) => { let (tx, _) = watch::channel(Screenshot::new(
log::info!( display_id,
"take_screenshot_loop: merged_screenshot_tx.send success. display#{}", 0,
display_id 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);
} }
Err(_) => { if let Err(err) = tx_for_send.send(screenshot.clone()) {
} log::warn!("display {} screenshot_tx.send failed: {}", display_id, err);
}
screenshot_tx.send(screenshot).unwrap();
// log::info!("take_screenshot_loop: send success. display#{}", display_id)
} else { } else {
warn!("take_screenshot_loop: {}", screenshot.err().unwrap()); log::debug!("screenshot: {:?}", screenshot);
} }
yield_now().await;
}
Ok(())
} }
pub fn get_sorted_colors(colors: &Vec<u8>, mappers: &Vec<SamplePointMapper>) -> Vec<u8> { 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> { pub async fn clone_merged_screenshot_rx(&self) -> broadcast::Receiver<Screenshot> {
self.merged_screenshot_tx.read().await.subscribe() 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" "icons/icon.ico"
], ],
"identifier": "cc.ivanli.ambient-light.desktop", "identifier": "cc.ivanli.ambient-light.desktop",
"targets": "all" "targets": "all",
"macOS": {
"minimumSystemVersion": "13"
}
}, },
"security": { "security": {
"csp": null "csp": null