GPU based picking with points (#1721)

* a working re_renderer example for picking in point clouds
* a debug overlay renderer
* a (incomplete) picking draw phase, fully implemented for point clouds
* the ability to schedule readback transfers for picking rects that are rendered on demand to a small texture
* refactor draw phase processing out to its own module (we should formalize things things some more later!)

crates/re_renderer/examples/framework.rs

@@ -36,7 +36,9 @@ pub trait Example {
         pixels_from_point: f32,
     ) -> Vec<ViewDrawResult>;
 
-    fn on_keyboard_input(&mut self, input: winit::event::KeyboardInput);
+    fn on_keyboard_input(&mut self, _input: winit::event::KeyboardInput) {}
+
+    fn on_cursor_moved(&mut self, _position_in_pixel: glam::UVec2) {}
 }
 
 #[allow(dead_code)]
@@ -205,6 +207,13 @@ impl<E: Example + 'static> Application<E> {
                     event: WindowEvent::KeyboardInput { input, .. },
                     ..
                 } => self.example.on_keyboard_input(input),
+                Event::WindowEvent {
+                    event: WindowEvent::CursorMoved { position, .. },
+                    ..
+                } => self.example.on_cursor_moved(glam::uvec2(
+                    position.x.round() as u32,
+                    position.y.round() as u32,
+                )),
                 Event::WindowEvent {
                     event:
                         WindowEvent::ScaleFactorChanged {

crates/re_renderer/examples/multiview.rs

@@ -216,17 +216,7 @@ impl Multiview {
                     {
                         let screenshot = self.scheduled_screenshots.swap_remove(index);
 
-                        // Need to do a memcpy to remove the padding.
-                        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)],
-                            );
-                        }
+                        re_log::info!("Received screenshot. Total bytes {}", data.len());
 
                         // Get next available file name.
                         let mut i = 1;
@@ -238,11 +228,12 @@ impl Multiview {
                             i += 1;
                         };
 
+                        #[cfg(not(target_arch = "wasm32"))]
                         image::save_buffer(
                             filename,
-                            &buffer,
-                            screenshot.width,
-                            screenshot.height,
+                            &screenshot.row_info.remove_padding(data),
+                            screenshot.extent.x,
+                            screenshot.extent.y,
                             image::ColorType::Rgba8,
                         )
                         .expect("Failed to save screenshot");

crates/re_renderer/examples/outlines.rs

@@ -1,7 +1,8 @@
 use itertools::Itertools;
 use re_renderer::{
-    renderer::{MeshInstance, OutlineConfig, OutlineMaskPreference},
+    renderer::MeshInstance,
     view_builder::{Projection, TargetConfiguration, ViewBuilder},
+    OutlineConfig, OutlineMaskPreference,
 };
 use winit::event::{ElementState, VirtualKeyCode};
 

crates/re_renderer/examples/picking.rs

@@ -0,0 +1,198 @@
+use ahash::HashMap;
+use itertools::Itertools as _;
+use rand::Rng;
+use re_renderer::{
+    view_builder::{Projection, TargetConfiguration, ViewBuilder},
+    Color32, GpuReadbackBufferIdentifier, IntRect, PickingLayerId, PickingLayerInstanceId,
+    PointCloudBuilder, RenderContext, ScheduledPickingRect, Size,
+};
+
+mod framework;
+
+struct PointSet {
+    positions: Vec<glam::Vec3>,
+    radii: Vec<Size>,
+    colors: Vec<Color32>,
+    picking_ids: Vec<PickingLayerInstanceId>,
+}
+
+struct Picking {
+    point_sets: Vec<PointSet>,
+    scheduled_picking_rects: HashMap<GpuReadbackBufferIdentifier, ScheduledPickingRect>,
+    picking_position: glam::UVec2,
+}
+
+fn random_color(rnd: &mut impl rand::Rng) -> Color32 {
+    ecolor::Hsva {
+        h: rnd.gen::<f32>(),
+        s: rnd.gen::<f32>() * 0.5 + 0.5,
+        v: rnd.gen::<f32>() * 0.5 + 0.5,
+        a: 1.0,
+    }
+    .into()
+}
+
+impl Picking {
+    #[allow(clippy::unused_self)]
+    fn handle_incoming_picking_data(&mut self, re_ctx: &mut RenderContext, _time: f32) {
+        re_ctx
+            .gpu_readback_belt
+            .lock()
+            .receive_data(|data, identifier| {
+                if let Some(picking_rect_info) = self.scheduled_picking_rects.remove(&identifier) {
+                    // TODO(andreas): Move this into a utility function?
+                    let picking_data_without_padding =
+                        picking_rect_info.row_info.remove_padding(data);
+                    let picking_data: &[PickingLayerId] =
+                        bytemuck::cast_slice(&picking_data_without_padding);
+
+                    // Grab the middle pixel. usually we'd want to do something clever that snaps the the closest object of interest.
+                    let picked_pixel = picking_data[(picking_rect_info.rect.extent.x / 2
+                        + (picking_rect_info.rect.extent.y / 2) * picking_rect_info.rect.extent.x)
+                        as usize];
+
+                    if picked_pixel.object.0 != 0 {
+                        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;
+                    }
+                } else {
+                    re_log::error!("Received picking data for unknown identifier");
+                }
+            });
+    }
+}
+
+impl framework::Example for Picking {
+    fn title() -> &'static str {
+        "Picking"
+    }
+
+    fn on_cursor_moved(&mut self, position_in_pixel: glam::UVec2) {
+        self.picking_position = position_in_pixel;
+    }
+
+    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;
+
+        // Split point cloud into several batches to test picking of multiple objects.
+        let point_sets = (0..2)
+            .map(|_| PointSet {
+                positions: (0..point_count)
+                    .map(|_| {
+                        glam::vec3(
+                            rnd.gen_range(random_point_range.clone()),
+                            rnd.gen_range(random_point_range.clone()),
+                            rnd.gen_range(random_point_range.clone()),
+                        )
+                    })
+                    .collect_vec(),
+                radii: std::iter::repeat(Size::new_scene(0.08))
+                    .take(point_count)
+                    .collect_vec(),
+                colors: (0..point_count)
+                    .map(|_| random_color(&mut rnd))
+                    .collect_vec(),
+                picking_ids: (0..point_count as u64)
+                    .map(PickingLayerInstanceId)
+                    .collect_vec(),
+            })
+            .collect_vec();
+
+        Picking {
+            point_sets,
+            scheduled_picking_rects: HashMap::default(),
+            picking_position: glam::UVec2::ZERO,
+        }
+    }
+
+    fn draw(
+        &mut self,
+        re_ctx: &mut re_renderer::RenderContext,
+        resolution: [u32; 2],
+        time: &framework::Time,
+        pixels_from_point: f32,
+    ) -> Vec<framework::ViewDrawResult> {
+        self.handle_incoming_picking_data(re_ctx, time.seconds_since_startup());
+
+        let mut view_builder = ViewBuilder::default();
+
+        // TODO(#1426): unify camera logic between examples.
+        let camera_position = glam::vec3(1.0, 3.5, 7.0);
+
+        view_builder
+            .setup_view(
+                re_ctx,
+                TargetConfiguration {
+                    name: "OutlinesDemo".into(),
+                    resolution_in_pixel: resolution,
+                    view_from_world: macaw::IsoTransform::look_at_rh(
+                        camera_position,
+                        glam::Vec3::ZERO,
+                        glam::Vec3::Y,
+                    )
+                    .unwrap(),
+                    projection_from_view: Projection::Perspective {
+                        vertical_fov: 70.0 * std::f32::consts::TAU / 360.0,
+                        near_plane_distance: 0.01,
+                    },
+                    pixels_from_point,
+                    outline_config: None,
+                    ..Default::default()
+                },
+            )
+            .unwrap();
+
+        // Use an uneven number of pixels for the picking rect so that there is a clearly defined middle-pixel.
+        // (for this sample a size of 1 would be sufficient, but for a real application you'd want to use a larger size to allow snapping)
+        let picking_rect_size = 31;
+
+        let picking_rect = view_builder
+            .schedule_picking_readback(
+                re_ctx,
+                IntRect::from_middle_and_extent(
+                    self.picking_position.as_ivec2(),
+                    glam::uvec2(picking_rect_size, picking_rect_size),
+                ),
+                false,
+            )
+            .unwrap();
+        self.scheduled_picking_rects
+            .insert(picking_rect.identifier, picking_rect);
+
+        let mut builder = PointCloudBuilder::<()>::new(re_ctx);
+
+        for (i, point_set) in self.point_sets.iter().enumerate() {
+            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(
+                    point_set.positions.len(),
+                    point_set.positions.iter().cloned(),
+                )
+                .radii(point_set.radii.iter().cloned())
+                .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(&re_renderer::renderer::GenericSkyboxDrawData::new(re_ctx));
+
+        let command_buffer = view_builder
+            .draw(re_ctx, ecolor::Rgba::TRANSPARENT)
+            .unwrap();
+
+        vec![framework::ViewDrawResult {
+            view_builder,
+            command_buffer,
+            target_location: glam::Vec2::ZERO,
+        }]
+    }
+
+    fn on_keyboard_input(&mut self, _input: winit::event::KeyboardInput) {}
+}
+
+fn main() {
+    framework::start::<Picking>();
+}

crates/re_renderer/shader/debug_overlay.wgsl

@@ -0,0 +1,68 @@
+// Debug overlay shader
+//
+// Works together with `debug_overlay.rs` to display a texture on top of the screen.
+// It is meant to be used as last part of the compositor phase in order to present the debug output unfiltered.
+// It's sole purpose is for developing new rendering features and it should not be used in production!
+//
+// The fragment shader is a blueprint for handling different texture outputs.
+// *Do* edit it on the fly for debugging purposes!
+
+#import <./types.wgsl>
+#import <./global_bindings.wgsl>
+
+struct UniformBuffer {
+    screen_resolution: Vec2,
+    position_in_pixel: Vec2,
+    extent_in_pixel: Vec2,
+    mode: u32,
+    _padding: u32,
+};
+@group(1) @binding(0)
+var<uniform> uniforms: UniformBuffer;
+
+@group(1) @binding(1)
+var debug_texture_float: texture_2d<f32>;
+@group(1) @binding(2)
+var debug_texture_uint: texture_2d<u32>;
+
+// Mode options, see `DebugOverlayMode` in `debug_overlay.rs`
+const ShowFloatTexture: u32 = 0u;
+const ShowUintTexture: u32 = 1u;
+
+struct VertexOutput {
+    @builtin(position) position: Vec4,
+    @location(0) texcoord: Vec2,
+};
+
+@vertex
+fn main_vs(@builtin(vertex_index) vertex_index: u32) -> VertexOutput {
+    let texcoord = Vec2(f32(vertex_index / 2u), f32(vertex_index % 2u));
+
+    // This calculation could be simplified by pre-computing things on the CPU.
+    // But this is not the point here - we want to debug this and other things rapidly by editing the shader.
+    let screen_fraction = texcoord * (uniforms.extent_in_pixel / uniforms.screen_resolution) +
+                        uniforms.position_in_pixel / uniforms.screen_resolution;
+    let screen_ndc = Vec2(screen_fraction.x * 2.0 - 1.0, 1.0 - screen_fraction.y * 2.0);
+
+    var out: VertexOutput;
+    out.position = Vec4(screen_ndc, 0.0, 1.0);
+    out.texcoord = texcoord;
+    return out;
+}
+
+@fragment
+fn main_fs(in: VertexOutput) -> @location(0) Vec4 {
+    if uniforms.mode == ShowFloatTexture {
+        return Vec4(textureSample(debug_texture_float, nearest_sampler, in.texcoord).rgb, 1.0);
+    } else if uniforms.mode == ShowUintTexture {
+        let coords = IVec2(in.texcoord * Vec2(textureDimensions(debug_texture_uint).xy));
+        let raw_values = textureLoad(debug_texture_uint, coords, 0);
+
+        let num_color_levels = 20u;
+        let mapped_values = Vec4(raw_values % num_color_levels) / f32(num_color_levels - 1u);
+
+        return Vec4(mapped_values.rgb, 1.0);
+    } else {
+        return Vec4(1.0, 0.0, 1.0, 1.0);
+    }
+}