Skip to content

Commit f1ba095

Browse files
committed
- Implement first draft of async streamed frame capturing for macos
- change the capturing to use `Recorder` that as a async stream api - introduce `Frame` that holds the image and the `Instant` when captured - change some args to work with `impl AsRef<>` in order to avoid tight type coupling
1 parent b5e2561 commit f1ba095

File tree

12 files changed

+169
-102
lines changed

12 files changed

+169
-102
lines changed

Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ env_logger = "0.9.0"
2727
simplerand = "1.3.0"
2828
humantime = "2.1.0"
2929
smol = "1.2.5"
30+
flume = { version = "0.10", default-features = false }
3031

3132
[dependencies.clap]
3233
version = "3.1.9"

src/capture.rs

Lines changed: 33 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,17 @@
1-
use anyhow::{Context, Result};
2-
use image::save_buffer;
3-
use image::ColorType::Rgba8;
41
use std::borrow::Borrow;
5-
use std::ops::{Add, Sub};
62
use std::sync::mpsc::Receiver;
73
use std::sync::{Arc, Mutex};
84
use std::time::{Duration, Instant};
5+
6+
use anyhow::{Context, Result};
7+
use image::save_buffer;
8+
use image::ColorType::Rgba8;
9+
use smol::stream::block_on;
910
use tempfile::TempDir;
1011

12+
use crate::common::{Frame, Recorder};
1113
use crate::utils::{file_name_for, IMG_EXT};
12-
use crate::{ImageOnHeap, PlatformApi, WindowId};
14+
use crate::{Image, PlatformApi, WindowId};
1315

1416
/// captures screenshots as file on disk
1517
/// collects also the timecodes when they have been captured
@@ -22,60 +24,49 @@ pub fn capture_thread(
2224
tempdir: Arc<Mutex<TempDir>>,
2325
force_natural: bool,
2426
) -> Result<()> {
25-
let duration = Duration::from_millis(250);
26-
let start = Instant::now();
27-
let mut idle_duration = Duration::from_millis(0);
28-
let mut last_frame: Option<ImageOnHeap> = None;
27+
let recorder = Recorder::new(api, win_id, 4);
28+
let mut last_frame: Option<Frame> = None;
2929
let mut identical_frames = 0;
30-
let mut last_now = Instant::now();
31-
loop {
32-
// blocks for a timeout
33-
if rx.recv_timeout(duration).is_ok() {
34-
break;
35-
}
36-
let now = Instant::now();
37-
let effective_now = now.sub(idle_duration);
38-
let tc = effective_now.saturating_duration_since(start).as_millis();
39-
let image = api.capture_window_screenshot(win_id)?;
30+
let start = Instant::now();
31+
32+
for frame in block_on(recorder) {
4033
if !force_natural {
41-
if last_frame.is_some()
42-
&& image
43-
.samples
44-
.as_slice()
45-
.eq(last_frame.as_ref().unwrap().samples.as_slice())
46-
{
47-
identical_frames += 1;
48-
} else {
49-
identical_frames = 0;
34+
let image: &Image = frame.as_ref();
35+
if let Some(last_image) = last_frame.as_ref() {
36+
let last_image: &Image = last_image.as_ref();
37+
if image.samples.as_slice().eq(last_image.samples.as_slice()) {
38+
identical_frames += 1;
39+
} else {
40+
identical_frames = 0;
41+
}
5042
}
5143
}
5244

53-
if identical_frames > 0 {
54-
// let's track now the duration as idle
55-
idle_duration = idle_duration.add(now.duration_since(last_now));
56-
} else {
57-
if let Err(e) = save_frame(&image, tc, tempdir.lock().unwrap().borrow(), file_name_for)
58-
{
59-
eprintln!("{}", &e);
60-
return Err(e);
61-
}
45+
if identical_frames == 0 {
46+
let tc: &Instant = frame.as_ref();
47+
let tc = tc.duration_since(start).as_millis();
48+
save_image(&frame, tc, tempdir.lock().unwrap().borrow(), file_name_for)?;
6249
time_codes.lock().unwrap().push(tc);
63-
last_frame = Some(image);
64-
identical_frames = 0;
50+
last_frame = Some(frame);
51+
}
52+
53+
// when there is a message we should just stop
54+
if rx.recv_timeout(Duration::from_millis(1)).is_ok() {
55+
break;
6556
}
66-
last_now = now;
6757
}
6858

6959
Ok(())
7060
}
7161

7262
/// saves a frame as a tga file
73-
pub fn save_frame(
74-
image: &ImageOnHeap,
63+
pub fn save_image(
64+
image: impl AsRef<Image>,
7565
time_code: u128,
7666
tempdir: &TempDir,
7767
file_name_for: fn(&u128, &str) -> String,
7868
) -> Result<()> {
69+
let image = image.as_ref();
7970
save_buffer(
8071
tempdir.path().join(file_name_for(&time_code, IMG_EXT)),
8172
&image.samples,

src/common/frame.rs

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
use std::time::Instant;
2+
3+
use image::flat::SampleLayout;
4+
use image::{ColorType, FlatSamples};
5+
6+
use crate::common::image::{convert_bgra_to_rgba, crop, CropMut};
7+
use crate::{Image, Margin, Result};
8+
9+
pub struct Frame {
10+
image: Image,
11+
timecode: Instant,
12+
}
13+
14+
impl Frame {
15+
pub fn from_bgra(raw_data: Vec<u8>, channels: u8, width: u32, height: u32) -> Self {
16+
let timecode = Instant::now();
17+
let mut raw_data = raw_data;
18+
convert_bgra_to_rgba(&mut raw_data[..]);
19+
20+
let color = ColorType::Rgba8;
21+
let image = FlatSamples {
22+
samples: raw_data,
23+
layout: SampleLayout::row_major_packed(channels, width, height),
24+
color_hint: Some(color),
25+
};
26+
27+
Self { image, timecode }
28+
}
29+
}
30+
31+
impl AsRef<Image> for Frame {
32+
fn as_ref(&self) -> &Image {
33+
&self.image
34+
}
35+
}
36+
37+
impl CropMut for Frame {
38+
fn crop(&mut self, margin: &Margin) -> Result<()> {
39+
self.image = crop(&self, margin)?;
40+
41+
Ok(())
42+
}
43+
}
44+
45+
impl AsRef<Instant> for Frame {
46+
fn as_ref(&self) -> &Instant {
47+
&self.timecode
48+
}
49+
}

src/common/identify_transparency.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@ use image::{GenericImageView, Rgba};
66
/// this helps to identify outer transparent regions
77
/// since some backends provides transparency either from a compositor effect like drop shadow on ubuntu / GNOME
88
/// or some strange right side strip on MacOS
9-
pub fn identify_transparency(image: Image) -> Result<Option<Margin>> {
10-
let image: View<_, Rgba<u8>> = image.as_view()?;
9+
pub fn identify_transparency(image: impl AsRef<Image>) -> Result<Option<Margin>> {
10+
let image: View<_, Rgba<u8>> = image.as_ref().as_view()?;
1111
let (width, height) = image.dimensions();
1212
let half_width = width / 2;
1313
let half_height = height / 2;

src/common/image.rs

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,19 @@
1-
use crate::common::Margin;
2-
use crate::{Image, ImageOnHeap, Result};
31
use image::flat::View;
42
use image::{imageops, GenericImageView, ImageBuffer, Rgba};
53

4+
use crate::common::Margin;
5+
use crate::{Image, ImageOnHeap, Result};
6+
7+
pub trait CropMut {
8+
fn crop(&mut self, margin: &Margin) -> Result<()>;
9+
}
10+
611
///
712
/// specialized version of crop for [`ImageOnHeap`] and [`Margin`]
813
///
914
#[cfg_attr(not(macos), allow(dead_code))]
10-
pub fn crop(image: Image, margin: &Margin) -> Result<ImageOnHeap> {
11-
let mut img2: View<_, Rgba<u8>> = image.as_view()?;
15+
pub fn crop(image: impl AsRef<Image>, margin: &Margin) -> Result<ImageOnHeap> {
16+
let mut img2: View<_, Rgba<u8>> = image.as_ref().as_view()?;
1217
let (width, height) = (
1318
img2.width() - (margin.left + margin.right) as u32,
1419
img2.height() - (margin.top + margin.bottom) as u32,
@@ -28,7 +33,7 @@ pub fn crop(image: Image, margin: &Margin) -> Result<ImageOnHeap> {
2833
}
2934
}
3035

31-
Ok(Box::new(buf.into_flat_samples()))
36+
Ok(buf.into_flat_samples())
3237
}
3338

3439
///
@@ -40,19 +45,20 @@ pub fn convert_bgra_to_rgba(buffer: &mut [u8]) {
4045

4146
#[cfg(test)]
4247
mod tests {
43-
use super::*;
4448
use image::open;
4549

50+
use super::*;
51+
4652
#[test]
4753
fn should_crop() -> Result<()> {
4854
// given
4955
let image_org = open("tests/frames/frame-macos-right-side-issue.tga")?;
5056
let image = image_org.into_rgba8();
51-
let image_raw = ImageOnHeap::new(image.into_flat_samples());
57+
let image_raw = image.into_flat_samples();
5258
let (width, height) = (image_raw.layout.width, image_raw.layout.height);
5359

5460
// when
55-
let cropped = crop(*image_raw, &Margin::new(1, 1, 1, 1))?;
61+
let cropped = crop(image_raw, &Margin::new(1, 1, 1, 1))?;
5662

5763
// then
5864
assert_eq!(cropped.layout.width, width - 2);

src/common/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,12 @@ pub mod identify_transparency;
22
pub mod image;
33
pub mod utils;
44

5+
mod frame;
56
mod margin;
67
mod platform_api;
78
mod recorder;
89

10+
pub use frame::*;
911
pub use margin::*;
1012
pub use platform_api::*;
1113
pub use recorder::*;

src/common/platform_api.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
1-
use crate::{ImageOnHeap, Result, WindowId, WindowList};
1+
use crate::common::Frame;
2+
use crate::{Result, WindowId, WindowList};
23

34
pub trait PlatformApi: Send + Unpin + Sized {
45
/// 1. it does check for the screenshot
56
/// 2. it checks for transparent margins and configures the api
67
/// to cut them away in further screenshots
78
fn calibrate(&mut self, window_id: WindowId) -> Result<()>;
89
fn window_list(&self) -> Result<WindowList>;
9-
fn capture_window_screenshot(&self, window_id: WindowId) -> Result<ImageOnHeap>;
10+
fn capture_window_screenshot(&self, window_id: WindowId) -> Result<Frame>;
1011
fn get_active_window(&self) -> Result<WindowId>;
1112
}

src/common/recorder.rs

Lines changed: 44 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,18 @@
1-
use crate::{ImageOnHeap, PlatformApi, WindowId};
2-
3-
use smol::future::FutureExt;
41
use std::pin::Pin;
52
use std::task::{Context, Poll};
6-
use std::time::{Duration, Instant};
3+
use std::time::Duration;
74

5+
use smol::future::FutureExt;
86
use smol::stream::Stream;
97
use smol::Timer;
108

11-
struct Recorder<A: PlatformApi> {
9+
use crate::common::Frame;
10+
use crate::{PlatformApi, WindowId};
11+
12+
pub struct Recorder<A: PlatformApi> {
1213
api: A,
1314
window_id: WindowId,
1415
timer: Timer,
15-
last_frame_timestamp: Option<Instant>,
1616
}
1717

1818
impl<A: PlatformApi> Recorder<A> {
@@ -22,48 +22,72 @@ impl<A: PlatformApi> Recorder<A> {
2222
api,
2323
window_id,
2424
timer: Timer::interval(fps),
25-
last_frame_timestamp: None,
2625
}
2726
}
2827
}
2928

3029
impl<A: PlatformApi> Stream for Recorder<A> {
31-
type Item = ImageOnHeap;
30+
type Item = Frame;
3231

3332
fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
3433
let this = self.get_mut();
35-
let now = Instant::now();
3634
match this.timer.poll(_cx) {
37-
Poll::Ready(_) => {
38-
let d = now.duration_since(this.last_frame_timestamp.unwrap_or(now));
39-
dbg!(d);
40-
this.last_frame_timestamp.replace(now);
41-
Poll::Ready(this.api.capture_window_screenshot(this.window_id).ok())
42-
}
35+
Poll::Ready(_) => Poll::Ready(this.api.capture_window_screenshot(this.window_id).ok()),
4336
Poll::Pending => Poll::Pending,
4437
}
4538
}
4639
}
4740

4841
#[cfg(test)]
4942
#[cfg(target_os = "macos")]
43+
#[cfg(feature = "e2e_tests")]
5044
mod tests {
51-
use super::*;
52-
use crate::macos::*;
5345
use smol::stream::block_on;
5446

47+
use crate::macos::*;
48+
49+
use super::*;
50+
5551
#[test]
5652
fn should_record_not() {
57-
let api = setup().unwrap();
58-
let rec = Recorder::new(api, 682, 4);
53+
let win = 9123;
54+
let mut api = setup().unwrap();
55+
api.calibrate(win).unwrap();
56+
let rec = Recorder::new(api, win, 10);
5957

6058
let mut i = 0;
6159
for _img in block_on(rec) {
6260
i += 1;
63-
if i >= 4 {
61+
if i >= 5 {
6462
println!("Done with testing..");
6563
break;
6664
}
6765
}
6866
}
67+
68+
// #[test]
69+
// fn should_queue_frames_for_saving() {
70+
// let win = 9123;
71+
// let mut api = setup().unwrap();
72+
// api.calibrate(win).unwrap();
73+
// let mut rec = Recorder::new(api, win, 10);
74+
//
75+
// let future = rec.next();
76+
// {
77+
// let (sender, receiver) = flume::unbounded();
78+
//
79+
// // A function that schedules the task when it gets woken up.
80+
// let schedule = move |runnable| sender.send(runnable).unwrap();
81+
//
82+
// // Construct a task.
83+
// let (runnable, task) = async_task::spawn(future, schedule);
84+
//
85+
// // Push the task into the queue by invoking its schedule function.
86+
// runnable.schedule();
87+
//
88+
// for runnable in receiver {
89+
// runnable.run();
90+
// }
91+
// }
92+
// }
6993
}

0 commit comments

Comments
 (0)