GPU based mesh picking in viewer (#1737)

* picking layer support for meshes
* move out view builder from render bridge into ui code
* remove now pointless create_scene_paint_callback method
* move screenshot taking out of renderer bridge, create view builder earlier in the ui build up process
* fix some issues in mesh picking and add mesh to picking example
* use gpu for mesh picking in the viewer
* debug option for debugging picking overlay
* no longer discard "classic" picking information
* placeholder picking layer implementations for remaining opaque primitives
* fix temporal gaps in gpu picking report

Author: Wumpf

crates/re_log_types/src/hash.rs

@@ -18,6 +18,12 @@ impl Hash64 {
         Self(hash(value))
     }
 
+    /// From an existing u64. Use this only for data conversions.
+    #[inline]
+    pub fn from_u64(i: u64) -> Self {
+        Self(i)
+    }
+
     #[inline]
     pub fn hash64(&self) -> u64 {
         self.0

crates/re_log_types/src/path/entity_path.rs

@@ -14,6 +14,12 @@ impl EntityPathHash {
     /// Sometimes used as the hash of `None`.
     pub const NONE: EntityPathHash = EntityPathHash(Hash64::ZERO);
 
+    /// From an existing u64. Use this only for data conversions.
+    #[inline]
+    pub fn from_u64(i: u64) -> Self {
+        Self(Hash64::from_u64(i))
+    }
+
     #[inline]
     pub fn hash64(&self) -> u64 {
         self.0.hash64()

crates/re_renderer/examples/picking.rs

@@ -1,9 +1,10 @@
 use itertools::Itertools as _;
 use rand::Rng;
 use re_renderer::{
+    renderer::MeshInstance,
     view_builder::{Projection, TargetConfiguration, ViewBuilder},
-    Color32, GpuReadbackIdentifier, IntRect, PickingLayerInstanceId, PickingLayerProcessor,
-    PointCloudBuilder, Size,
+    Color32, GpuReadbackIdentifier, IntRect, PickingLayerId, PickingLayerInstanceId,
+    PickingLayerProcessor, PointCloudBuilder, Size,
 };
 
 mod framework;
@@ -18,6 +19,8 @@ struct PointSet {
 struct Picking {
     point_sets: Vec<PointSet>,
     picking_position: glam::UVec2,
+    model_mesh_instances: Vec<MeshInstance>,
+    mesh_is_hovered: bool,
 }
 
 fn random_color(rnd: &mut impl rand::Rng) -> Color32 {
@@ -34,6 +37,12 @@ fn random_color(rnd: &mut impl rand::Rng) -> Color32 {
 /// Identifiers don't need to be unique and we don't have anything interesting to distinguish here!
 const READBACK_IDENTIFIER: GpuReadbackIdentifier = 0;
 
+/// Mesh ID used for picking. Uses the entire 64bit range for testing.
+const MESH_ID: PickingLayerId = PickingLayerId {
+    object: re_renderer::PickingLayerObjectId(0x1234_5678_9012_3456),
+    instance: re_renderer::PickingLayerInstanceId(0x3456_1234_5678_9012),
+};
+
 impl framework::Example for Picking {
     fn title() -> &'static str {
         "Picking"
@@ -43,10 +52,10 @@ impl framework::Example for Picking {
         self.picking_position = position_in_pixel;
     }
 
-    fn new(_re_ctx: &mut re_renderer::RenderContext) -> Self {
+    fn new(re_ctx: &mut re_renderer::RenderContext) -> Self {
         let mut rnd = <rand::rngs::StdRng as rand::SeedableRng>::seed_from_u64(42);
         let random_point_range = -5.0_f32..5.0_f32;
-        let point_count = 10000;
+        let point_count = 1000;
 
         // Split point cloud into several batches to test picking of multiple objects.
         let point_sets = (0..2)
@@ -72,9 +81,13 @@ impl framework::Example for Picking {
             })
             .collect_vec();
 
+        let model_mesh_instances = crate::framework::load_rerun_mesh(re_ctx);
+
         Picking {
             point_sets,
+            model_mesh_instances,
             picking_position: glam::UVec2::ZERO,
+            mesh_is_hovered: false,
         }
     }
 
@@ -93,7 +106,12 @@ impl framework::Example for Picking {
                 + (picking_result.rect.extent.y / 2) * picking_result.rect.extent.x)
                 as usize];
 
-            if picked_pixel.object.0 != 0 {
+            self.mesh_is_hovered = false;
+            if picked_pixel == MESH_ID {
+                self.mesh_is_hovered = true;
+            } else if picked_pixel.object.0 != 0
+                && picked_pixel.object.0 <= self.point_sets.len() as u64
+            {
                 let point_set = &mut self.point_sets[picked_pixel.object.0 as usize - 1];
                 point_set.radii[picked_pixel.instance.0 as usize] = Size::new_scene(0.1);
                 point_set.colors[picked_pixel.instance.0 as usize] = Color32::DEBUG_COLOR;
@@ -139,10 +157,9 @@ impl framework::Example for Picking {
             .schedule_picking_rect(re_ctx, picking_rect, READBACK_IDENTIFIER, (), false)
             .unwrap();
 
-        let mut builder = PointCloudBuilder::<()>::new(re_ctx);
-
+        let mut point_builder = PointCloudBuilder::<()>::new(re_ctx);
         for (i, point_set) in self.point_sets.iter().enumerate() {
-            builder
+            point_builder
                 .batch(format!("Random Points {i}"))
                 .picking_object_id(re_renderer::PickingLayerObjectId(i as u64 + 1)) // offset by one since 0=default=no hit
                 .add_points(
@@ -153,7 +170,29 @@ impl framework::Example for Picking {
                 .colors(point_set.colors.iter().cloned())
                 .picking_instance_ids(point_set.picking_ids.iter().cloned());
         }
-        view_builder.queue_draw(&builder.to_draw_data(re_ctx).unwrap());
+        view_builder.queue_draw(&point_builder.to_draw_data(re_ctx).unwrap());
+
+        let instances = self
+            .model_mesh_instances
+            .iter()
+            .map(|instance| MeshInstance {
+                gpu_mesh: instance.gpu_mesh.clone(),
+                mesh: None,
+                world_from_mesh: glam::Affine3A::from_translation(glam::vec3(0.0, 0.0, 0.0)),
+                picking_layer_id: MESH_ID,
+                additive_tint: if self.mesh_is_hovered {
+                    Color32::DEBUG_COLOR
+                } else {
+                    Color32::TRANSPARENT
+                },
+                ..Default::default()
+            })
+            .collect_vec();
+
+        view_builder.queue_draw(&re_renderer::renderer::GenericSkyboxDrawData::new(re_ctx));
+        view_builder
+            .queue_draw(&re_renderer::renderer::MeshDrawData::new(re_ctx, &instances).unwrap());
+
         view_builder.queue_draw(&re_renderer::renderer::GenericSkyboxDrawData::new(re_ctx));
 
         let command_buffer = view_builder

crates/re_renderer/shader/depth_cloud.wgsl

@@ -139,6 +139,15 @@ fn fs_main(in: VertexOut) -> @location(0) Vec4 {
     return vec4(in.point_color.rgb, coverage);
 }
 
+@fragment
+fn fs_main_picking_layer(in: VertexOut) -> @location(0) UVec4 {
+    let coverage = sphere_quad_coverage(in.pos_in_world, in.point_radius, in.point_pos_in_world);
+    if coverage <= 0.5 {
+        discard;
+    }
+    return UVec4(0u, 0u, 0u, 0u); // TODO(andreas): Implement picking layer id pass-through.
+}
+
 @fragment
 fn fs_main_outline_mask(in: VertexOut) -> @location(0) UVec2 {
     // Output is an integer target, can't use coverage therefore.

crates/re_renderer/shader/instanced_mesh.wgsl

@@ -15,13 +15,26 @@ struct MaterialUniformBuffer {
 var<uniform> material: MaterialUniformBuffer;
 
 struct VertexOut {
-    @builtin(position) position: Vec4,
-    @location(0) color: Vec4, // 0-1 linear space with unmultiplied/separate alpha
-    @location(1) texcoord: Vec2,
-    @location(2) normal_world_space: Vec3,
-    @location(3) additive_tint_rgb: Vec3, // 0-1 linear space
+    @builtin(position)
+    position: Vec4,
+
+    @location(0)
+    color: Vec4, // 0-1 linear space with unmultiplied/separate alpha
+
+    @location(1)
+    texcoord: Vec2,
+
+    @location(2)
+    normal_world_space: Vec3,
+
+    @location(3) @interpolate(flat)
+    additive_tint_rgb: Vec3, // 0-1 linear space
+
     @location(4) @interpolate(flat)
     outline_mask_ids: UVec2,
+
+    @location(5) @interpolate(flat)
+    picking_layer_id: UVec4,
 };
 
 @vertex
@@ -44,6 +57,7 @@ fn vs_main(in_vertex: VertexIn, in_instance: InstanceIn) -> VertexOut {
     out.normal_world_space = world_normal;
     out.additive_tint_rgb = linear_from_srgb(in_instance.additive_tint_srgb.rgb);
     out.outline_mask_ids = in_instance.outline_mask_ids;
+    out.picking_layer_id = in_instance.picking_layer_id;
 
     return out;
 }
@@ -70,6 +84,11 @@ fn fs_main_shaded(in: VertexOut) -> @location(0) Vec4 {
     }
 }
 
+@fragment
+fn fs_main_picking_layer(in: VertexOut) -> @location(0) UVec4 {
+    return in.picking_layer_id;
+}
+
 @fragment
 fn fs_main_outline_mask(in: VertexOut) -> @location(0) UVec2 {
     return in.outline_mask_ids;

crates/re_renderer/shader/lines.wgsl

@@ -299,6 +299,15 @@ fn fs_main(in: VertexOut) -> @location(0) Vec4 {
     return Vec4(in.color.rgb * shading, coverage);
 }
 
+@fragment
+fn fs_main_picking_layer(in: VertexOut) -> @location(0) UVec4 {
+    var coverage = compute_coverage(in);
+    if coverage < 0.5 {
+        discard;
+    }
+    return UVec4(0u, 0u, 0u, 0u); // TODO(andreas): Implement picking layer id pass-through.
+}
+
 @fragment
 fn fs_main_outline_mask(in: VertexOut) -> @location(0) UVec2 {
     // Output is an integer target, can't use coverage therefore.

crates/re_renderer/shader/mesh_vertex.wgsl

@@ -17,5 +17,6 @@ struct InstanceIn {
     @location(8) world_from_mesh_normal_row_1: Vec3,
     @location(9) world_from_mesh_normal_row_2: Vec3,
     @location(10) additive_tint_srgb: Vec4,
-    @location(11) outline_mask_ids: UVec2,
+    @location(11) picking_layer_id: UVec4,
+    @location(12) outline_mask_ids: UVec2,
 };

crates/re_renderer/shader/rectangle.wgsl

@@ -54,6 +54,11 @@ fn fs_main(in: VertexOut) -> @location(0) Vec4 {
     return texture_color * rect_info.multiplicative_tint;
 }
 
+@fragment
+fn fs_main_picking_layer(in: VertexOut) -> @location(0) UVec4 {
+    return UVec4(0u, 0u, 0u, 0u); // TODO(andreas): Implement picking layer id pass-through.
+}
+
 @fragment
 fn fs_main_outline_mask(in: VertexOut) -> @location(0) UVec2 {
     return rect_info.outline_mask;

crates/re_renderer/src/draw_phases/picking_layer.rs

@@ -5,6 +5,7 @@
 //!
 //! The picking layer is a RGBA texture with 32bit per channel, the red & green channel are used for the [`PickingLayerObjectId`],
 //! the blue & alpha channel are used for the [`PickingLayerInstanceId`].
+//! (Keep in mind that GPUs are little endian, so R will have the lower bytes and G the higher ones)
 //!
 //! In order to accomplish small render targets, the projection matrix is cropped to only render the area of interest.
 
@@ -43,27 +44,38 @@ struct ReadbackBeltMetadata<T: 'static + Send + Sync> {
 /// Typically used to identify higher level objects
 /// Some renderers might allow to change this part of the picking identifier only at a coarse grained level.
 #[repr(C)]
-#[derive(Clone, Copy, bytemuck::Zeroable, bytemuck::Pod, Default, Debug)]
+#[derive(Clone, Copy, bytemuck::Zeroable, bytemuck::Pod, Default, Debug, PartialEq, Eq)]
 pub struct PickingLayerObjectId(pub u64);
 
 /// The second 64bit of the picking layer.
 ///
 /// Typically used to identify instances.
 /// Some renderers might allow to change only this part of the picking identifier at a fine grained level.
 #[repr(C)]
-#[derive(Clone, Copy, bytemuck::Zeroable, bytemuck::Pod, Default, Debug)]
+#[derive(Clone, Copy, bytemuck::Zeroable, bytemuck::Pod, Default, Debug, PartialEq, Eq)]
 pub struct PickingLayerInstanceId(pub u64);
 
 /// Combination of `PickingLayerObjectId` and `PickingLayerInstanceId`.
 ///
 /// This is the same memory order as it is found in the GPU picking layer texture.
 #[repr(C)]
-#[derive(Clone, Copy, bytemuck::Zeroable, bytemuck::Pod, Default, Debug)]
+#[derive(Clone, Copy, bytemuck::Zeroable, bytemuck::Pod, Default, Debug, PartialEq, Eq)]
 pub struct PickingLayerId {
     pub object: PickingLayerObjectId,
     pub instance: PickingLayerInstanceId,
 }
 
+impl From<PickingLayerId> for [u32; 4] {
+    fn from(val: PickingLayerId) -> Self {
+        [
+            val.object.0 as u32,
+            (val.object.0 >> 32) as u32,
+            val.instance.0 as u32,
+            (val.instance.0 >> 32) as u32,
+        ]
+    }
+}
+
 /// Manages the rendering of the picking layer pass, its render targets & readback buffer.
 ///
 /// The view builder creates this for every frame that requests a picking result.

crates/re_renderer/src/renderer/depth_cloud.rs

@@ -24,7 +24,7 @@ use crate::{
         GpuRenderPipelineHandle, GpuTexture, PipelineLayoutDesc, RenderPipelineDesc, TextureDesc,
         TextureRowDataInfo,
     },
-    ColorMap, OutlineMaskPreference,
+    ColorMap, OutlineMaskPreference, PickingLayerProcessor,
 };
 
 use super::{
@@ -387,6 +387,7 @@ fn create_and_upload_texture<T: bytemuck::Pod>(
 
 pub struct DepthCloudRenderer {
     render_pipeline_color: GpuRenderPipelineHandle,
+    render_pipeline_picking_layer: GpuRenderPipelineHandle,
     render_pipeline_outline_mask: GpuRenderPipelineHandle,
     bind_group_layout: GpuBindGroupLayoutHandle,
 }
@@ -395,7 +396,11 @@ impl Renderer for DepthCloudRenderer {
     type RendererDrawData = DepthCloudDrawData;
 
     fn participated_phases() -> &'static [DrawPhase] {
-        &[DrawPhase::OutlineMask, DrawPhase::Opaque]
+        &[
+            DrawPhase::Opaque,
+            DrawPhase::PickingLayer,
+            DrawPhase::OutlineMask,
+        ]
     }
 
     fn create_renderer<Fs: FileSystem>(
@@ -480,7 +485,19 @@ impl Renderer for DepthCloudRenderer {
             &pools.pipeline_layouts,
             &pools.shader_modules,
         );
-
+        let render_pipeline_picking_layer = pools.render_pipelines.get_or_create(
+            device,
+            &RenderPipelineDesc {
+                label: "DepthCloudRenderer::render_pipeline_picking_layer".into(),
+                fragment_entrypoint: "fs_main_picking_layer".into(),
+                render_targets: smallvec![Some(PickingLayerProcessor::PICKING_LAYER_FORMAT.into())],
+                depth_stencil: PickingLayerProcessor::PICKING_LAYER_DEPTH_STATE,
+                multisample: PickingLayerProcessor::PICKING_LAYER_MSAA_STATE,
+                ..render_pipeline_desc_color.clone()
+            },
+            &pools.pipeline_layouts,
+            &pools.shader_modules,
+        );
         let render_pipeline_outline_mask = pools.render_pipelines.get_or_create(
             device,
             &RenderPipelineDesc {
@@ -500,6 +517,7 @@ impl Renderer for DepthCloudRenderer {
 
         DepthCloudRenderer {
             render_pipeline_color,
+            render_pipeline_picking_layer,
             render_pipeline_outline_mask,
             bind_group_layout,
         }
@@ -518,8 +536,9 @@ impl Renderer for DepthCloudRenderer {
         }
 
         let pipeline_handle = match phase {
-            DrawPhase::OutlineMask => self.render_pipeline_outline_mask,
             DrawPhase::Opaque => self.render_pipeline_color,
+            DrawPhase::PickingLayer => self.render_pipeline_picking_layer,
+            DrawPhase::OutlineMask => self.render_pipeline_outline_mask,
             _ => unreachable!("We were called on a phase we weren't subscribed to: {phase:?}"),
         };
         let pipeline = pools.render_pipelines.get_resource(pipeline_handle)?;