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 <[email protected]>

.vscode/settings.json

@@ -25,6 +25,7 @@
         "andreas",
         "bbox",
         "bindgroup",
+        "colormap",
         "emath",
         "framebuffer",
         "hoverable",

crates/re_log_types/src/component_types/instance_key.rs

@@ -63,6 +63,16 @@ impl InstanceKey {
     pub fn specific_index(self) -> Option<InstanceKey> {
         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 {

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,
             },

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,19 +54,31 @@ var<uniform> depth_cloud_info: DepthCloudInfo;
 var depth_texture: texture_2d<f32>;
 
 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,
 };
 
 // ---
 
 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

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 {

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.

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,
 

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<PickingRayHit> {
     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,
             });
         }

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<InstancePathHash>,
+    pub textured_rectangles_ids: Vec<EntityPathHash>,
     pub textured_rectangles: Vec<re_renderer::renderer::TexturedRect>,
 
     pub line_strips: LineStripSeriesBuilder,