Integrate GPU color maps into depth clouds (#1486)

* integrate color maps into depth clouds

* clean mess

* pad

Author: teh-cmc

crates/re_data_store/src/entity_properties.rs

@@ -47,10 +47,16 @@ pub struct EntityProperties {
     pub visible_history: ExtraQueryHistory,
     pub interactive: bool,
 
+    /// Enable color mapping?
+    ///
+    /// See [`Self::color_mapper`] to select an actual mapping.
+    pub color_mapping: bool,
+    /// What kind of color mapping should be applied (none, map, texture, transfer..)?
+    pub color_mapper: EditableAutoValue<ColorMapper>,
+
     /// Distance of the projection plane (frustum far plane).
     ///
     /// Only applies to pinhole cameras when in a spatial view, using 3D navigation.
-    ///
     pub pinhole_image_plane_distance: EditableAutoValue<f32>,
 
     /// Should the depth texture be backprojected into a point cloud?
@@ -76,10 +82,15 @@ impl EntityProperties {
             visible: self.visible && child.visible,
             visible_history: self.visible_history.with_child(&child.visible_history),
             interactive: self.interactive && child.interactive,
+
+            color_mapping: self.color_mapping || child.color_mapping,
+            color_mapper: self.color_mapper.or(&child.color_mapper).clone(),
+
             pinhole_image_plane_distance: self
                 .pinhole_image_plane_distance
                 .or(&child.pinhole_image_plane_distance)
                 .clone(),
+
             backproject_depth: self.backproject_depth || child.backproject_depth,
             backproject_pinhole_ent_path: self
                 .backproject_pinhole_ent_path
@@ -101,6 +112,8 @@ impl Default for EntityProperties {
             visible: true,
             visible_history: ExtraQueryHistory::default(),
             interactive: true,
+            color_mapping: false,
+            color_mapper: EditableAutoValue::default(),
             pinhole_image_plane_distance: EditableAutoValue::default(),
             backproject_depth: false,
             backproject_pinhole_ent_path: None,
@@ -133,6 +146,53 @@ impl ExtraQueryHistory {
             sequences: self.sequences.max(child.sequences),
         }
     }
+} // ----------------------------------------------------------------------------
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+pub enum ColorMap {
+    Grayscale,
+    Turbo,
+    Viridis,
+    Plasma,
+    Magma,
+    Inferno,
+}
+
+impl std::fmt::Display for ColorMap {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(match self {
+            ColorMap::Grayscale => "Grayscale",
+            ColorMap::Turbo => "Turbo",
+            ColorMap::Viridis => "Viridis",
+            ColorMap::Plasma => "Plasma",
+            ColorMap::Magma => "Magma",
+            ColorMap::Inferno => "Inferno",
+        })
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+pub enum ColorMapper {
+    /// Use a well-known color map, pre-implemented as a wgsl module.
+    ColorMap(ColorMap),
+    // TODO(cmc): support textures.
+    // TODO(cmc): support custom transfer functions.
+}
+
+impl std::fmt::Display for ColorMapper {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ColorMapper::ColorMap(colormap) => colormap.fmt(f),
+        }
+    }
+}
+
+impl Default for ColorMapper {
+    fn default() -> Self {
+        Self::ColorMap(ColorMap::Grayscale)
+    }
 }
 
 // ----------------------------------------------------------------------------

crates/re_renderer/examples/depth_cloud.rs

@@ -178,6 +178,7 @@ impl RenderDepthClouds {
                 radius_scale: *radius_scale,
                 depth_dimensions: depth.dimensions,
                 depth_data: depth.data.clone(),
+                colormap: re_renderer::ColorMap::ColorMapTurbo,
             }],
         )
         .unwrap();

crates/re_renderer/shader/colormap.wgsl

@@ -11,15 +11,15 @@ const COLORMAP_INFERNO: u32 = 5u;
 
 fn colormap_srgb(which: u32, t: f32) -> Vec3 {
     if which == COLORMAP_TURBO {
-        return colormap_turbo(t);
+        return colormap_turbo_srgb(t);
     } else if which == COLORMAP_VIRIDIS {
-        return colormap_viridis(t);
+        return colormap_viridis_srgb(t);
     } else if which == COLORMAP_PLASMA {
-        return colormap_plasma(t);
+        return colormap_plasma_srgb(t);
     } else if which == COLORMAP_MAGMA {
-        return colormap_magma(t);
+        return colormap_magma_srgb(t);
     } else if which == COLORMAP_INFERNO {
-        return colormap_inferno(t);
+        return colormap_inferno_srgb(t);
     } else { // assume grayscale
         return linear_from_srgb(Vec3(t));
     }

crates/re_renderer/shader/depth_cloud.wgsl

@@ -2,6 +2,7 @@
 //!
 //! See `src/renderer/depth_cloud.rs` for more documentation.
 
+#import <./colormap.wgsl>
 #import <./global_bindings.wgsl>
 #import <./types.wgsl>
 #import <./utils/camera.wgsl>
@@ -26,8 +27,8 @@ fn compute_point_data(quad_idx: i32) -> PointData {
     // TODO(cmc): expose knobs to linearize/normalize/flip/cam-to-plane depth.
     let norm_linear_depth = textureLoad(depth_texture, texcoords, 0).x;
 
-    // TODO(cmc): support color maps & albedo textures
-    let color = Vec4(linear_from_srgb(Vec3(norm_linear_depth)), 1.0);
+    // TODO(cmc): albedo textures
+    let color = Vec4(colormap_srgb(depth_cloud_info.colormap, norm_linear_depth), 1.0);
 
     // TODO(cmc): This assumes a pinhole camera; need to support other kinds at some point.
     let intrinsics = depth_cloud_info.depth_camera_intrinsics;
@@ -62,6 +63,12 @@ struct DepthCloudInfo {
 
     /// The scale to apply to the radii of the backprojected points.
     radius_scale: f32,
+    radius_scale_row_pad0: f32,
+    radius_scale_row_pad1: f32,
+    radius_scale_row_pad2: f32,
+
+    /// Configures color mapping mode, see `colormap.wgsl`.
+    colormap: u32,
 };
 @group(1) @binding(0)
 var<uniform> depth_cloud_info: DepthCloudInfo;

crates/re_renderer/src/renderer/depth_cloud.rs

@@ -24,6 +24,7 @@ use crate::{
         GpuRenderPipelineHandle, GpuTexture, PipelineLayoutDesc, RenderPipelineDesc,
         ShaderModuleDesc, TextureDesc,
     },
+    ColorMap,
 };
 
 use super::{
@@ -42,8 +43,9 @@ mod gpu_data {
         pub depth_camera_extrinsics: crate::wgpu_buffer_types::Mat4,
         pub depth_camera_intrinsics: crate::wgpu_buffer_types::Mat3,
         pub radius_scale: crate::wgpu_buffer_types::F32RowPadded,
+        pub colormap: crate::wgpu_buffer_types::U32RowPadded,
 
-        pub end_padding: [crate::wgpu_buffer_types::PaddingRow; 16 - 8],
+        pub end_padding: [crate::wgpu_buffer_types::PaddingRow; 16 - 9],
     }
 }
 
@@ -89,6 +91,9 @@ pub struct DepthCloud {
     ///
     /// See [`DepthCloudDepthData`] for more information.
     pub depth_data: DepthCloudDepthData,
+
+    /// Configures color mapping mode.
+    pub colormap: ColorMap,
 }
 
 impl Default for DepthCloud {
@@ -99,6 +104,7 @@ impl Default for DepthCloud {
             radius_scale: 1.0,
             depth_dimensions: glam::UVec2::ZERO,
             depth_data: DepthCloudDepthData::default(),
+            colormap: ColorMap::ColorMapTurbo,
         }
     }
 }
@@ -144,6 +150,7 @@ impl DepthCloudDrawData {
                 depth_camera_extrinsics: info.depth_camera_extrinsics.into(),
                 depth_camera_intrinsics: info.depth_camera_intrinsics.into(),
                 radius_scale: info.radius_scale.into(),
+                colormap: (info.colormap as u32).into(),
                 end_padding: Default::default(),
             }),
         );

crates/re_viewer/src/ui/selection_panel.rs

@@ -1,5 +1,7 @@
 use egui::NumExt as _;
-use re_data_store::{query_latest_single, EditableAutoValue, EntityPath, EntityProperties};
+use re_data_store::{
+    query_latest_single, ColorMap, ColorMapper, EditableAutoValue, EntityPath, EntityProperties,
+};
 use re_log_types::{
     component_types::{Tensor, TensorDataMeaning},
     TimeType, Transform,
@@ -418,6 +420,45 @@ fn entity_props_ui(
         });
 }
 
+fn colormap_props_ui(ui: &mut egui::Ui, entity_props: &mut EntityProperties) {
+    ui.checkbox(&mut entity_props.color_mapping, "Color mapping")
+        .on_hover_text("Toggles color mapping");
+    ui.end_row();
+
+    if !entity_props.color_mapping {
+        return;
+    }
+
+    let current = *entity_props.color_mapper.get();
+
+    ui.label("Color map");
+    egui::ComboBox::from_id_source("color_mapper")
+        .selected_text(current.to_string())
+        .show_ui(ui, |ui| {
+            ui.style_mut().wrap = Some(false);
+            ui.set_min_width(64.0);
+
+            // TODO(cmc): that is not ideal but I don't want to import yet another proc-macro...
+            let mut add_label = |proposed| {
+                if ui
+                    .selectable_label(current == proposed, proposed.to_string())
+                    .clicked()
+                {
+                    entity_props.color_mapper = EditableAutoValue::Auto(proposed);
+                }
+            };
+
+            add_label(ColorMapper::ColorMap(ColorMap::Grayscale));
+            add_label(ColorMapper::ColorMap(ColorMap::Turbo));
+            add_label(ColorMapper::ColorMap(ColorMap::Viridis));
+            add_label(ColorMapper::ColorMap(ColorMap::Plasma));
+            add_label(ColorMapper::ColorMap(ColorMap::Magma));
+            add_label(ColorMapper::ColorMap(ColorMap::Inferno));
+        });
+
+    ui.end_row();
+}
+
 fn pinhole_props_ui(
     ctx: &mut ViewerContext<'_>,
     ui: &mut egui::Ui,
@@ -520,6 +561,10 @@ fn depth_props_ui(
                 entity_props.backproject_radius_scale = EditableAutoValue::UserEdited(radius_scale);
             }
             ui.end_row();
+
+            // TODO(cmc): This should apply to the depth map entity as a whole, but for that we
+            // need to get the current hardcoded colormapping out of the image cache first.
+            colormap_props_ui(ui, entity_props);
         } else {
             entity_props.backproject_pinhole_ent_path = None;
         }

crates/re_viewer/src/ui/view_spatial/scene/scene_part/images.rs

@@ -13,7 +13,7 @@ use re_log_types::{
 use re_query::{query_primary_with_history, EntityView, QueryError};
 use re_renderer::{
     renderer::{DepthCloud, DepthCloudDepthData},
-    Size,
+    ColorMap, Size,
 };
 
 use crate::{
@@ -307,12 +307,27 @@ impl ImagesPart {
 
         let world_from_obj = extrinsics * glam::Mat4::from_scale(glam::Vec3::splat(scale));
 
+        let colormap = properties
+            .color_mapping
+            .then(|| match *properties.color_mapper.get() {
+                re_data_store::ColorMapper::ColorMap(colormap) => match colormap {
+                    re_data_store::ColorMap::Grayscale => ColorMap::Grayscale,
+                    re_data_store::ColorMap::Turbo => ColorMap::ColorMapTurbo,
+                    re_data_store::ColorMap::Viridis => ColorMap::ColorMapViridis,
+                    re_data_store::ColorMap::Plasma => ColorMap::ColorMapPlasma,
+                    re_data_store::ColorMap::Magma => ColorMap::ColorMapMagma,
+                    re_data_store::ColorMap::Inferno => ColorMap::ColorMapInferno,
+                },
+            })
+            .unwrap_or(ColorMap::Grayscale);
+
         scene.primitives.depth_clouds.push(DepthCloud {
             depth_camera_extrinsics: world_from_obj,
             depth_camera_intrinsics: intrinsics.image_from_cam.into(),
             radius_scale,
             depth_dimensions: dimensions,
             depth_data: data,
+            colormap,
         });
     }
 }