From 43aa22c426ba845ed9384db7b07d5ca0de85581f Mon Sep 17 00:00:00 2001 From: Andreas Reich Date: Sun, 16 Apr 2023 15:41:24 +0200 Subject: [PATCH] Gpu picking for depth clouds (#1849) * wip * allow for hovering depth clouds via gpu picking * Use `[x, y]: [u32; 2]` as argument --------- Co-authored-by: Emil Ernerfeldt --- .vscode/settings.json | 1 + .../src/component_types/instance_key.rs | 10 +++ crates/re_renderer/examples/depth_cloud.rs | 1 + crates/re_renderer/shader/depth_cloud.wgsl | 31 +++++++--- .../re_renderer/src/renderer/depth_cloud.rs | 20 +++--- .../src/resource_managers/texture_manager.rs | 10 +++ .../src/ui/view_spatial/scene/mod.rs | 3 +- .../src/ui/view_spatial/scene/picking.rs | 22 +++++-- .../src/ui/view_spatial/scene/primitives.rs | 3 +- .../view_spatial/scene/scene_part/images.rs | 61 ++++++++----------- crates/re_viewer/src/ui/view_spatial/ui.rs | 45 +++++++++----- 11 files changed, 132 insertions(+), 75 deletions(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index ed03ba4db684..c52ea2c7178d 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -25,6 +25,7 @@ "andreas", "bbox", "bindgroup", + "colormap", "emath", "framebuffer", "hoverable", diff --git a/crates/re_log_types/src/component_types/instance_key.rs b/crates/re_log_types/src/component_types/instance_key.rs index 30d6b6b7893a..d2cbbee1e096 100644 --- a/crates/re_log_types/src/component_types/instance_key.rs +++ b/crates/re_log_types/src/component_types/instance_key.rs @@ -63,6 +63,16 @@ impl InstanceKey { pub fn specific_index(self) -> Option { self.is_specific().then_some(self) } + + /// Creates a new [`InstanceKey`] that identifies a 2d coordinate. + pub fn from_2d_image_coordinate([x, y]: [u32; 2], image_width: u64) -> Self { + Self((x as u64) + (y as u64) * image_width) + } + + /// Retrieves 2d image coordinates (x, y) encoded in an instance key + pub fn to_2d_image_coordinate(self, image_width: u64) -> [u32; 2] { + [(self.0 % image_width) as u32, (self.0 / image_width) as u32] + } } impl std::fmt::Display for InstanceKey { diff --git a/crates/re_renderer/examples/depth_cloud.rs b/crates/re_renderer/examples/depth_cloud.rs index 2003292731e6..8885ba0d4369 100644 --- a/crates/re_renderer/examples/depth_cloud.rs +++ b/crates/re_renderer/examples/depth_cloud.rs @@ -182,6 +182,7 @@ impl RenderDepthClouds { depth_data: depth.data.clone(), colormap: re_renderer::Colormap::Turbo, outline_mask_id: Default::default(), + picking_object_id: Default::default(), }], radius_boost_in_ui_points_for_outlines: 2.5, }, diff --git a/crates/re_renderer/shader/depth_cloud.wgsl b/crates/re_renderer/shader/depth_cloud.wgsl index 14caa6d8e8d4..1e7f7afdf0d7 100644 --- a/crates/re_renderer/shader/depth_cloud.wgsl +++ b/crates/re_renderer/shader/depth_cloud.wgsl @@ -28,6 +28,9 @@ struct DepthCloudInfo { /// Outline mask id for the outline mask pass. outline_mask_id: UVec2, + /// Picking object id that applies for the entire depth cloud. + picking_layer_object_id: UVec2, + /// Multiplier to get world-space depth from whatever is in the texture. world_depth_from_texture_value: f32, @@ -51,11 +54,23 @@ var depth_cloud_info: DepthCloudInfo; var depth_texture: texture_2d; struct VertexOut { - @builtin(position) pos_in_clip: Vec4, - @location(0) pos_in_world: Vec3, - @location(1) point_pos_in_world: Vec3, - @location(2) point_color: Vec4, - @location(3) point_radius: f32, + @builtin(position) + pos_in_clip: Vec4, + + @location(0) @interpolate(perspective) + pos_in_world: Vec3, + + @location(1) @interpolate(flat) + point_pos_in_world: Vec3, + + @location(2) @interpolate(flat) + point_color: Vec4, + + @location(3) @interpolate(flat) + point_radius: f32, + + @location(4) @interpolate(flat) + quad_idx: u32, }; // --- @@ -63,7 +78,7 @@ struct VertexOut { struct PointData { pos_in_world: Vec3, unresolved_radius: f32, - color: Vec4 + color: Vec4, } // Backprojects the depth texture using the intrinsics passed in the uniform buffer. @@ -75,6 +90,7 @@ fn compute_point_data(quad_idx: i32) -> PointData { let world_space_depth = depth_cloud_info.world_depth_from_texture_value * textureLoad(depth_texture, texcoords, 0).x; var data: PointData; + if 0.0 < world_space_depth && world_space_depth < f32max { // TODO(cmc): albedo textures let color = Vec4(colormap_linear(depth_cloud_info.colormap, world_space_depth / depth_cloud_info.max_depth_in_world), 1.0); @@ -113,6 +129,7 @@ fn vs_main(@builtin(vertex_index) vertex_idx: u32) -> VertexOut { var out: VertexOut; out.point_pos_in_world = point_data.pos_in_world; out.point_color = point_data.color; + out.quad_idx = u32(quad_idx); if 0.0 < point_data.unresolved_radius { // Span quad @@ -145,7 +162,7 @@ fn fs_main_picking_layer(in: VertexOut) -> @location(0) UVec4 { if coverage <= 0.5 { discard; } - return UVec4(0u, 0u, 0u, 0u); // TODO(andreas): Implement picking layer id pass-through. + return UVec4(depth_cloud_info.picking_layer_object_id, in.quad_idx, 0u); } @fragment diff --git a/crates/re_renderer/src/renderer/depth_cloud.rs b/crates/re_renderer/src/renderer/depth_cloud.rs index 6c56428ba1c4..285c0a2f9fd0 100644 --- a/crates/re_renderer/src/renderer/depth_cloud.rs +++ b/crates/re_renderer/src/renderer/depth_cloud.rs @@ -24,7 +24,7 @@ use crate::{ GpuRenderPipelineHandle, GpuTexture, PipelineLayoutDesc, RenderPipelineDesc, Texture2DBufferInfo, TextureDesc, }, - Colormap, OutlineMaskPreference, PickingLayerProcessor, + Colormap, OutlineMaskPreference, PickingLayerObjectId, PickingLayerProcessor, }; use super::{ @@ -35,7 +35,7 @@ use super::{ // --- mod gpu_data { - use crate::wgpu_buffer_types; + use crate::{wgpu_buffer_types, PickingLayerObjectId}; /// Keep in sync with mirror in `depth_cloud.wgsl.` #[repr(C, align(256))] @@ -47,6 +47,7 @@ mod gpu_data { pub depth_camera_intrinsics: wgpu_buffer_types::Mat3, pub outline_mask_id: wgpu_buffer_types::UVec2, + pub picking_layer_object_id: PickingLayerObjectId, /// Multiplier to get world-space depth from whatever is in the texture. pub world_depth_from_texture_value: f32, @@ -57,14 +58,13 @@ mod gpu_data { /// The maximum depth value in world-space, for use with the colormap. pub max_depth_in_world: f32, + /// Which colormap should be used. pub colormap: u32, /// Changes over different draw-phases. - pub radius_boost_in_ui_points: f32, + pub radius_boost_in_ui_points: wgpu_buffer_types::F32RowPadded, - pub row_pad: f32, - - pub end_padding: [wgpu_buffer_types::PaddingRow; 16 - 4 - 3 - 1 - 1], + pub end_padding: [wgpu_buffer_types::PaddingRow; 16 - 4 - 3 - 1 - 1 - 1], } impl DepthCloudInfoUBO { @@ -82,6 +82,7 @@ mod gpu_data { depth_data, colormap, outline_mask_id, + picking_object_id, } = depth_cloud; let user_depth_from_texture_value = match depth_data { @@ -99,8 +100,8 @@ mod gpu_data { point_radius_from_world_depth: *point_radius_from_world_depth, max_depth_in_world: *max_depth_in_world, colormap: *colormap as u32, - radius_boost_in_ui_points, - row_pad: Default::default(), + radius_boost_in_ui_points: radius_boost_in_ui_points.into(), + picking_layer_object_id: *picking_object_id, end_padding: Default::default(), } } @@ -164,6 +165,9 @@ pub struct DepthCloud { /// Option outline mask id preference. pub outline_mask_id: OutlineMaskPreference, + + /// Picking object id that applies for the entire depth cloud. + pub picking_object_id: PickingLayerObjectId, } impl DepthCloud { diff --git a/crates/re_renderer/src/resource_managers/texture_manager.rs b/crates/re_renderer/src/resource_managers/texture_manager.rs index b31f2e01f0dc..d4449d4d7c00 100644 --- a/crates/re_renderer/src/resource_managers/texture_manager.rs +++ b/crates/re_renderer/src/resource_managers/texture_manager.rs @@ -21,6 +21,16 @@ impl GpuTexture2DHandle { pub fn invalid() -> Self { Self(None) } + + /// Width of the texture, defaults to 1 if invalid since fallback textures are typically one pixel. + pub fn width(&self) -> u32 { + self.0.as_ref().map_or(1, |t| t.texture.width()) + } + + /// Height of the texture, defaults to 1 if invalid since fallback textures are typically one pixel. + pub fn height(&self) -> u32 { + self.0.as_ref().map_or(1, |t| t.texture.height()) + } } /// Data required to create a texture 2d resource. diff --git a/crates/re_viewer/src/ui/view_spatial/scene/mod.rs b/crates/re_viewer/src/ui/view_spatial/scene/mod.rs index 6ef4ef4cb50c..ef5082ad2fa8 100644 --- a/crates/re_viewer/src/ui/view_spatial/scene/mod.rs +++ b/crates/re_viewer/src/ui/view_spatial/scene/mod.rs @@ -56,7 +56,8 @@ pub struct MeshSource { } pub struct Image { - pub instance_path_hash: InstancePathHash, + /// Path to the image (note image instance ids would refer to pixels!) + pub ent_path: EntityPath, pub tensor: Tensor, diff --git a/crates/re_viewer/src/ui/view_spatial/scene/picking.rs b/crates/re_viewer/src/ui/view_spatial/scene/picking.rs index 59996c8e803c..fd20df7c733a 100644 --- a/crates/re_viewer/src/ui/view_spatial/scene/picking.rs +++ b/crates/re_viewer/src/ui/view_spatial/scene/picking.rs @@ -1,6 +1,7 @@ //! Handles picking in 2D & 3D spaces. use re_data_store::InstancePathHash; +use re_log_types::{component_types::InstanceKey, EntityPathHash}; use re_renderer::PickingLayerProcessor; use super::{SceneSpatialPrimitives, SceneSpatialUiData}; @@ -9,10 +10,10 @@ use crate::{ ui::view_spatial::eye::Eye, }; -#[derive(Clone)] +#[derive(Clone, PartialEq, Eq)] pub enum PickingHitType { - /// The hit was a textured rect at the given uv coordinates (ranging from 0 to 1) - TexturedRect(glam::Vec2), + /// The hit was a textured rect. + TexturedRect, /// The result came from GPU based picking. GpuPickingResult, @@ -217,7 +218,7 @@ fn picking_gpu( fn picking_textured_rects( context: &PickingContext, textured_rectangles: &[re_renderer::renderer::TexturedRect], - textured_rectangles_ids: &[InstancePathHash], + textured_rectangles_ids: &[EntityPathHash], ) -> Vec { crate::profile_function!(); @@ -249,9 +250,18 @@ fn picking_textured_rects( if (0.0..=1.0).contains(&u) && (0.0..=1.0).contains(&v) { hits.push(PickingRayHit { - instance_path_hash: *id, + instance_path_hash: InstancePathHash { + entity_path_hash: *id, + instance_key: InstanceKey::from_2d_image_coordinate( + [ + (u * rect.colormapped_texture.texture.width() as f32) as u32, + (v * rect.colormapped_texture.texture.height() as f32) as u32, + ], + rect.colormapped_texture.texture.width() as u64, + ), + }, space_position: intersection_world, - hit_type: PickingHitType::TexturedRect(glam::vec2(u, v)), + hit_type: PickingHitType::TexturedRect, depth_offset: rect.depth_offset, }); } diff --git a/crates/re_viewer/src/ui/view_spatial/scene/primitives.rs b/crates/re_viewer/src/ui/view_spatial/scene/primitives.rs index d8748a6bb36b..5fbd4448ff79 100644 --- a/crates/re_viewer/src/ui/view_spatial/scene/primitives.rs +++ b/crates/re_viewer/src/ui/view_spatial/scene/primitives.rs @@ -1,5 +1,6 @@ use egui::Color32; use re_data_store::InstancePathHash; +use re_log_types::EntityPathHash; use re_renderer::{ renderer::{DepthClouds, MeshInstance}, LineStripSeriesBuilder, PointCloudBuilder, @@ -21,7 +22,7 @@ pub struct SceneSpatialPrimitives { // TODO(andreas): Storing extra data like so is unsafe and not future proof either // (see also above comment on the need to separate cpu-readable data) - pub textured_rectangles_ids: Vec, + pub textured_rectangles_ids: Vec, pub textured_rectangles: Vec, pub line_strips: LineStripSeriesBuilder, diff --git a/crates/re_viewer/src/ui/view_spatial/scene/scene_part/images.rs b/crates/re_viewer/src/ui/view_spatial/scene/scene_part/images.rs index 0c2c29156e50..6bd65a7aa9f5 100644 --- a/crates/re_viewer/src/ui/view_spatial/scene/scene_part/images.rs +++ b/crates/re_viewer/src/ui/view_spatial/scene/scene_part/images.rs @@ -4,7 +4,7 @@ use egui::NumExt; use glam::Vec3; use itertools::Itertools; -use re_data_store::{query_latest_single, EntityPath, EntityProperties, InstancePathHash}; +use re_data_store::{query_latest_single, EntityPath, EntityProperties}; use re_log_types::{ component_types::{ColorRGBA, InstanceKey, Tensor, TensorData, TensorDataMeaning}, Component, Transform, @@ -19,7 +19,7 @@ use crate::{ misc::{SpaceViewHighlights, SpaceViewOutlineMasks, TransformCache, ViewerContext}, ui::{ scene::SceneQuery, - view_spatial::{scene::scene_part::instance_path_hash_for_picking, Image, SceneSpatial}, + view_spatial::{Image, SceneSpatial}, Annotations, DefaultColor, }, }; @@ -32,8 +32,7 @@ fn push_tensor_texture( ctx: &mut ViewerContext<'_>, annotations: &Arc, world_from_obj: glam::Mat4, - entity_path: &EntityPath, - instance_path_hash: InstancePathHash, + ent_path: &EntityPath, tensor: &Tensor, multiplicative_tint: egui::Rgba, outline_mask: OutlineMaskPreference, @@ -42,7 +41,7 @@ fn push_tensor_texture( let Some([height, width, _]) = tensor.image_height_width_channels() else { return; }; - let debug_name = entity_path.to_string(); + let debug_name = ent_path.to_string(); let tensor_stats = ctx.cache.tensor_stats(tensor); match crate::gpu_bridge::tensor_to_gpu( @@ -68,7 +67,7 @@ fn push_tensor_texture( scene .primitives .textured_rectangles_ids - .push(instance_path_hash); + .push(ent_path.hash()); } Err(err) => { re_log::error_once!("Failed to create texture from tensor for {debug_name:?}: {err}"); @@ -150,8 +149,8 @@ impl ImagesPart { ) -> Result<(), QueryError> { crate::profile_function!(); - for (instance_key, tensor, color) in itertools::izip!( - entity_view.iter_instance_keys()?, + // Instance ids of tensors refer to entries inside the tensor. + for (tensor, color) in itertools::izip!( entity_view.iter_primary()?, entity_view.iter_component::()? ) { @@ -161,6 +160,17 @@ impl ImagesPart { return Ok(()); } + let annotations = scene.annotation_map.find(ent_path); + + // TODO(jleibs): Meter should really be its own component + let meter = tensor.meter; + scene.ui.images.push(Image { + ent_path: ent_path.clone(), + tensor: tensor.clone(), + meter, + annotations: annotations.clone(), + }); + let entity_highlight = highlights.entity_outline_mask(ent_path.hash()); if *properties.backproject_depth.get() && tensor.meaning == TensorDataMeaning::Depth @@ -179,6 +189,7 @@ impl ImagesPart { transforms, properties, &tensor, + ent_path, &pinhole_ent_path, entity_highlight, ) { @@ -191,15 +202,14 @@ impl ImagesPart { } Self::process_entity_view_as_image( - entity_view, scene, ctx, ent_path, world_from_obj, entity_highlight, - instance_key, tensor, color, + &annotations, ); } } @@ -209,57 +219,34 @@ impl ImagesPart { #[allow(clippy::too_many_arguments)] fn process_entity_view_as_image( - entity_view: &EntityView, scene: &mut SceneSpatial, ctx: &mut ViewerContext<'_>, ent_path: &EntityPath, world_from_obj: glam::Mat4, entity_highlight: &SpaceViewOutlineMasks, - instance_key: InstanceKey, tensor: Tensor, color: Option, + annotations: &Arc, ) { crate::profile_function!(); - let instance_path_hash = instance_path_hash_for_picking( - ent_path, - instance_key, - entity_view, - entity_highlight.any_selection_highlight, - ); - - let annotations = scene.annotation_map.find(ent_path); - let color = annotations.class_description(None).annotation_info().color( color.map(|c| c.to_array()).as_ref(), DefaultColor::OpaqueWhite, ); - let outline_mask = entity_highlight.index_outline_mask(instance_path_hash.instance_key); - match ctx.cache.decode.try_decode_tensor_if_necessary(tensor) { Ok(tensor) => { push_tensor_texture( scene, ctx, - &annotations, + annotations, world_from_obj, ent_path, - instance_path_hash, &tensor, color.into(), - outline_mask, + entity_highlight.overall, ); - - // TODO(jleibs): Meter should really be its own component - let meter = tensor.meter; - - scene.ui.images.push(Image { - instance_path_hash, - tensor, - meter, - annotations, - }); } Err(err) => { // TODO(jleibs): Would be nice to surface these through the UI instead @@ -279,6 +266,7 @@ impl ImagesPart { transforms: &TransformCache, properties: &EntityProperties, tensor: &Tensor, + ent_path: &EntityPath, pinhole_ent_path: &EntityPath, entity_highlight: &SpaceViewOutlineMasks, ) -> Result<(), String> { @@ -360,6 +348,7 @@ impl ImagesPart { depth_data: data, colormap, outline_mask_id: entity_highlight.overall, + picking_object_id: re_renderer::PickingLayerObjectId(ent_path.hash64()), }); Ok(()) diff --git a/crates/re_viewer/src/ui/view_spatial/ui.rs b/crates/re_viewer/src/ui/view_spatial/ui.rs index 577175b7be48..2a02c662467b 100644 --- a/crates/re_viewer/src/ui/view_spatial/ui.rs +++ b/crates/re_viewer/src/ui/view_spatial/ui.rs @@ -742,31 +742,45 @@ pub fn picking( // TODO(#1818): Depth at pointer only works for depth images so far. let mut depth_at_pointer = None; for hit in &picking_result.hits { - let Some(instance_path) = hit.instance_path_hash.resolve(&ctx.log_db.entity_db) + let Some(mut instance_path) = hit.instance_path_hash.resolve(&ctx.log_db.entity_db) else { continue; }; - if !entity_properties - .get(&instance_path.entity_path) - .interactive - { + + let ent_properties = entity_properties.get(&instance_path.entity_path); + if !ent_properties.interactive { continue; } - hovered_items.push(crate::misc::Item::InstancePath( - Some(space_view_id), - instance_path.clone(), - )); // Special hover ui for images. - let picked_image_with_uv = if let PickingHitType::TexturedRect(uv) = hit.hit_type { + let picked_image_with_coords = if hit.hit_type == PickingHitType::TexturedRect + || *ent_properties.backproject_depth.get() + { + // We don't support selecting pixels yet. + instance_path.instance_key = re_log_types::component_types::InstanceKey::SPLAT; scene .ui .images .iter() - .find(|image| image.instance_path_hash == hit.instance_path_hash) - .map(|image| (image, uv)) + .find(|image| image.ent_path == instance_path.entity_path) + .and_then(|image| { + image.tensor.image_height_width_channels().map(|[_, w, _]| { + ( + image, + hit.instance_path_hash + .instance_key + .to_2d_image_coordinate(w), + ) + }) + }) } else { None }; - response = if let Some((image, uv)) = picked_image_with_uv { + + hovered_items.push(crate::misc::Item::InstancePath( + Some(space_view_id), + instance_path.clone(), + )); + + response = if let Some((image, coords)) = picked_image_with_coords { if let Some(meter) = image.meter { if let Some(raw_value) = image.tensor.get(&[ picking_context.pointer_in_space2d.y.round() as _, @@ -801,7 +815,6 @@ pub fn picking( ui.separator(); ui.horizontal(|ui| { let (w, h) = (w.size as f32, h.size as f32); - let center = [(uv.x * w) as isize, (uv.y * h) as isize]; if *state.nav_mode.get() == SpatialNavigationMode::TwoD { let rect = egui::Rect::from_min_size( egui::Pos2::ZERO, @@ -810,14 +823,14 @@ pub fn picking( data_ui::image::show_zoomed_image_region_area_outline( ui, &tensor_view, - center, + [coords[0] as _, coords[1] as _], space_from_ui.inverse().transform_rect(rect), ); } data_ui::image::show_zoomed_image_region( ui, &tensor_view, - center, + [coords[0] as _, coords[1] as _], image.meter, ); });