comfy_wgpu/
screenshot.rs

1use chrono::Utc;
2use comfy_core::chrono::{DateTime, Timelike};
3use image::RgbaImage;
4
5use crate::*;
6
7pub struct ScreenshotItem {
8    pub image: RgbaImage,
9    pub time: DateTime<Utc>,
10}
11
12#[derive(Copy, Clone, Debug)]
13pub struct ScreenshotParams {
14    pub record_screenshots: bool,
15    /// When set to 1, a screenshot will be taken every frame.
16    /// When set to a higher number, a screenshot will be taken every n frames.
17    pub screenshot_interval_n: usize,
18    pub history_length: usize,
19
20    counter: usize,
21}
22
23impl Default for ScreenshotParams {
24    fn default() -> Self {
25        Self {
26            record_screenshots: false,
27            screenshot_interval_n: 1,
28            history_length: 10,
29            counter: 0,
30        }
31    }
32}
33
34#[cfg(target_arch = "wasm32")]
35pub fn record_screenshot_history(
36    _screen: UVec2,
37    _context: &GraphicsContext,
38    _screenshot_buffer: &SizedBuffer,
39    _output: &wgpu::SurfaceTexture,
40    _params: &mut ScreenshotParams,
41    _screenshot_history_buffer: &mut VecDeque<ScreenshotItem>,
42) {
43}
44
45pub fn save_screenshots_to_folder(
46    folder: &str,
47    screenshot_history_buffer: &VecDeque<ScreenshotItem>,
48) {
49    if std::fs::create_dir_all(&folder).log_err_ok().is_none() {
50        return;
51    }
52
53    for (i, screenshot) in screenshot_history_buffer.iter().enumerate() {
54        let t = &screenshot.time;
55
56        let name = format!(
57            "{}/image {} {}-{}-{}.{} [{}].png",
58            &folder,
59            t.date_naive(),
60            t.hour(),
61            t.minute(),
62            t.second(),
63            t.timestamp_subsec_millis(),
64            i,
65        );
66
67        screenshot.image.save(name).log_err_ok();
68    }
69}
70
71#[cfg(not(target_arch = "wasm32"))]
72pub fn record_screenshot_history(
73    screen: UVec2,
74    context: &GraphicsContext,
75    screenshot_buffer: &SizedBuffer,
76    output: &wgpu::SurfaceTexture,
77    params: &mut ScreenshotParams,
78    screenshot_history_buffer: &mut VecDeque<ScreenshotItem>,
79) {
80    let start_time = Instant::now();
81
82    {
83        let mut encoder = context.device.create_command_encoder(
84            &wgpu::CommandEncoderDescriptor {
85                label: Some("Copy output texture Encoder"),
86            },
87        );
88
89        encoder.copy_texture_to_buffer(
90            wgpu::ImageCopyTexture {
91                texture: &output.texture,
92                mip_level: 0,
93                origin: wgpu::Origin3d::ZERO,
94                aspect: wgpu::TextureAspect::All,
95            },
96            wgpu::ImageCopyBuffer {
97                buffer: &screenshot_buffer.buffer,
98                layout: wgpu::ImageDataLayout {
99                    offset: 0,
100                    bytes_per_row: Some(
101                        std::mem::size_of::<u32>() as u32 * screen.x,
102                    ),
103                    rows_per_image: Some(screen.y),
104                },
105            },
106            output.texture.size(),
107        );
108
109        context.queue.submit(std::iter::once(encoder.finish()));
110    }
111
112    let screenshot_image = pollster::block_on(async {
113        let buffer_slice = screenshot_buffer.buffer.slice(..);
114
115        let (tx, rx) = futures_intrusive::channel::shared::oneshot_channel();
116
117        buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
118            tx.send(result).unwrap();
119        });
120
121        context.device.poll(wgpu::Maintain::Wait);
122        rx.receive().await.unwrap().unwrap();
123
124        let data = buffer_slice.get_mapped_range();
125
126        let mut rgba_data: Vec<u8> = Vec::with_capacity(data.len());
127
128        for chunk in data.chunks_exact(4) {
129            let b = chunk[0];
130            let g = chunk[1];
131            let r = chunk[2];
132            let a = chunk[3];
133
134            rgba_data.push(r);
135            rgba_data.push(g);
136            rgba_data.push(b);
137            rgba_data.push(a);
138        }
139
140        let buffer = image::ImageBuffer::<image::Rgba<u8>, _>::from_raw(
141            screen.x, screen.y, rgba_data,
142        )
143        .unwrap();
144
145        buffer
146    });
147
148    if params.counter % params.screenshot_interval_n == 0 {
149        if screenshot_history_buffer.len() == params.history_length {
150            screenshot_history_buffer.pop_front();
151        }
152
153        screenshot_history_buffer.push_back(ScreenshotItem {
154            image: screenshot_image,
155            time: Utc::now(),
156        });
157    }
158
159    params.counter += 1;
160
161    perf_counter(
162        "screenshots in buffer",
163        screenshot_history_buffer.len() as u64,
164    );
165    perf_counter("screenshot time", start_time.elapsed().as_micros() as u64);
166
167    screenshot_buffer.buffer.unmap();
168}
169
170#[cfg(feature = "record-pngs")]
171pub fn record_pngs(
172    screen: UVec2,
173    context: &GraphicsContext,
174    screenshot_buffer: &SizedBuffer,
175    output: &wgpu::SurfaceTexture,
176) {
177    {
178        let mut encoder = context.device.create_command_encoder(
179            &wgpu::CommandEncoderDescriptor {
180                label: Some("Copy output texture Encoder"),
181            },
182        );
183
184        encoder.copy_texture_to_buffer(
185            wgpu::ImageCopyTexture {
186                texture: &output.texture,
187                mip_level: 0,
188                origin: wgpu::Origin3d::ZERO,
189                aspect: wgpu::TextureAspect::All,
190            },
191            wgpu::ImageCopyBuffer {
192                buffer: &screenshot_buffer.buffer,
193                layout: wgpu::ImageDataLayout {
194                    offset: 0,
195                    bytes_per_row: Some(
196                        std::mem::size_of::<u32>() as u32 * screen.x,
197                    ),
198                    rows_per_image: Some(screen.y),
199                },
200            },
201            output.texture.size(),
202        );
203
204        context.queue.submit(std::iter::once(encoder.finish()));
205    }
206
207    pollster::block_on(async {
208        let buffer_slice = screenshot_buffer.buffer.slice(..);
209
210        let (tx, rx) = futures_intrusive::channel::shared::oneshot_channel();
211
212        buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
213            tx.send(result).unwrap();
214        });
215
216        context.device.poll(wgpu::Maintain::Wait);
217        rx.receive().await.unwrap().unwrap();
218
219        let data = buffer_slice.get_mapped_range();
220
221        let path = std::env::current_exe().unwrap();
222        let example_name = path.file_stem().unwrap().to_string_lossy();
223
224        let images_dir = format!("target/images/{}", example_name);
225
226        let videos_dir = "target/videos";
227        let screenshots_dir = "target/screenshots";
228
229        std::fs::create_dir_all(&images_dir).unwrap();
230        std::fs::create_dir_all(&videos_dir).unwrap();
231        std::fs::create_dir_all(&screenshots_dir).unwrap();
232
233        let name = format!("{}/image-{:03}.png", &images_dir, get_frame());
234
235        {
236            let mut rgba_data: Vec<u8> = Vec::with_capacity(data.len());
237
238            for chunk in data.chunks_exact(4) {
239                let b = chunk[0];
240                let g = chunk[1];
241                let r = chunk[2];
242                let a = chunk[3];
243
244                rgba_data.push(r);
245                rgba_data.push(g);
246                rgba_data.push(b);
247                rgba_data.push(a);
248            }
249
250
251            let buffer = image::ImageBuffer::<image::Rgba<u8>, _>::from_raw(
252                screen.x, screen.y, rgba_data,
253            )
254            .unwrap();
255
256            let resized = image::imageops::resize(
257                &buffer,
258                screen.x / 3,
259                screen.y / 3,
260                image::imageops::FilterType::Nearest,
261            );
262
263            resized.save(name).unwrap();
264
265            if get_frame() > 60 {
266                resized
267                    .save(format!("{}/{}.png", screenshots_dir, example_name))
268                    .unwrap();
269
270                let ffmpeg_command = "ffmpeg";
271                let framerate = "30";
272                let input_pattern = "image-%03d.png";
273                let output_file =
274                    format!("{}/{}.webm", videos_dir, example_name);
275
276                let output = std::process::Command::new(ffmpeg_command)
277                    .arg("-framerate")
278                    .arg(framerate)
279                    .arg("-y")
280                    .arg("-i")
281                    .arg(format!("{}/{}", images_dir, input_pattern))
282                    .arg(output_file)
283                    .output()
284                    .expect("Failed to execute ffmpeg command");
285
286                if output.status.success() {
287                    println!("Successfully generated the GIF.");
288                } else {
289                    println!("Error generating the GIF:");
290                    println!(
291                        "stdout: {}",
292                        String::from_utf8_lossy(&output.stdout)
293                    );
294                    println!(
295                        "stderr: {}",
296                        String::from_utf8_lossy(&output.stderr)
297                    );
298                }
299
300                std::process::exit(0);
301            }
302        }
303    });
304
305    screenshot_buffer.buffer.unmap();
306}