Central GpuReadback handling for re_viewer, experimental space view s…

…creenshots (#1717)

EXPERIMENTAL: screenshot in context menu on spatial views

.vscode/settings.json

@@ -35,6 +35,7 @@
         "nyud",
         "objectron",
         "Readback",
+        "readbacks",
         "Skybox",
         "smallvec",
         "swapchain",

crates/re_renderer/src/view_builder.rs

@@ -180,6 +180,7 @@ impl Default for TargetConfiguration {
     }
 }
 
+#[derive(Clone)]
 pub struct ScheduledScreenshot {
     pub identifier: GpuReadbackBufferIdentifier,
     pub width: u32,

crates/re_renderer/src/wgpu_resources/mod.rs

@@ -115,6 +115,7 @@ impl WgpuResourcePools {
     }
 }
 
+#[derive(Clone, Copy)]
 pub struct TextureRowDataInfo {
     /// How many bytes per row contain actual data.
     pub bytes_per_row_unpadded: u32,

crates/re_viewer/src/app.rs

@@ -11,13 +11,13 @@ use re_arrow_store::DataStoreStats;
 use re_data_store::log_db::LogDb;
 use re_format::format_number;
 use re_log_types::{ApplicationId, LogMsg, RecordingId};
-use re_renderer::WgpuResourcePoolStatistics;
+use re_renderer::{GpuReadbackBufferIdentifier, WgpuResourcePoolStatistics};
 use re_smart_channel::Receiver;
 use re_ui::{toasts, Command};
 
 use crate::{
     app_icon::setup_app_icon,
-    misc::{AppOptions, Caches, RecordingConfig, ViewerContext},
+    misc::{AppOptions, Caches, RecordingConfig, ScheduledGpuReadback, ViewerContext},
     ui::{data_ui::ComponentUiRegistry, Blueprint},
     viewer_analytics::ViewerAnalytics,
 };
@@ -87,6 +87,10 @@ pub struct App {
 
     latest_queue_interest: instant::Instant,
 
+    /// List of all data we're currently waiting for from the GPU for each application id
+    scheduled_gpu_readbacks_per_application:
+        HashMap<ApplicationId, HashMap<GpuReadbackBufferIdentifier, ScheduledGpuReadback>>,
+
     /// Measures how long a frame takes to paint
     frame_time_history: egui::util::History<f32>,
 
@@ -143,6 +147,8 @@ impl App {
 
             latest_queue_interest: instant::Instant::now(), // TODO(emilk): `Instant::MIN` when we have our own `Instant` that supports it.
 
+            scheduled_gpu_readbacks_per_application: HashMap::default(),
+
             frame_time_history: egui::util::History::new(1..100, 0.5),
 
             pending_commands: Default::default(),
@@ -404,6 +410,37 @@ impl App {
                 );
             });
     }
+
+    fn process_gpu_readback_data(&mut self, data: &[u8], identifier: GpuReadbackBufferIdentifier) {
+        for (application_id, scheduled_gpu_readbacks) in
+            &mut self.scheduled_gpu_readbacks_per_application
+        {
+            if let Some((_, scheduled_readback)) = scheduled_gpu_readbacks.remove_entry(&identifier)
+            {
+                match scheduled_readback {
+                    ScheduledGpuReadback::SpaceViewScreenshot {
+                        screenshot,
+                        space_view_id,
+                        mode,
+                    } => {
+                        if let Some(blueprint) = self.state.blueprints.get_mut(application_id) {
+                            blueprint.viewport.save_spaceview_screenshot(
+                                &screenshot,
+                                data,
+                                space_view_id,
+                                mode,
+                            );
+                        }
+                    }
+                }
+                return;
+            }
+        }
+        re_log::warn_once!(
+            "Received unexpected GPU readback. (Size: {}, identifier {identifier}).",
+            data.len()
+        );
+    }
 }
 
 impl eframe::App for App {
@@ -449,8 +486,8 @@ impl eframe::App for App {
             egui_ctx.set_pixels_per_point(pixels_per_point);
         }
 
+        // TODO(andreas): store the re_renderer somewhere else.
         let gpu_resource_stats = {
-            // TODO(andreas): store the re_renderer somewhere else.
             let egui_renderer = {
                 let render_state = frame.wgpu_render_state().unwrap();
                 &mut render_state.renderer.read()
@@ -459,6 +496,13 @@ impl eframe::App for App {
                 .paint_callback_resources
                 .get::<re_renderer::RenderContext>()
                 .unwrap();
+
+            // Handle GPU readback data.
+            render_ctx
+                .gpu_readback_belt
+                .lock()
+                .receive_data(|data, identifier| self.process_gpu_readback_data(data, identifier));
+
             // Query statistics before begin_frame as this might be more accurate if there's resources that we recreate every frame.
             render_ctx.gpu_resources.statistics()
         };
@@ -505,7 +549,7 @@ impl eframe::App for App {
                 let blueprint = self
                     .state
                     .blueprints
-                    .entry(selected_app_id)
+                    .entry(selected_app_id.clone())
                     .or_insert_with(|| Blueprint::new(egui_ctx));
 
                 recording_config_entry(
@@ -535,6 +579,9 @@ impl eframe::App for App {
                         self.state.show(
                             ui,
                             render_ctx,
+                            self.scheduled_gpu_readbacks_per_application
+                                .entry(selected_app_id)
+                                .or_default(),
                             log_db,
                             &self.re_ui,
                             &self.component_ui_registry,
@@ -936,10 +983,12 @@ struct AppState {
 }
 
 impl AppState {
+    #[allow(clippy::too_many_arguments)]
     fn show(
         &mut self,
         ui: &mut egui::Ui,
         render_ctx: &mut re_renderer::RenderContext,
+        scheduled_gpu_readbacks: &mut HashMap<GpuReadbackBufferIdentifier, ScheduledGpuReadback>,
         log_db: &LogDb,
         re_ui: &re_ui::ReUi,
         component_ui_registry: &ComponentUiRegistry,
@@ -977,6 +1026,7 @@ impl AppState {
             rec_cfg,
             re_ui,
             render_ctx,
+            scheduled_gpu_readbacks,
         };
 
         let blueprint = blueprints
@@ -1587,6 +1637,17 @@ fn options_menu_ui(ui: &mut egui::Ui, _frame: &mut eframe::Frame, options: &mut
         ui.close_menu();
     }
 
+    #[cfg(not(target_arch = "wasm32"))]
+    {
+        if ui
+            .checkbox(&mut options.experimental_space_view_screenshots, "(experimental) Space View screenshots")
+            .on_hover_text("Allow taking screenshots of 2D & 3D space views via their context menu. Does not contain labels.")
+            .clicked()
+        {
+            ui.close_menu();
+        }
+    }
+
     #[cfg(debug_assertions)]
     {
         ui.separator();

crates/re_viewer/src/misc/app_options.rs

@@ -18,6 +18,10 @@ pub struct AppOptions {
     /// (Since this is serialized, even between sessions!)
     #[cfg(not(target_arch = "wasm32"))]
     pub zoom_factor: f32,
+
+    /// Enable the experimental feature for space view screenshots.
+    #[cfg(not(target_arch = "wasm32"))]
+    pub experimental_space_view_screenshots: bool,
 }
 
 impl Default for AppOptions {
@@ -32,6 +36,9 @@ impl Default for AppOptions {
 
             #[cfg(not(target_arch = "wasm32"))]
             zoom_factor: 1.0,
+
+            #[cfg(not(target_arch = "wasm32"))]
+            experimental_space_view_screenshots: false,
         }
     }
 }

crates/re_viewer/src/misc/gpu_readbacks.rs

@@ -0,0 +1,26 @@
+use re_renderer::ScheduledScreenshot;
+
+use crate::ui::SpaceViewId;
+
+#[derive(PartialEq, Eq, Clone, Copy)]
+#[allow(dead_code)] // Not used on the web.
+pub enum ScreenshotMode {
+    /// The screenshot will be saved to disc and copied to the clipboard.
+    SaveAndCopyToClipboard,
+
+    /// The screenshot will be copied to the clipboard.
+    CopyToClipboard,
+}
+
+/// A previously scheduled GPU readback, waiting for getting the result.
+pub enum ScheduledGpuReadback {
+    SpaceViewScreenshot {
+        screenshot: ScheduledScreenshot,
+        space_view_id: SpaceViewId,
+        mode: ScreenshotMode,
+    },
+    // Picking {
+    //     picking: ScheduledPicking,
+    //     space_view_id: SpaceViewId,
+    // },
+}

crates/re_viewer/src/misc/mod.rs

@@ -35,6 +35,9 @@ pub use {
     },
 };
 
+mod gpu_readbacks;
+pub use gpu_readbacks::{ScheduledGpuReadback, ScreenshotMode};
+
 // ----------------------------------------------------------------------------
 
 pub fn help_hover_button(ui: &mut egui::Ui) -> egui::Response {

crates/re_viewer/src/misc/viewer_context.rs

@@ -1,5 +1,7 @@
+use ahash::HashMap;
 use re_data_store::{log_db::LogDb, InstancePath};
 use re_log_types::{ComponentPath, EntityPath, MsgId, TimeInt, Timeline};
+use re_renderer::GpuReadbackBufferIdentifier;
 
 use crate::ui::{
     data_ui::{ComponentUiRegistry, DataUi},
@@ -8,7 +10,7 @@ use crate::ui::{
 
 use super::{
     item::{Item, ItemCollection},
-    HoverHighlight,
+    HoverHighlight, ScheduledGpuReadback,
 };
 
 /// Common things needed by many parts of the viewer.
@@ -32,6 +34,9 @@ pub struct ViewerContext<'a> {
     pub re_ui: &'a re_ui::ReUi,
 
     pub render_ctx: &'a mut re_renderer::RenderContext,
+
+    /// List of all data we're currently waiting for from the GPU.
+    pub scheduled_gpu_readbacks: &'a mut HashMap<GpuReadbackBufferIdentifier, ScheduledGpuReadback>,
 }
 
 impl<'a> ViewerContext<'a> {

crates/re_viewer/src/ui/view_spatial/ui.rs

@@ -9,7 +9,8 @@ use re_renderer::renderer::OutlineConfig;
 
 use crate::{
     misc::{
-        space_info::query_view_coordinates, SelectionHighlight, SpaceViewHighlights, ViewerContext,
+        space_info::query_view_coordinates, ScreenshotMode, SelectionHighlight,
+        SpaceViewHighlights, ViewerContext,
     },
     ui::{data_blueprint::DataBlueprintTree, view_spatial::UiLabelTarget, SpaceViewId},
 };
@@ -620,3 +621,31 @@ pub fn outline_config(gui_ctx: &egui::Context) -> OutlineConfig {
         color_layer_b: selection_outline_color,
     }
 }
+
+pub fn screenshot_context_menu(
+    _ctx: &ViewerContext<'_>,
+    response: egui::Response,
+) -> (egui::Response, Option<ScreenshotMode>) {
+    #[cfg(not(target_arch = "wasm32"))]
+    {
+        if _ctx.app_options.experimental_space_view_screenshots {
+            let mut take_screenshot = None;
+            let response = response.context_menu(|ui| {
+                if ui.button("Screenshot (save to disk)").clicked() {
+                    take_screenshot = Some(ScreenshotMode::SaveAndCopyToClipboard);
+                    ui.close_menu();
+                } else if ui.button("Screenshot (clipboard only)").clicked() {
+                    take_screenshot = Some(ScreenshotMode::CopyToClipboard);
+                    ui.close_menu();
+                }
+            });
+            (response, take_screenshot)
+        } else {
+            (response, None)
+        }
+    }
+    #[cfg(target_arch = "wasm32")]
+    {
+        (response, None)
+    }
+}

crates/re_viewer/src/ui/view_spatial/ui_2d.rs

@@ -8,11 +8,13 @@ use re_log_types::component_types::TensorTrait;
 use re_renderer::view_builder::TargetConfiguration;
 
 use super::{
-    eye::Eye, scene::AdditionalPickingInfo, ui::create_labels, SpatialNavigationMode,
-    ViewSpatialState,
+    eye::Eye,
+    scene::AdditionalPickingInfo,
+    ui::{create_labels, screenshot_context_menu},
+    SpatialNavigationMode, ViewSpatialState,
 };
 use crate::{
-    misc::{HoveredSpace, Item, SpaceViewHighlights},
+    misc::{HoveredSpace, Item, ScheduledGpuReadback, SpaceViewHighlights},
     ui::{
         data_ui::{self, DataUi},
         view_spatial::{
@@ -427,6 +429,8 @@ fn view_2d_scrollable(
 
     // ------------------------------------------------------------------------
 
+    let (response, screenshot_action) = screenshot_context_menu(ctx, response);
+
     // Draw a re_renderer driven view.
     // Camera & projection are configured to ingest space coordinates directly.
     {
@@ -445,14 +449,25 @@ fn view_2d_scrollable(
             return response;
         };
 
-        let Ok(callback) = create_scene_paint_callback(
+        let Ok((callback, screenshot)) = create_scene_paint_callback(
             ctx.render_ctx,
             target_config, painter.clip_rect(),
             scene.primitives,
             &ScreenBackground::ClearColor(parent_ui.visuals().extreme_bg_color.into()),
+            screenshot_action.is_some(),
         ) else {
             return response;
         };
+        if let (Some(screenshot), Some(screenshot_action)) = (screenshot, screenshot_action) {
+            ctx.scheduled_gpu_readbacks.insert(
+                screenshot.identifier,
+                ScheduledGpuReadback::SpaceViewScreenshot {
+                    space_view_id,
+                    screenshot,
+                    mode: screenshot_action,
+                },
+            );
+        }
 
         painter.add(callback);
     }