From f0c2fdea77633e92b891ed6db6bff28d4661e973 Mon Sep 17 00:00:00 2001 From: Andreas Reich Date: Wed, 29 Mar 2023 11:28:01 +0200 Subject: [PATCH] Central GpuReadback handling for re_viewer, experimental space view screenshots (#1717) EXPERIMENTAL: screenshot in context menu on spatial views --- .vscode/settings.json | 1 + crates/re_renderer/src/view_builder.rs | 1 + crates/re_renderer/src/wgpu_resources/mod.rs | 1 + crates/re_viewer/src/app.rs | 69 ++++++++++++++++- crates/re_viewer/src/misc/app_options.rs | 7 ++ crates/re_viewer/src/misc/gpu_readbacks.rs | 26 +++++++ crates/re_viewer/src/misc/mod.rs | 3 + crates/re_viewer/src/misc/viewer_context.rs | 7 +- crates/re_viewer/src/ui/view_spatial/ui.rs | 31 +++++++- crates/re_viewer/src/ui/view_spatial/ui_2d.rs | 23 +++++- crates/re_viewer/src/ui/view_spatial/ui_3d.rs | 36 ++++++--- .../src/ui/view_spatial/ui_renderer_bridge.rs | 40 +++++++--- crates/re_viewer/src/ui/viewport.rs | 76 ++++++++++++++++++- 13 files changed, 290 insertions(+), 31 deletions(-) create mode 100644 crates/re_viewer/src/misc/gpu_readbacks.rs diff --git a/.vscode/settings.json b/.vscode/settings.json index 416a398ed8f0..ed03ba4db684 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -35,6 +35,7 @@ "nyud", "objectron", "Readback", + "readbacks", "Skybox", "smallvec", "swapchain", diff --git a/crates/re_renderer/src/view_builder.rs b/crates/re_renderer/src/view_builder.rs index 36dc0f036dcd..643f4fadb578 100644 --- a/crates/re_renderer/src/view_builder.rs +++ b/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, diff --git a/crates/re_renderer/src/wgpu_resources/mod.rs b/crates/re_renderer/src/wgpu_resources/mod.rs index 87dc4520550b..b008dfaff4c8 100644 --- a/crates/re_renderer/src/wgpu_resources/mod.rs +++ b/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, diff --git a/crates/re_viewer/src/app.rs b/crates/re_viewer/src/app.rs index 4a33c4a073da..7ac13c7552ba 100644 --- a/crates/re_viewer/src/app.rs +++ b/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>, + /// Measures how long a frame takes to paint frame_time_history: egui::util::History, @@ -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::() .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, 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(); diff --git a/crates/re_viewer/src/misc/app_options.rs b/crates/re_viewer/src/misc/app_options.rs index 05ea62536cf7..e71d784f2696 100644 --- a/crates/re_viewer/src/misc/app_options.rs +++ b/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, } } } diff --git a/crates/re_viewer/src/misc/gpu_readbacks.rs b/crates/re_viewer/src/misc/gpu_readbacks.rs new file mode 100644 index 000000000000..260e72e89e3a --- /dev/null +++ b/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, + // }, +} diff --git a/crates/re_viewer/src/misc/mod.rs b/crates/re_viewer/src/misc/mod.rs index 83abb356d071..ddd25425642c 100644 --- a/crates/re_viewer/src/misc/mod.rs +++ b/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 { diff --git a/crates/re_viewer/src/misc/viewer_context.rs b/crates/re_viewer/src/misc/viewer_context.rs index 631f17231112..9ac87cc0f2f7 100644 --- a/crates/re_viewer/src/misc/viewer_context.rs +++ b/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, } impl<'a> ViewerContext<'a> { diff --git a/crates/re_viewer/src/ui/view_spatial/ui.rs b/crates/re_viewer/src/ui/view_spatial/ui.rs index 24ab3178167a..75ae6eeea592 100644 --- a/crates/re_viewer/src/ui/view_spatial/ui.rs +++ b/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) { + #[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) + } +} diff --git a/crates/re_viewer/src/ui/view_spatial/ui_2d.rs b/crates/re_viewer/src/ui/view_spatial/ui_2d.rs index fdc9ee30ae31..20629958c0d5 100644 --- a/crates/re_viewer/src/ui/view_spatial/ui_2d.rs +++ b/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); } diff --git a/crates/re_viewer/src/ui/view_spatial/ui_3d.rs b/crates/re_viewer/src/ui/view_spatial/ui_3d.rs index ca3fdeb68a74..29eac391df25 100644 --- a/crates/re_viewer/src/ui/view_spatial/ui_3d.rs +++ b/crates/re_viewer/src/ui/view_spatial/ui_3d.rs @@ -6,17 +6,17 @@ use macaw::{vec3, BoundingBox, Quat, Vec3}; use re_data_store::{InstancePath, InstancePathHash}; use re_log_types::{EntityPath, ViewCoordinates}; use re_renderer::{ - view_builder::{Projection, TargetConfiguration}, + view_builder::{Projection, ScheduledScreenshot, TargetConfiguration}, RenderContext, Size, }; use crate::{ - misc::{HoveredSpace, Item, SpaceViewHighlights}, + misc::{HoveredSpace, Item, ScheduledGpuReadback, SpaceViewHighlights}, ui::{ data_ui::{self, DataUi}, view_spatial::{ scene::AdditionalPickingInfo, - ui::{create_labels, outline_config}, + ui::{create_labels, outline_config, screenshot_context_menu}, ui_renderer_bridge::{create_scene_paint_callback, get_viewport, ScreenBackground}, SceneSpatial, SpaceCamera3D, SpatialNavigationMode, }, @@ -496,7 +496,9 @@ pub fn view_3d( } } - paint_view( + let (_, screenshot_action) = screenshot_context_menu(ctx, response); + + let screenshot = paint_view( ui, eye, rect, @@ -504,13 +506,25 @@ pub fn view_3d( ctx.render_ctx, &space.to_string(), state.auto_size_config(), + screenshot_action.is_some(), ); + 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, + }, + ); + } // Add egui driven labels on top of re_renderer content. let painter = ui.painter().with_clip_rect(ui.max_rect()); painter.extend(label_shapes); } +#[allow(clippy::too_many_arguments)] fn paint_view( ui: &mut egui::Ui, eye: Eye, @@ -519,14 +533,15 @@ fn paint_view( render_ctx: &mut RenderContext, name: &str, auto_size_config: re_renderer::AutoSizeConfig, -) { + take_screenshot: bool, +) -> Option { crate::profile_function!(); // Determine view port resolution and position. let pixels_from_point = ui.ctx().pixels_per_point(); let resolution_in_pixel = get_viewport(rect, pixels_from_point); if resolution_in_pixel[0] == 0 || resolution_in_pixel[1] == 0 { - return; + return None; } let target_config = TargetConfiguration { @@ -549,15 +564,18 @@ fn paint_view( .then(|| outline_config(ui.ctx())), }; - let Ok(callback) = create_scene_paint_callback( + let Ok((callback, scheduled_screenshot)) = create_scene_paint_callback( render_ctx, target_config, rect, - scene.primitives, &ScreenBackground::GenericSkybox) + scene.primitives, &ScreenBackground::GenericSkybox, + take_screenshot) else { - return; + return None; }; ui.painter().add(callback); + + scheduled_screenshot } fn show_projections_from_2d_space( diff --git a/crates/re_viewer/src/ui/view_spatial/ui_renderer_bridge.rs b/crates/re_viewer/src/ui/view_spatial/ui_renderer_bridge.rs index 73d0e2172569..7d7dabad8f7a 100644 --- a/crates/re_viewer/src/ui/view_spatial/ui_renderer_bridge.rs +++ b/crates/re_viewer/src/ui/view_spatial/ui_renderer_bridge.rs @@ -1,7 +1,7 @@ use egui::mutex::Mutex; use re_renderer::{ renderer::{DepthCloudDrawData, GenericSkyboxDrawData, MeshDrawData, RectangleDrawData}, - view_builder::{TargetConfiguration, ViewBuilder}, + view_builder::{ScheduledScreenshot, TargetConfiguration, ViewBuilder}, RenderContext, }; @@ -20,16 +20,25 @@ pub fn create_scene_paint_callback( clip_rect: egui::Rect, primitives: SceneSpatialPrimitives, background: &ScreenBackground, -) -> anyhow::Result { + take_screenshot: bool, +) -> anyhow::Result<(egui::PaintCallback, Option)> { let pixels_from_point = target_config.pixels_from_point; - let (command_buffer, view_builder) = - create_and_fill_view_builder(render_ctx, target_config, primitives, background)?; - Ok(renderer_paint_callback( + let (command_buffer, view_builder, scheduled_screenshot) = create_and_fill_view_builder( render_ctx, - command_buffer, - view_builder, - clip_rect, - pixels_from_point, + target_config, + primitives, + background, + take_screenshot, + )?; + Ok(( + renderer_paint_callback( + render_ctx, + command_buffer, + view_builder, + clip_rect, + pixels_from_point, + ), + scheduled_screenshot, )) } @@ -43,7 +52,12 @@ fn create_and_fill_view_builder( target_config: TargetConfiguration, primitives: SceneSpatialPrimitives, background: &ScreenBackground, -) -> anyhow::Result<(wgpu::CommandBuffer, ViewBuilder)> { + take_screenshot: bool, +) -> anyhow::Result<( + wgpu::CommandBuffer, + ViewBuilder, + Option, +)> { let mut view_builder = ViewBuilder::default(); view_builder.setup_view(render_ctx, target_config)?; @@ -61,6 +75,10 @@ fn create_and_fill_view_builder( view_builder.queue_draw(&GenericSkyboxDrawData::new(render_ctx)); } + let scheduled_screenshot = take_screenshot + .then(|| view_builder.schedule_screenshot(render_ctx).ok()) + .flatten(); + let command_buffer = view_builder.draw( render_ctx, match background { @@ -69,7 +87,7 @@ fn create_and_fill_view_builder( }, )?; - Ok((command_buffer, view_builder)) + Ok((command_buffer, view_builder, scheduled_screenshot)) } slotmap::new_key_type! { pub struct ViewBuilderHandle; } diff --git a/crates/re_viewer/src/ui/viewport.rs b/crates/re_viewer/src/ui/viewport.rs index 37bf1a84124a..9bcc1861e5f3 100644 --- a/crates/re_viewer/src/ui/viewport.rs +++ b/crates/re_viewer/src/ui/viewport.rs @@ -6,9 +6,12 @@ use ahash::HashMap; use itertools::Itertools as _; use re_data_store::EntityPath; +use re_renderer::ScheduledScreenshot; use crate::{ - misc::{space_info::SpaceInfoCollection, Item, SpaceViewHighlights, ViewerContext}, + misc::{ + space_info::SpaceInfoCollection, Item, ScreenshotMode, SpaceViewHighlights, ViewerContext, + }, ui::space_view_heuristics::default_created_space_views, }; @@ -545,6 +548,77 @@ impl Viewport { }) .collect() } + + pub fn save_spaceview_screenshot( + &self, + screenshot: &ScheduledScreenshot, + data: &[u8], + space_view_id: SpaceViewId, + mode: ScreenshotMode, + ) { + let Some(space_view_display_name) = self.space_views + .get(&space_view_id) + .map(|space_view| &space_view.display_name) else { + re_log::warn!("Failed to take screenshot of space view, since it was closed before the operation was finished."); + return; + }; + + // Need to do a memcpy to remove the padding. + // TODO(andreas): This should be a utility in the render crate. + let row_info = screenshot.row_info; + let mut buffer = + Vec::with_capacity((row_info.bytes_per_row_unpadded * screenshot.height) as usize); + for row in 0..screenshot.height { + let offset = (row_info.bytes_per_row_padded * row) as usize; + buffer.extend_from_slice( + &data[offset..(offset + row_info.bytes_per_row_unpadded as usize)], + ); + } + + // Set to clipboard. + #[cfg(not(target_arch = "wasm32"))] + crate::misc::Clipboard::with(|clipboard| { + clipboard.set_image([screenshot.width as _, screenshot.height as _], &buffer); + }); + if mode == ScreenshotMode::CopyToClipboard { + return; + } + + // Get next available file name. + let safe_space_view_name = + space_view_display_name.replace(|c: char| !c.is_alphanumeric() && c != ' ', ""); + let mut i = 1; + let filename = loop { + let filename = format!("Screenshot {safe_space_view_name} - {i}.png"); + if !std::path::Path::new(&filename).exists() { + break filename; + } + i += 1; + }; + let filename = std::path::Path::new(&filename); + + match image::save_buffer( + filename, + &buffer, + screenshot.width, + screenshot.height, + image::ColorType::Rgba8, + ) { + Ok(_) => { + re_log::info!( + "Saved screenshot to {:?}.", + filename.canonicalize().unwrap_or(filename.to_path_buf()) + ); + } + Err(err) => { + re_log::error!( + "Failed to safe screenshot to {:?}: {}", + filename.canonicalize().unwrap_or(filename.to_path_buf()), + err + ); + } + } + } } /// Show a single button (`add_content`), justified,