GPU colormapping, first step (#1835)

* Add TextureManager2D::get_or_create_with

* Small code cleanup

* Add code to upload a Tensor to a GPU texture

* Add helper method Tensor::image_height_width_depth

* Minor code cleanup (multiplicative_tint)

* Hook up color textures via the new path

* Refactor: introduce ColormappedTexture

* Start working on an uint sampler

* merge fix

* Dumb colormapping of depth textures!

* Use turbo for depth maps (and single-channel images :grimace:)

* Use grayscale for luminance

* ColorMap -> Colormap

* Apply annotation context colormaps

* Support sint textures too

* cleanup

* merge fix

* Fix RGB images

* More cleanup

* Better error-handlign and nicer error message

* Clean up the SAMPLE_TYPE with constants

* Nicer shader interface

* Better error handling

* Remove dead code

* Self-review cleanup

* Fix bug in shader when sampling sint textures

* Use textureSampleLevel

* Apply a gamma to the image (unused as of now)

* image_height_width_channels

* fix various review comments

* Optimize narrow_f64_to_f32s: avoid one allocation

crates/re_data_store/src/entity_properties.rs

@@ -140,7 +140,8 @@ impl ExtraQueryHistory {
 
 #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
 #[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
-pub enum ColorMap {
+pub enum Colormap {
+    /// Perceptually even
     Grayscale,
     #[default]
     Turbo,
@@ -150,15 +151,15 @@ pub enum ColorMap {
     Inferno,
 }
 
-impl std::fmt::Display for ColorMap {
+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",
+            Colormap::Grayscale => "Grayscale",
+            Colormap::Turbo => "Turbo",
+            Colormap::Viridis => "Viridis",
+            Colormap::Plasma => "Plasma",
+            Colormap::Magma => "Magma",
+            Colormap::Inferno => "Inferno",
         })
     }
 }
@@ -167,23 +168,23 @@ impl std::fmt::Display for ColorMap {
 #[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),
+    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),
+            ColorMapper::Colormap(colormap) => colormap.fmt(f),
         }
     }
 }
 
 impl Default for ColorMapper {
     #[inline]
     fn default() -> Self {
-        Self::ColorMap(ColorMap::default())
+        Self::Colormap(Colormap::default())
     }
 }
 

crates/re_log_types/src/component_types/tensor.rs

@@ -369,6 +369,23 @@ impl Tensor {
         self.shape.len()
     }
 
+    /// If this tensor is shaped as an image, return the height, width, and channels/depth of it.
+    pub fn image_height_width_channels(&self) -> Option<[u64; 3]> {
+        if self.shape.len() == 2 {
+            Some([self.shape[0].size, self.shape[1].size, 1])
+        } else if self.shape.len() == 3 {
+            let channels = self.shape[2].size;
+            // gray, rgb, rgba
+            if matches!(channels, 1 | 3 | 4) {
+                Some([self.shape[0].size, self.shape[1].size, channels])
+            } else {
+                None
+            }
+        } else {
+            None
+        }
+    }
+
     pub fn is_shaped_like_an_image(&self) -> bool {
         self.num_dim() == 2
             || self.num_dim() == 3 && {

crates/re_renderer/examples/2d.rs

@@ -1,7 +1,8 @@
 use ecolor::Hsva;
 use re_renderer::{
     renderer::{
-        LineStripFlags, RectangleDrawData, TextureFilterMag, TextureFilterMin, TexturedRect,
+        ColormappedTexture, LineStripFlags, RectangleDrawData, TextureFilterMag, TextureFilterMin,
+        TexturedRect,
     },
     resource_managers::{GpuTexture2DHandle, Texture2DCreationDesc},
     view_builder::{self, Projection, TargetConfiguration, ViewBuilder},
@@ -39,7 +40,7 @@ impl framework::Example for Render2D {
             &mut re_ctx.gpu_resources.textures,
             &Texture2DCreationDesc {
                 label: "rerun logo".into(),
-                data: &image_data,
+                data: image_data.into(),
                 format: wgpu::TextureFormat::Rgba8UnormSrgb,
                 width: rerun_logo.width(),
                 height: rerun_logo.height(),
@@ -193,7 +194,9 @@ impl framework::Example for Render2D {
                     top_left_corner_position: glam::vec3(500.0, 120.0, -0.05),
                     extent_u: self.rerun_logo_texture_width as f32 * image_scale * glam::Vec3::X,
                     extent_v: self.rerun_logo_texture_height as f32 * image_scale * glam::Vec3::Y,
-                    texture: self.rerun_logo_texture.clone(),
+                    colormapped_texture: ColormappedTexture::from_unorm_srgba(
+                        self.rerun_logo_texture.clone(),
+                    ),
                     texture_filter_magnification: TextureFilterMag::Nearest,
                     texture_filter_minification: TextureFilterMin::Linear,
                     ..Default::default()
@@ -207,7 +210,9 @@ impl framework::Example for Render2D {
                     ),
                     extent_u: self.rerun_logo_texture_width as f32 * image_scale * glam::Vec3::X,
                     extent_v: self.rerun_logo_texture_height as f32 * image_scale * glam::Vec3::Y,
-                    texture: self.rerun_logo_texture.clone(),
+                    colormapped_texture: ColormappedTexture::from_unorm_srgba(
+                        self.rerun_logo_texture.clone(),
+                    ),
                     texture_filter_magnification: TextureFilterMag::Linear,
                     texture_filter_minification: TextureFilterMin::Linear,
                     depth_offset: 1,

crates/re_renderer/examples/depth_cloud.rs

@@ -20,8 +20,8 @@ use itertools::Itertools;
 use macaw::IsoTransform;
 use re_renderer::{
     renderer::{
-        DepthCloud, DepthCloudDepthData, DepthCloudDrawData, DepthClouds, DrawData,
-        GenericSkyboxDrawData, RectangleDrawData, TexturedRect,
+        ColormappedTexture, DepthCloud, DepthCloudDepthData, DepthCloudDrawData, DepthClouds,
+        DrawData, GenericSkyboxDrawData, RectangleDrawData, TexturedRect,
     },
     resource_managers::{GpuTexture2DHandle, Texture2DCreationDesc},
     view_builder::{self, Projection, ViewBuilder},
@@ -183,7 +183,7 @@ impl RenderDepthClouds {
                     max_depth_in_world: 5.0,
                     depth_dimensions: depth.dimensions,
                     depth_data: depth.data.clone(),
-                    colormap: re_renderer::ColorMap::ColorMapTurbo,
+                    colormap: re_renderer::Colormap::Turbo,
                     outline_mask_id: Default::default(),
                 }],
                 radius_boost_in_ui_points_for_outlines: 2.5,
@@ -245,7 +245,7 @@ impl framework::Example for RenderDepthClouds {
             &mut re_ctx.gpu_resources.textures,
             &Texture2DCreationDesc {
                 label: "albedo".into(),
-                data: bytemuck::cast_slice(&albedo.rgba8),
+                data: bytemuck::cast_slice(&albedo.rgba8).into(),
                 format: wgpu::TextureFormat::Rgba8UnormSrgb,
                 width: albedo.dimensions.x,
                 height: albedo.dimensions.y,
@@ -331,7 +331,7 @@ impl framework::Example for RenderDepthClouds {
                     .transform_point3(glam::Vec3::new(1.0, 1.0, 0.0)),
                 extent_u: world_from_model.transform_vector3(-glam::Vec3::X),
                 extent_v: world_from_model.transform_vector3(-glam::Vec3::Y),
-                texture: albedo_handle.clone(),
+                colormapped_texture: ColormappedTexture::from_unorm_srgba(albedo_handle.clone()),
                 texture_filter_magnification: re_renderer::renderer::TextureFilterMag::Nearest,
                 texture_filter_minification: re_renderer::renderer::TextureFilterMin::Linear,
                 multiplicative_tint: Rgba::from_white_alpha(0.5),

crates/re_renderer/shader/colormap.wgsl

@@ -2,18 +2,20 @@
 #import <./utils/srgb.wgsl>
 
 // NOTE: Keep in sync with `colormap.rs`!
-const GRAYSCALE:        u32 = 0u;
-const COLORMAP_TURBO:   u32 = 1u;
-const COLORMAP_VIRIDIS: u32 = 2u;
-const COLORMAP_PLASMA:  u32 = 3u;
-const COLORMAP_MAGMA:   u32 = 4u;
-const COLORMAP_INFERNO: u32 = 5u;
+const COLORMAP_GRAYSCALE: u32 = 1u;
+const COLORMAP_TURBO:     u32 = 2u;
+const COLORMAP_VIRIDIS:   u32 = 3u;
+const COLORMAP_PLASMA:    u32 = 4u;
+const COLORMAP_MAGMA:     u32 = 5u;
+const COLORMAP_INFERNO:   u32 = 6u;
 
 /// Returns a gamma-space sRGB in 0-1 range.
 ///
 /// The input will be saturated to [0, 1] range.
 fn colormap_srgb(which: u32, t: f32) -> Vec3 {
-    if which == COLORMAP_TURBO {
+    if which == COLORMAP_GRAYSCALE {
+        return linear_from_srgb(Vec3(t));
+    } else if which == COLORMAP_TURBO {
         return colormap_turbo_srgb(t);
     } else if which == COLORMAP_VIRIDIS {
         return colormap_viridis_srgb(t);
@@ -23,8 +25,8 @@ fn colormap_srgb(which: u32, t: f32) -> Vec3 {
         return colormap_magma_srgb(t);
     } else if which == COLORMAP_INFERNO {
         return colormap_inferno_srgb(t);
-    } else { // assume grayscale
-        return linear_from_srgb(Vec3(t));
+    } else {
+        return ERROR_RGBA.rgb;
     }
 }
 

crates/re_renderer/shader/rectangle.wgsl

@@ -1,14 +1,33 @@
 #import <./types.wgsl>
+#import <./colormap.wgsl>
 #import <./global_bindings.wgsl>
 #import <./utils/depth_offset.wgsl>
 
+// Keep in sync with mirror in rectangle.rs
+
+// Which texture to read from?
+const SAMPLE_TYPE_FLOAT = 1u;
+const SAMPLE_TYPE_SINT  = 2u;
+const SAMPLE_TYPE_UINT  = 3u;
+
+// How do we do colormapping?
+const COLOR_MAPPER_OFF      = 1u;
+const COLOR_MAPPER_FUNCTION = 2u;
+const COLOR_MAPPER_TEXTURE  = 3u;
+
 struct UniformBuffer {
     /// Top left corner position in world space.
     top_left_corner_position: Vec3,
 
+    /// Which colormap to use, if any
+    colormap_function: u32,
+
     /// Vector that spans up the rectangle from its top left corner along the u axis of the texture.
     extent_u: Vec3,
 
+    /// Which texture sample to use
+    sample_type: u32,
+
     /// Vector that spans up the rectangle from its top left corner along the v axis of the texture.
     extent_v: Vec3,
 
@@ -18,16 +37,35 @@ struct UniformBuffer {
     multiplicative_tint: Vec4,
 
     outline_mask: UVec2,
+
+    /// Range of the texture values.
+    /// Will be mapped to the [0, 1] range before we colormap.
+    range_min_max: Vec2,
+
+    color_mapper: u32,
+
+    /// Exponent to raise the normalized texture value.
+    /// Inverse brightness.
+    gamma: f32,
 };
 
 @group(1) @binding(0)
 var<uniform> rect_info: UniformBuffer;
 
 @group(1) @binding(1)
-var texture: texture_2d<f32>;
+var texture_sampler: sampler;
 
 @group(1) @binding(2)
-var texture_sampler: sampler;
+var texture_float: texture_2d<f32>;
+
+@group(1) @binding(3)
+var texture_sint: texture_2d<i32>;
+
+@group(1) @binding(4)
+var texture_uint: texture_2d<u32>;
+
+@group(1) @binding(5)
+var colormap_texture: texture_2d<f32>;
 
 
 struct VertexOut {
@@ -50,7 +88,46 @@ fn vs_main(@builtin(vertex_index) v_idx: u32) -> VertexOut {
 
 @fragment
 fn fs_main(in: VertexOut) -> @location(0) Vec4 {
-    let texture_color = textureSample(texture, texture_sampler, in.texcoord);
+    // Sample the main texture:
+    var sampled_value: Vec4;
+    if rect_info.sample_type == SAMPLE_TYPE_FLOAT {
+        sampled_value = textureSampleLevel(texture_float, texture_sampler, in.texcoord, 0.0); // TODO(emilk): support mipmaps
+    } else if rect_info.sample_type == SAMPLE_TYPE_SINT {
+        let icoords = IVec2(in.texcoord * Vec2(textureDimensions(texture_sint).xy));
+        sampled_value = Vec4(textureLoad(texture_sint, icoords, 0));
+    } else if rect_info.sample_type == SAMPLE_TYPE_UINT {
+        let icoords = IVec2(in.texcoord * Vec2(textureDimensions(texture_uint).xy));
+        sampled_value = Vec4(textureLoad(texture_uint, icoords, 0));
+    } else {
+        return ERROR_RGBA; // unknown sample type
+    }
+
+    // Normalize the sample:
+    let range = rect_info.range_min_max;
+    var normalized_value: Vec4 = (sampled_value - range.x) / (range.y - range.x);
+
+    // Apply gamma:
+    normalized_value = vec4(pow(normalized_value.rgb, vec3(rect_info.gamma)), normalized_value.a); // TODO(emilk): handle premultiplied alpha
+
+    // Apply colormap, if any:
+    var texture_color: Vec4;
+    if rect_info.color_mapper == COLOR_MAPPER_OFF {
+        texture_color = normalized_value;
+    } else if rect_info.color_mapper == COLOR_MAPPER_FUNCTION {
+        let rgb = colormap_linear(rect_info.colormap_function, normalized_value.r);
+        texture_color = Vec4(rgb, 1.0);
+    } else if rect_info.color_mapper == COLOR_MAPPER_TEXTURE {
+        let colormap_size = textureDimensions(colormap_texture).xy;
+        let color_index = normalized_value.r * f32(colormap_size.x * colormap_size.y);
+        // TODO(emilk): interpolate between neighboring colors for non-integral color indices
+        let color_index_i32 = i32(color_index);
+        let x = color_index_i32 % colormap_size.x;
+        let y = color_index_i32 / colormap_size.x;
+        texture_color = textureLoad(colormap_texture, IVec2(x, y), 0);
+    } else {
+        return ERROR_RGBA; // unknown color mapper
+    }
+
     return texture_color * rect_info.multiplicative_tint;
 }
 

crates/re_renderer/shader/types.wgsl

@@ -48,3 +48,7 @@ const ONE  = Vec4(1.0, 1.0, 1.0, 1.0);
 // fn inf() -> f32 {
 //     return 1.0 / 0.0;
 // }
+
+
+/// The color to use when we encounter an error.
+const ERROR_RGBA = Vec4(1.0, 0.0, 1.0, 1.0);

crates/re_renderer/src/colormap.rs

@@ -5,25 +5,28 @@ use glam::{Vec2, Vec3A, Vec4, Vec4Swizzles};
 // ---
 
 // NOTE: Keep in sync with `colormap.wgsl`!
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
 #[repr(u32)]
-pub enum ColorMap {
-    Grayscale = 0,
-    ColorMapTurbo = 1,
-    ColorMapViridis = 2,
-    ColorMapPlasma = 3,
-    ColorMapMagma = 4,
-    ColorMapInferno = 5,
+pub enum Colormap {
+    // Reserve 0 for "disabled"
+    /// Perceptually even
+    #[default]
+    Grayscale = 1,
+    Turbo = 2,
+    Viridis = 3,
+    Plasma = 4,
+    Magma = 5,
+    Inferno = 6,
 }
 
-pub fn colormap_srgb(which: ColorMap, t: f32) -> [u8; 4] {
+pub fn colormap_srgb(which: Colormap, t: f32) -> [u8; 4] {
     match which {
-        ColorMap::Grayscale => grayscale_srgb(t),
-        ColorMap::ColorMapTurbo => colormap_turbo_srgb(t),
-        ColorMap::ColorMapViridis => colormap_viridis_srgb(t),
-        ColorMap::ColorMapPlasma => colormap_plasma_srgb(t),
-        ColorMap::ColorMapMagma => colormap_magma_srgb(t),
-        ColorMap::ColorMapInferno => colormap_inferno_srgb(t),
+        Colormap::Grayscale => grayscale_srgb(t),
+        Colormap::Turbo => colormap_turbo_srgb(t),
+        Colormap::Viridis => colormap_viridis_srgb(t),
+        Colormap::Plasma => colormap_plasma_srgb(t),
+        Colormap::Magma => colormap_magma_srgb(t),
+        Colormap::Inferno => colormap_inferno_srgb(t),
     }
 }
 

crates/re_renderer/src/importer/gltf.rs

@@ -62,7 +62,7 @@ pub fn load_gltf_from_buffer(
                 format!("gltf image used by {texture_names} in {mesh_name}")
             }
             .into(),
-            data: &data,
+            data: data.into(),
             format,
             width: image.width,
             height: image.height,

crates/re_renderer/src/lib.rs

@@ -33,7 +33,7 @@ pub use allocator::GpuReadbackIdentifier;
 pub use color::Rgba32Unmul;
 pub use colormap::{
     colormap_inferno_srgb, colormap_magma_srgb, colormap_plasma_srgb, colormap_srgb,
-    colormap_turbo_srgb, colormap_viridis_srgb, grayscale_srgb, ColorMap,
+    colormap_turbo_srgb, colormap_viridis_srgb, grayscale_srgb, Colormap,
 };
 pub use context::RenderContext;
 pub use debug_label::DebugLabel;