Add command to screenshot the application (#2293)

Closes #2117
Closes #482

### What
Cmd-P for command palette, then select "Screenshot" command. Rerun will
re-style itself for Web for a frame, screenshot that, and then copy it
to the clipboard. I only implemented this for the native viewer.

You can set the resolution at startup with `--window-size`. All
screenshots are captured at 2x pixels-per-point, i.e. pretending that
you are on a high-dpi screen, wether you are or not.

You can also trigger this from the command line:

```
❯ cargo rerun ../fiat.rrd --screenshot-to fiat.png
    Finished dev [optimized + debuginfo] target(s) in 0.46s
     Running `target/debug/rerun ../fiat.rrd --screenshot-to fiat.png --window-size 1024x768`
[2023-05-31T16:11:21Z INFO  rerun::run] Loading "../fiat.rrd"…
[2023-05-31T16:11:22Z INFO  re_viewer::screenshotter] Screenshot saved to "fiat.png"
```

We can use this to generate screenshots for our examples.

### Result:


![fiat](https://github.com/rerun-io/rerun/assets/1148717/98cc125e-6cb5-4d84-81ff-062f54c7fb97)



### Checklist
* [x] I have read and agree to [Contributor
Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and
the [Code of
Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md)
* [x] I've included a screenshot or gif (if applicable)

<!-- This line will get updated when the PR build summary job finishes.
-->
PR Build Summary: https://build.rerun.io/pr/2293

<!-- pr-link-docs:start -->
Docs preview: https://rerun.io/preview/11b16c5/docs
<!-- pr-link-docs:end -->

crates/re_data_ui/src/image.rs

@@ -713,7 +713,6 @@ fn save_image(tensor: &re_components::Tensor, dynamic_image: &image::DynamicImag
                 .save_file()
             {
                 match dynamic_image.save(&path) {
-                    // TODO(emilk): show a popup instead of logging result
                     Ok(()) => {
                         re_log::info!("Image saved to {path:?}");
                     }

crates/re_ui/examples/re_ui_example.rs

@@ -238,7 +238,7 @@ impl ExampleApp {
         };
         let top_bar_style = self
             .re_ui
-            .top_bar_style(native_pixels_per_point, fullscreen);
+            .top_bar_style(native_pixels_per_point, fullscreen, false);
 
         egui::TopBottomPanel::top("top_bar")
             .frame(self.re_ui.top_panel_frame())

crates/re_ui/src/command.rs

@@ -47,6 +47,10 @@ pub enum Command {
     PlaybackStepBack,
     PlaybackStepForward,
     PlaybackRestart,
+
+    // Dev-tools:
+    #[cfg(not(target_arch = "wasm32"))]
+    ScreenshotWholeApp,
 }
 
 impl Command {
@@ -128,6 +132,12 @@ impl Command {
                 "Move the time marker to the next point in time with any data",
             ),
             Command::PlaybackRestart => ("Restart", "Restart from beginning of timeline"),
+
+            #[cfg(not(target_arch = "wasm32"))]
+            Command::ScreenshotWholeApp => (
+                "Screenshot",
+                "Copy screenshot of the whole app to clipboard",
+            ),
         }
     }
 
@@ -189,6 +199,9 @@ impl Command {
             Command::PlaybackStepBack => Some(key(Key::ArrowLeft)),
             Command::PlaybackStepForward => Some(key(Key::ArrowRight)),
             Command::PlaybackRestart => Some(cmd(Key::ArrowLeft)),
+
+            #[cfg(not(target_arch = "wasm32"))]
+            Command::ScreenshotWholeApp => None,
         }
     }
 

crates/re_ui/src/lib.rs

@@ -236,6 +236,7 @@ impl ReUi {
         &self,
         native_pixels_per_point: Option<f32>,
         fullscreen: bool,
+        style_like_web: bool,
     ) -> TopBarStyle {
         let gui_zoom = if let Some(native_pixels_per_point) = native_pixels_per_point {
             native_pixels_per_point / self.egui_ctx.pixels_per_point()
@@ -245,7 +246,7 @@ impl ReUi {
 
         // On Mac, we share the same space as the native red/yellow/green close/minimize/maximize buttons.
         // This means we need to make room for them.
-        let make_room_for_window_buttons = {
+        let make_room_for_window_buttons = !style_like_web && {
             #[cfg(target_os = "macos")]
             {
                 crate::FULLSIZE_CONTENT && !fullscreen

crates/re_viewer/Cargo.toml

@@ -70,6 +70,7 @@ ahash.workspace = true
 anyhow.workspace = true
 arrow2.workspace = true
 arrow2_convert.workspace = true
+bytemuck.workspace = true
 cfg-if.workspace = true
 eframe = { workspace = true, default-features = false, features = [
   "default_fonts",

crates/re_viewer/src/app.rs

@@ -41,10 +41,35 @@ enum TimeControlCommand {
 // ----------------------------------------------------------------------------
 
 /// Settings set once at startup (e.g. via command-line options) and not serialized.
-#[derive(Clone, Copy, Default)]
+#[derive(Clone)]
 pub struct StartupOptions {
     pub memory_limit: re_memory::MemoryLimit,
+
     pub persist_state: bool,
+
+    /// Take a screenshot of the app and quit.
+    /// We use this to generate screenshots of our exmples.
+    #[cfg(not(target_arch = "wasm32"))]
+    pub screenshot_to_path_then_quit: Option<std::path::PathBuf>,
+
+    /// Set the screen resolution in logical points.
+    #[cfg(not(target_arch = "wasm32"))]
+    pub resolution_in_points: Option<[f32; 2]>,
+}
+
+impl Default for StartupOptions {
+    fn default() -> Self {
+        Self {
+            memory_limit: re_memory::MemoryLimit::default(),
+            persist_state: true,
+
+            #[cfg(not(target_arch = "wasm32"))]
+            screenshot_to_path_then_quit: None,
+
+            #[cfg(not(target_arch = "wasm32"))]
+            resolution_in_points: None,
+        }
+    }
 }
 
 // ----------------------------------------------------------------------------
@@ -60,6 +85,7 @@ pub struct App {
     startup_options: StartupOptions,
     ram_limit_warner: re_memory::RamLimitWarner,
     re_ui: re_ui::ReUi,
+    screenshotter: crate::screenshotter::Screenshotter,
 
     /// Listens to the local text log stream
     text_log_rx: std::sync::mpsc::Receiver<re_log::LogMsg>,
@@ -146,11 +172,21 @@ impl App {
             );
         }
 
+        #[allow(unused_mut, clippy::needless_update)] // false positive on web
+        let mut screenshotter = crate::screenshotter::Screenshotter::default();
+
+        #[cfg(not(target_arch = "wasm32"))]
+        if let Some(screenshot_path) = startup_options.screenshot_to_path_then_quit.clone() {
+            screenshotter.screenshot_to_path_then_quit(screenshot_path);
+        }
+
         Self {
             build_info,
             startup_options,
             ram_limit_warner: re_memory::RamLimitWarner::warn_at_fraction_of_max(0.75),
             re_ui,
+            screenshotter,
+
             text_log_rx,
             component_ui_registry: re_data_ui::create_component_ui_registry(),
             rx,
@@ -389,6 +425,11 @@ impl App {
             Command::PlaybackRestart => {
                 self.run_time_control_command(TimeControlCommand::Restart);
             }
+
+            #[cfg(not(target_arch = "wasm32"))]
+            Command::ScreenshotWholeApp => {
+                self.screenshotter.request_screenshot();
+            }
         }
     }
 
@@ -507,13 +548,27 @@ impl eframe::App for App {
     fn update(&mut self, egui_ctx: &egui::Context, frame: &mut eframe::Frame) {
         let frame_start = Instant::now();
 
+        #[cfg(not(target_arch = "wasm32"))]
+        if let Some(resolution_in_points) = self.startup_options.resolution_in_points.take() {
+            frame.set_window_size(resolution_in_points.into());
+        }
+
+        #[cfg(not(target_arch = "wasm32"))]
+        if self.screenshotter.update(egui_ctx, frame).quit {
+            frame.close();
+            return;
+        }
+
         if self.startup_options.memory_limit.limit.is_none() {
             // we only warn about high memory usage if the user hasn't specified a limit
             self.ram_limit_warner.update();
         }
 
         #[cfg(not(target_arch = "wasm32"))]
-        {
+        if self.screenshotter.is_screenshotting() {
+            // Make screenshots high-quality by pretending we have a high-dpi display, whether we do or not:
+            egui_ctx.set_pixels_per_point(2.0);
+        } else {
             // Ensure zoom factor is sane and in 10% steps at all times before applying it.
             {
                 let mut zoom_factor = self.state.app_options.zoom_factor;
@@ -679,7 +734,10 @@ impl eframe::App for App {
         }
 
         self.handle_dropping_files(egui_ctx);
-        self.toasts.show(egui_ctx);
+
+        if !self.screenshotter.is_screenshotting() {
+            self.toasts.show(egui_ctx);
+        }
 
         if let Some(cmd) = self.cmd_palette.show(egui_ctx) {
             self.pending_commands.push(cmd);
@@ -702,6 +760,13 @@ impl eframe::App for App {
             re_log::warn_once!("Blueprint unexpectedly missing from store.");
         }
     }
+
+    #[cfg(not(target_arch = "wasm32"))]
+    fn post_rendering(&mut self, _window_size: [u32; 2], frame: &eframe::Frame) {
+        if let Some(screenshot) = frame.screenshot() {
+            self.screenshotter.save(&screenshot);
+        }
+    }
 }
 
 fn paint_background_fill(ui: &mut egui::Ui) {
@@ -1237,7 +1302,10 @@ fn top_panel(
             frame.info().window_info.fullscreen
         }
     };
-    let top_bar_style = app.re_ui.top_bar_style(native_pixels_per_point, fullscreen);
+    let style_like_web = app.screenshotter.is_screenshotting();
+    let top_bar_style =
+        app.re_ui
+            .top_bar_style(native_pixels_per_point, fullscreen, style_like_web);
 
     egui::TopBottomPanel::top("top_bar")
         .frame(app.re_ui.top_panel_frame())

crates/re_viewer/src/lib.rs

@@ -9,6 +9,7 @@ pub mod env_vars;
 #[cfg(not(target_arch = "wasm32"))]
 mod profiler;
 mod remote_viewer_app;
+mod screenshotter;
 mod ui;
 mod viewer_analytics;
 

crates/re_viewer/src/remote_viewer_app.rs

@@ -75,7 +75,7 @@ impl RemoteViewerApp {
                 let app = crate::App::from_receiver(
                     self.build_info,
                     &self.app_env,
-                    self.startup_options,
+                    self.startup_options.clone(),
                     self.re_ui.clone(),
                     storage,
                     rx,

crates/re_viewer/src/screenshotter.rs

@@ -0,0 +1,100 @@
+//! Screenshotting not implemented on web yet because we
+//! haven't implemented "copy image to clipboard" there.
+
+/// Helper for screenshotting the entire app
+#[cfg(not(target_arch = "wasm32"))]
+#[derive(Default)]
+pub struct Screenshotter {
+    countdown: Option<usize>,
+    target_path: Option<std::path::PathBuf>,
+    quit: bool,
+}
+
+#[cfg(not(target_arch = "wasm32"))]
+#[must_use]
+pub struct ScreenshotterOutput {
+    /// If true, the screenshotter was told at startup to quit after its donw.
+    pub quit: bool,
+}
+
+#[cfg(not(target_arch = "wasm32"))]
+impl Screenshotter {
+    /// Used for generating screenshots in dev builds.
+    ///
+    /// Should only be called at startup.
+    pub fn screenshot_to_path_then_quit(&mut self, path: std::path::PathBuf) {
+        assert!(self.countdown.is_none(), "screenshotter misused");
+        self.request_screenshot();
+        self.target_path = Some(path);
+    }
+
+    pub fn request_screenshot(&mut self) {
+        // Give app time to change the style, and then wait for animations to finish:
+        self.countdown = Some(10);
+    }
+
+    /// Call once per frame
+    pub fn update(
+        &mut self,
+        egui_ctx: &egui::Context,
+        frame: &mut eframe::Frame,
+    ) -> ScreenshotterOutput {
+        if let Some(countdown) = &mut self.countdown {
+            if *countdown == 0 {
+                frame.request_screenshot();
+            } else {
+                *countdown -= 1;
+            }
+
+            egui_ctx.request_repaint(); // Make sure we keep counting down
+        }
+
+        ScreenshotterOutput { quit: self.quit }
+    }
+
+    /// If true, temporarily re-style the UI to make it suitable for capture!
+    ///
+    /// We do the re-styling to create consistent screenshots across platforms.
+    /// In particular, we style the UI to look like the web viewer.
+    pub fn is_screenshotting(&self) -> bool {
+        self.countdown.is_some()
+    }
+
+    pub fn save(&mut self, image: &egui::ColorImage) {
+        self.countdown = None;
+        if let Some(path) = self.target_path.take() {
+            let w = image.width() as _;
+            let h = image.height() as _;
+            let image =
+                image::RgbaImage::from_raw(w, h, bytemuck::pod_collect_to_vec(&image.pixels))
+                    .expect("Failed to create image");
+            match image.save(&path) {
+                Ok(()) => {
+                    re_log::info!("Screenshot saved to {path:?}");
+                    self.quit = true;
+                }
+                Err(err) => {
+                    panic!("Failed saving screenshot to {path:?}: {err}");
+                }
+            }
+        } else {
+            re_viewer_context::Clipboard::with(|cb| {
+                cb.set_image(image.size, bytemuck::cast_slice(&image.pixels));
+            });
+        }
+    }
+}
+
+// ----------------------------------------------------------------------------
+
+#[cfg(target_arch = "wasm32")]
+#[derive(Default)]
+pub struct Screenshotter {}
+
+#[cfg(target_arch = "wasm32")]
+impl Screenshotter {
+    #[allow(clippy::unused_self)]
+    pub fn is_screenshotting(&self) -> bool {
+        false
+    }
+}