GPU based picking with points

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,205 @@
+use itertools::Itertools as _;
+use rand::Rng;
+use re_renderer::{
+    view_builder::{Projection, TargetConfiguration, ViewBuilder},
+    Color32, PickingLayerId, PickingLayerInstanceId, PointCloudBuilder, RenderContext,
+    ScheduledPickingRect, Size,
+};
+
+mod framework;
+
+struct Picking {
+    points_positions: Vec<glam::Vec3>,
+    points_radii: Vec<Size>,
+    points_colors: Vec<Color32>,
+    points_picking_ids: Vec<PickingLayerInstanceId>,
+
+    scheduled_picking_rects: Vec<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(index) = self
+                    .scheduled_picking_rects
+                    .iter()
+                    .position(|s| s.identifier == identifier)
+                {
+                    let picking_rect_info = self.scheduled_picking_rects.swap_remove(index);
+
+                    // 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.extent.x / 2
+                        + (picking_rect_info.extent.y / 2) * picking_rect_info.extent.x)
+                        as usize];
+                    if picked_pixel.object.0 != 0 {
+                        let index = picked_pixel.instance.0
+                            + (self.points_radii.len() as u64 / 2) * (picked_pixel.object.0 - 1);
+                        self.points_radii[index as usize] = Size::new_scene(0.1);
+                        self.points_colors[index 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;
+        let points_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();
+        let points_radii = std::iter::repeat(Size::new_scene(0.08))
+            .take(point_count)
+            .collect_vec();
+        let points_colors = (0..point_count)
+            .map(|_| random_color(&mut rnd))
+            .collect_vec();
+        let points_picking_ids = (0..point_count as u64)
+            .map(PickingLayerInstanceId)
+            .collect_vec();
+
+        Picking {
+            points_positions,
+            points_radii,
+            points_colors,
+            points_picking_ids,
+            scheduled_picking_rects: Vec::new(),
+            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();
+
+        let picking_rect_size = 32;
+        self.scheduled_picking_rects.push(
+            view_builder
+                .schedule_picking_readback(
+                    re_ctx,
+                    self.picking_position.as_ivec2()
+                        - glam::ivec2(picking_rect_size / 2, picking_rect_size / 2),
+                    picking_rect_size as u32,
+                    false,
+                )
+                .unwrap(),
+        );
+
+        let mut builder = PointCloudBuilder::<()>::new(re_ctx);
+
+        // Split into two batches to test picking of multiple batches.
+        let num_per_batch = self.points_positions.len() / 2;
+        builder
+            .batch("Random Points 1")
+            .picking_object_id(re_renderer::PickingLayerObjectId(1))
+            .add_points(
+                num_per_batch,
+                self.points_positions.iter().take(num_per_batch).cloned(),
+            )
+            .radii(self.points_radii.iter().take(num_per_batch).cloned())
+            .colors(self.points_colors.iter().take(num_per_batch).cloned())
+            .picking_instance_ids(self.points_picking_ids.iter().cloned());
+        builder
+            .batch("Random Points 2")
+            .picking_object_id(re_renderer::PickingLayerObjectId(2))
+            .add_points(
+                num_per_batch,
+                self.points_positions.iter().skip(num_per_batch).cloned(),
+            )
+            .radii(self.points_radii.iter().skip(num_per_batch).cloned())
+            .colors(self.points_colors.iter().skip(num_per_batch).cloned())
+            .picking_instance_ids(self.points_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,55 @@
+#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`
+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_portion = texcoord * (uniforms.extent_in_pixel / uniforms.screen_resolution) +
+                        uniforms.position_in_pixel / uniforms.screen_resolution;
+    let screen_ndc = Vec2(screen_portion.x * 2.0 - 1.0, 1.0 - screen_portion.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 num_color_levels = 20u;
+        let coords = IVec2(in.texcoord * Vec2(textureDimensions(debug_texture_uint).xy));
+        return Vec4(Vec3(textureLoad(debug_texture_uint, coords, 0).rga % num_color_levels) / f32(num_color_levels - 1u), 1.0);
+    } else {
+        return Vec4(1.0, 0.0, 1.0, 1.0);
+    }
+}