stereokit_rust/tools/
screenshot.rs

1use std::{
2    env::{current_dir, set_current_dir},
3    fs::File,
4    io::{Read, Write},
5    sync::Mutex,
6};
7
8use stereokit_macros::IStepper;
9
10use crate::{
11    maths::{Pose, Quat, Vec2, Vec3, units::CM},
12    prelude::*,
13    system::Renderer,
14    tex::{Tex, TexFormat},
15    ui::Ui,
16    util::{PickerMode, Platform},
17};
18
19use crate::sprite::Sprite;
20
21use super::{
22    file_browser::{FILE_BROWSER_OPEN, FILE_BROWSER_SAVE, FileBrowser},
23    os_api::get_external_path,
24};
25
26/// Somewhere to store the selected filename
27static FILE_NAME: Mutex<String> = Mutex::new(String::new());
28
29pub const SHOW_SCREENSHOT_WINDOW: &str = "Tool_ShowScreenshotWindow";
30pub const SCREENSHOT_FORMATS: [&str; 2] = [".raw", ".rgba"];
31pub const CAPTURE_TEXTURE_ID: &str = "Uniq_ScreenshotTexture";
32const BROWSER_SUFFIX: &str = "_file_browser";
33
34/// A simple screenshot viewer to take / save / display screenshots.
35/// ### Fields that can be changed before initialization:
36/// * `picture_size` - The size of the picture to take. Default is Vec2::new(800.0, 600.0).
37/// * `field_of_view` - The field of view of the camera. Default is 90.0.
38/// * `windows_pose` - The initial pose of the window.
39/// * `window_size` - The size of the window. Default is Vec2::new(42.0, 37.0) * CM.
40/// * `enabled` - If the screenshot viewer is enabled at start. Default is `true`
41///
42/// ### Events this stepper is listening to:
43/// * `SHOW_SCREENSHOT_WINDOW` - Event that triggers when the window is visible ("true") or hidden ("false").
44/// * `FILE_BROWSER_OPEN` - Event that triggers when a file as been selected with the file browser. You can use this
45///   event too if you want to load a screenshot.
46///
47/// ### Examples
48/// ```
49/// # stereokit_rust::test_init_sk!(); // !!!! Get a proper way to initialize sk !!!!
50/// use stereokit_rust::{maths::Vec3, sk::SkInfo, ui::Ui,
51///                      tools::{file_browser::FILE_BROWSER_OPEN,
52///                              screenshot::{ScreenshotViewer, SHOW_SCREENSHOT_WINDOW}}};
53///
54/// let sk_info  = Some(sk.get_sk_info_clone());
55///
56/// let mut screenshot_viewer = ScreenshotViewer::default();
57/// screenshot_viewer.window_pose = Ui::popup_pose([0.0, 0.15, 1.3]);
58/// sk.send_event(StepperAction::add("ScrViewer", screenshot_viewer));
59///
60/// let screenshot_path = std::env::current_dir().unwrap().join("assets/textures/screenshot.raw");
61/// assert!(screenshot_path.exists());
62/// let scr_file = screenshot_path.to_str().expect("String should be valid");
63///
64/// number_of_steps = 4;
65/// filename_scr = "screenshots/screenshot_viewer.jpeg";
66/// test_screenshot!( // !!!! Get a proper main loop !!!!
67///     if iter == 0 {
68///        sk.send_event(StepperAction::event( "main", SHOW_SCREENSHOT_WINDOW,"false",));
69///     } else if iter == 1 {
70///        sk.send_event(StepperAction::event( "main", SHOW_SCREENSHOT_WINDOW,"true",));
71///        // The image is not visible at the next step, but at the step after.
72///        sk.send_event(StepperAction::event( "ScrViewer", FILE_BROWSER_OPEN, scr_file));
73///     }
74/// );
75/// ```
76/// <img src="https://raw.githubusercontent.com/mvvvv/StereoKit-rust/refs/heads/master/screenshots/screenshot_viewer.jpeg" alt="screenshot" width="200">
77#[derive(IStepper)]
78pub struct ScreenshotViewer {
79    id: StepperId,
80    sk_info: Option<Rc<RefCell<SkInfo>>>,
81    pub enabled: bool,
82    shutdown_completed: bool,
83
84    pub picture_size: Vec2,
85    pub field_of_view: f32,
86    pub window_pose: Pose,
87    pub window_size: Vec2,
88    tex: Tex,
89    screen: Option<Sprite>,
90}
91
92unsafe impl Send for ScreenshotViewer {}
93
94impl Default for ScreenshotViewer {
95    fn default() -> Self {
96        let picture_size = Vec2::new(800.0, 600.0);
97        let tex = Tex::default();
98
99        Self {
100            id: "ScreenshotStepper".to_string(),
101            sk_info: None,
102            enabled: false,
103            shutdown_completed: false,
104
105            picture_size,
106            field_of_view: 90.0,
107            window_pose: Pose::new(Vec3::new(-0.7, 1.0, -0.3), Some(Quat::look_dir(Vec3::new(1.0, 0.0, 1.0)))),
108            window_size: Vec2::new(42.0, 37.0) * CM,
109            tex,
110            screen: None,
111        }
112    }
113}
114
115impl ScreenshotViewer {
116    /// Called from IStepper::initialize here you can abort the initialization by returning false
117    fn start(&mut self) -> bool {
118        // self.tex = Tex::gen_color(
119        //     Color128::WHITE,
120        //     self.picture_size.x as i32,
121        //     self.picture_size.y as i32,
122        //     TexType::Rendertarget,
123        //     TexFormat::RGBA32,
124        // );
125        self.tex = Tex::render_target(
126            self.picture_size.x as usize,
127            self.picture_size.y as usize,
128            None,
129            Some(TexFormat::RGBA32),
130            Some(TexFormat::Depth32),
131        )
132        .unwrap_or_default();
133        self.tex.id(CAPTURE_TEXTURE_ID);
134        true
135    }
136
137    /// Called from IStepper::step, here you can check the event report
138    fn check_event(&mut self, id: &StepperId, key: &str, value: &str) {
139        if key.eq(SHOW_SCREENSHOT_WINDOW) {
140            self.enabled = value.parse().unwrap_or(false);
141            if !self.enabled {
142                self.close_file_browser()
143            }
144        } else if id == &self.id {
145            if key.eq(FILE_BROWSER_OPEN) {
146                let mut file_name = FILE_NAME.lock().unwrap();
147                file_name.clear();
148                file_name.push_str(value);
149                self.screen = None;
150            } else if key.eq(FILE_BROWSER_SAVE) {
151                save_screenshot(value);
152            }
153        }
154    }
155
156    /// Called from IStepper::step after check_event, here you can draw your UI and scene
157    fn draw(&mut self, token: &MainThreadToken) {
158        if !self.enabled {
159            return;
160        };
161
162        Ui::window_begin("Screenshot", &mut self.window_pose, Some(self.window_size), None, None);
163        if let Some(sprite) = &self.screen {
164            Ui::image(sprite, Vec2::new(0.4, 0.3));
165        } else {
166            Ui::vspace(30.0 * CM);
167            let mut file_name_lock = FILE_NAME.lock().unwrap();
168            let file_name = file_name_lock.to_string();
169            if !file_name.is_empty() {
170                if let Ok(mut file) = File::open(&file_name) {
171                    if let Ok(mut tex) = Tex::find(CAPTURE_TEXTURE_ID) {
172                        let mut buf = [0u8; 12];
173                        if file.read_exact(&mut buf).is_ok() {
174                            // Vive le format RGBA !!! https://github.com/bzotto/rgba_bitmap
175                            let rgba_tag = format!("{:?}", &buf[0..4]);
176                            let mut four_u8 = [0u8; 4];
177                            four_u8.copy_from_slice(&buf[4..8]);
178                            let width = u32::from_be_bytes(four_u8) as usize;
179                            four_u8.copy_from_slice(&buf[8..12]);
180                            let height = u32::from_be_bytes(four_u8) as usize;
181                            Log::diag(format!("RGBA file {} with size is {}x{}", &file_name, width, height));
182                            if rgba_tag != "RGBA" {
183                                let mut data = vec![];
184                                match file.read_to_end(&mut data) {
185                                    Ok(mut _size) => {
186                                        let data_slice = data.as_slice();
187                                        tex.set_colors_u8(width, height, data_slice, 4);
188                                        self.screen = Sprite::from_tex(&self.tex, None, None).ok();
189                                    }
190                                    Err(err) => {
191                                        Log::warn(format!("Screenshoot Error when reading file {file_name} : {err:?}"))
192                                    }
193                                }
194                            } else {
195                                Log::warn(format!("File is not an RGBA {file_name}"));
196                            }
197                        } else {
198                            Log::warn(format!("Screenshoot Error unable to read rgba file infos {}", &file_name));
199                        }
200                    } else {
201                        Log::warn(format!("Screenshoot Error unable to get texture ScreenshotTex {}", &file_name));
202                    }
203                } else {
204                    Log::err(format!("ScreenshotViewer : file {} is not valid", &file_name))
205                }
206                file_name_lock.clear();
207            }
208        }
209        Ui::hseparator();
210        if Ui::button("Open", None) {
211            if true {
212                let mut file_browser = FileBrowser::default();
213
214                if cfg!(target_os = "android")
215                    && let Some(img_dir) = get_external_path(&self.sk_info)
216                {
217                    file_browser.dir = img_dir;
218                }
219                if !file_browser.dir.exists() {
220                    file_browser.dir = current_dir().unwrap_or_default();
221                }
222                file_browser.caller = self.id.clone();
223                file_browser.window_pose = Ui::popup_pose(Vec3::ZERO);
224                self.close_file_browser();
225                SkInfo::send_event(&self.sk_info, StepperAction::add(self.id.clone() + BROWSER_SUFFIX, file_browser));
226            } else if !Platform::get_file_picker_visible() {
227                Platform::file_picker_sz(
228                    PickerMode::Open,
229                    move |ok, file_name| {
230                        let mut name = FILE_NAME.lock().unwrap();
231                        name.clear();
232                        if ok {
233                            Log::diag(format!("Open screenshot {file_name}"));
234                            name.push_str(file_name);
235                            Platform::file_picker_close();
236                        } else {
237                            // großen tricherie
238                            name.push_str("aaa.raw");
239                        }
240                    },
241                    &SCREENSHOT_FORMATS,
242                )
243            }
244        }
245        Ui::same_line();
246        if Ui::button("Take Screenshot", None) {
247            let mut camera_at = self.window_pose;
248            camera_at.orientation = Quat::look_dir(camera_at.get_forward() * -1.0);
249            let width_i = self.picture_size.x as i32;
250            let height_i = self.picture_size.y as i32;
251
252            Renderer::screenshot_capture(
253                token,
254                move |dots, width, height| {
255                    Log::info(format!("data length {} -> size {}/{}", dots.len(), width, height));
256                    let tex = Tex::find(CAPTURE_TEXTURE_ID).ok();
257                    match tex {
258                        Some(mut tex) => tex.set_colors32(width, height, dots),
259                        None => todo!(),
260                    };
261                },
262                camera_at,
263                width_i,
264                height_i,
265                Some(self.field_of_view),
266                Some(TexFormat::RGBA32),
267            );
268
269            self.screen = Sprite::from_tex(&self.tex, None, None).ok();
270        }
271        Ui::same_line();
272        Ui::push_enabled(self.screen.is_some(), None);
273        if Ui::button("Save", None) && !Platform::get_file_picker_visible() {
274            if cfg!(target_os = "android")
275                && let Some(img_dir) = get_external_path(&self.sk_info)
276                && let Err(err) = set_current_dir(&img_dir)
277            {
278                Log::err(format!("Unable to move current_dir to {img_dir:?} : {err:?}"))
279            }
280            if true {
281                let mut file_browser = FileBrowser::default();
282
283                if cfg!(target_os = "android")
284                    && let Some(img_dir) = get_external_path(&self.sk_info)
285                {
286                    file_browser.dir = img_dir;
287                }
288                if !file_browser.dir.exists() {
289                    file_browser.dir = current_dir().unwrap_or_default();
290                }
291                file_browser.picker_mode = PickerMode::Save;
292                file_browser.caller = self.id.clone();
293                file_browser.window_pose = Ui::popup_pose(Vec3::ZERO);
294                file_browser.file_name_to_save = "scr_.rgba".into();
295                file_browser.exts = vec![".rgba".into(), ".raw".into()];
296                self.close_file_browser();
297                SkInfo::send_event(&self.sk_info, StepperAction::add(self.id.clone() + BROWSER_SUFFIX, file_browser));
298            } else {
299                Platform::file_picker_sz(
300                    PickerMode::Save,
301                    move |ok, file_name| {
302                        if ok {
303                            save_screenshot(file_name);
304                        }
305                    },
306                    &SCREENSHOT_FORMATS,
307                )
308            }
309        }
310        Ui::pop_enabled();
311        Ui::window_end();
312    }
313
314    fn close_file_browser(&mut self) {
315        SkInfo::send_event(&self.sk_info, StepperAction::remove(self.id.clone() + BROWSER_SUFFIX));
316    }
317
318    fn close(&mut self, triggering: bool) -> bool {
319        if triggering {
320            self.close_file_browser();
321            self.shutdown_completed = true;
322        }
323        self.shutdown_completed
324    }
325}
326
327fn save_screenshot(file_name: &str) {
328    let mut name = file_name.to_string();
329    if !file_name.ends_with(".rgba") && !file_name.ends_with(".raw") {
330        name += ".raw";
331    }
332
333    if let Ok(tex) = Tex::find(CAPTURE_TEXTURE_ID) {
334        if let Some((width, height, size)) = tex.get_data_infos(0) {
335            Log::diag(format!("size is {}", size * 4));
336            let data = vec![0u8; size * 4];
337            let data_slice = data.as_slice();
338            if tex.get_color_data_u8(data_slice, 4, 0) {
339                match File::create(&name) {
340                    // Vive le format RGBA !!! https://github.com/bzotto/rgba_bitmap
341                    Ok(mut file) => {
342                        if let Err(err) = file.write_fmt(format_args!("RGBA")) {
343                            Log::warn(format!("Screenshoot Error when writing RGBA {} : {:?}", &name, err));
344                        }
345                        if let Err(err) = file.write(&width.to_be_bytes()[4..]) {
346                            Log::warn(format!("Screenshoot Error when writing width {} : {:?}", &name, err));
347                        }
348                        if let Err(err) = file.write(&height.to_be_bytes()[4..]) {
349                            Log::warn(format!("Screenshoot Error when writing height {} : {:?}", &name, err));
350                        }
351                        if let Err(err) = file.write_all(data_slice) {
352                            Log::warn(format!("Screenshoot Error when writing raw image {} : {:?}", &name, err));
353                        }
354                    }
355                    Err(err) => Log::warn(format!("Screenshoot Error when creating file {name} : {err:?}")),
356                }
357            } else {
358                Log::warn(format!("Screenshoot Error when getting texture data {file_name}"));
359            }
360        } else {
361            Log::warn(format!("Screenshoot Error unable to get texture infos {file_name}"));
362        }
363    } else {
364        Log::warn(format!("Screenshoot Error unable to get texture ScreenshotTex {file_name}"));
365    }
366}
stereokit_rust/tools/screenshot.rs

stereokit_rust/tools/
screenshot.rs
