Fix grayscale images being too dark (#3999)

### What
* Closes #3994

We previously would apply a grayscale colormap to a grayscale image -
this was not only a bit redundant; it actually caused the wrong colors.
Now we only apply the grayscale color map for things like depth images
and tensors (when the user has selected such a colormap) and in all
other cases we preserve the raw grayscale values as best we can.

We tested:
* grayscale images of different bit depth and ranges (see below)
* tensor selection colormap preview
* tensor with grayscale colormap
* depth point cloud selection colormap preview
* depth point cloud with grayscale colormap
* NV12 test


![image](https://github.com/rerun-io/rerun/assets/1148717/94886ddb-61e2-4356-96b7-191015256b02)

### Checklist
* [x] I have read and agree to [Contributor
Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and
the [Code of
Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md)
* [x] I've included a screenshot or gif (if applicable)
* [x] I have tested [demo.rerun.io](https://demo.rerun.io/pr/3999) (if
applicable)
* [x] The PR title and labels are set such as to maximize their
usefulness for the next release's CHANGELOG

- [PR Build Summary](https://build.rerun.io/pr/3999)
- [Docs
preview](https://rerun.io/preview/4edf5f6b5d8ccf6ec33c359f85347ac7b0d40866/docs)
<!--DOCS-PREVIEW-->
- [Examples
preview](https://rerun.io/preview/4edf5f6b5d8ccf6ec33c359f85347ac7b0d40866/examples)
<!--EXAMPLES-PREVIEW-->
- [Recent benchmark results](https://ref.rerun.io/dev/bench/)
- [Wasm size tracking](https://ref.rerun.io/dev/sizes/)

crates/re_renderer/shader/rectangle.wgsl

@@ -6,14 +6,15 @@
 const SAMPLE_TYPE_FLOAT = 1u;
 const SAMPLE_TYPE_SINT  = 2u;
 const SAMPLE_TYPE_UINT  = 3u;
-const SAMPLE_TYPE_NV12 = 4u;
+const SAMPLE_TYPE_NV12  = 4u;
 
 // How do we do colormapping?
-const COLOR_MAPPER_OFF      = 1u;
-const COLOR_MAPPER_FUNCTION = 2u;
-const COLOR_MAPPER_TEXTURE  = 3u;
+const COLOR_MAPPER_OFF_GRAYSCALE = 1u;
+const COLOR_MAPPER_OFF_RGB       = 2u;
+const COLOR_MAPPER_FUNCTION      = 3u;
+const COLOR_MAPPER_TEXTURE       = 4u;
 
-const FILTER_NEAREST = 1u;
+const FILTER_NEAREST  = 1u;
 const FILTER_BILINEAR = 2u;
 
 struct UniformBuffer {

crates/re_renderer/shader/rectangle_fs.wgsl

@@ -145,7 +145,9 @@ fn fs_main(in: VertexOut) -> @location(0) Vec4 {
 
     // Apply colormap, if any:
     var texture_color: Vec4;
-    if rect_info.color_mapper == COLOR_MAPPER_OFF {
+    if rect_info.color_mapper == COLOR_MAPPER_OFF_GRAYSCALE {
+        texture_color = Vec4(normalized_value.rrr, 1.0);
+    } else if rect_info.color_mapper == COLOR_MAPPER_OFF_RGB {
         texture_color = normalized_value;
     } else if rect_info.color_mapper == COLOR_MAPPER_FUNCTION {
         let rgb = colormap_linear(rect_info.colormap_function, normalized_value.r);

crates/re_renderer/src/renderer/rectangles.rs

@@ -88,15 +88,21 @@ pub struct ColormappedTexture {
     ///
     /// Setting a color mapper for a four-component texture is an error.
     /// Failure to set a color mapper for a one-component texture is an error.
-    pub color_mapper: Option<ColorMapper>,
+    pub color_mapper: ColorMapper,
 
     /// For textures that need decoding in the shader, for example NV12 encoded images.
     pub shader_decoding: Option<ShaderDecoding>,
 }
 
 /// How to map the normalized `.r` component to a color.
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub enum ColorMapper {
+    /// Colormapping is off. Take the .r color and splat onto rgb.
+    OffGrayscale,
+
+    /// Colormapping is off. Keep rgb as is.
+    OffRGB,
+
     /// Apply the given function.
     Function(Colormap),
 
@@ -110,6 +116,16 @@ pub enum ColorMapper {
     Texture(GpuTexture2D),
 }
 
+impl ColorMapper {
+    #[inline]
+    pub fn is_on(&self) -> bool {
+        match self {
+            ColorMapper::OffGrayscale | ColorMapper::OffRGB => false,
+            ColorMapper::Function(_) | ColorMapper::Texture(_) => true,
+        }
+    }
+}
+
 impl ColormappedTexture {
     /// Assumes a separate/unmultiplied alpha.
     pub fn from_unorm_rgba(texture: GpuTexture2D) -> Self {
@@ -121,7 +137,7 @@ impl ColormappedTexture {
             range: [0.0, 1.0],
             gamma: 1.0,
             multiply_rgb_with_alpha: true,
-            color_mapper: None,
+            color_mapper: ColorMapper::OffRGB,
             shader_decoding: None,
         }
     }
@@ -221,9 +237,10 @@ mod gpu_data {
     const SAMPLE_TYPE_NV12: u32 = 4;
 
     // How do we do colormapping?
-    const COLOR_MAPPER_OFF: u32 = 1;
-    const COLOR_MAPPER_FUNCTION: u32 = 2;
-    const COLOR_MAPPER_TEXTURE: u32 = 3;
+    const COLOR_MAPPER_OFF_GRAYSCALE: u32 = 1;
+    const COLOR_MAPPER_OFF_RGB: u32 = 2;
+    const COLOR_MAPPER_FUNCTION: u32 = 3;
+    const COLOR_MAPPER_TEXTURE: u32 = 4;
 
     const FILTER_NEAREST: u32 = 1;
     const FILTER_BILINEAR: u32 = 2;
@@ -309,33 +326,27 @@ mod gpu_data {
             };
 
             let mut colormap_function = 0;
-            let color_mapper_int;
-
-            match texture_format.components() {
+            let color_mapper_int = match texture_format.components() {
                 1 => match color_mapper {
-                    Some(ColorMapper::Function(colormap)) => {
-                        color_mapper_int = COLOR_MAPPER_FUNCTION;
+                    ColorMapper::OffGrayscale => COLOR_MAPPER_OFF_GRAYSCALE,
+                    ColorMapper::OffRGB => COLOR_MAPPER_OFF_RGB,
+                    ColorMapper::Function(colormap) => {
                         colormap_function = *colormap as u32;
+                        COLOR_MAPPER_FUNCTION
                     }
-                    Some(ColorMapper::Texture(_)) => {
-                        color_mapper_int = COLOR_MAPPER_TEXTURE;
-                    }
-                    None => match shader_decoding {
-                        Some(super::ShaderDecoding::Nv12) => color_mapper_int = COLOR_MAPPER_OFF,
-                        _ => return Err(RectangleError::MissingColorMapper),
-                    },
+                    ColorMapper::Texture(_) => COLOR_MAPPER_TEXTURE,
                 },
-                4 => {
-                    if color_mapper.is_some() {
+                4 => match color_mapper {
+                    ColorMapper::OffGrayscale => COLOR_MAPPER_OFF_GRAYSCALE, // This is a bit weird, but why not
+                    ColorMapper::OffRGB => COLOR_MAPPER_OFF_RGB,
+                    ColorMapper::Function(_) | ColorMapper::Texture(_) => {
                         return Err(RectangleError::ColormappingRgbaTexture);
-                    } else {
-                        color_mapper_int = COLOR_MAPPER_OFF;
                     }
-                }
+                },
                 num_components => {
                     return Err(RectangleError::UnsupportedComponentCount(num_components));
                 }
-            }
+            };
 
             let minification_filter = match rectangle.options.texture_filter_minification {
                 super::TextureFilterMin::Linear => FILTER_BILINEAR,
@@ -448,17 +459,16 @@ impl RectangleDrawData {
             }
 
             // We also set up an optional colormap texture.
-            let colormap_texture = if let Some(ColorMapper::Texture(handle)) =
-                &rectangle.colormapped_texture.color_mapper
-            {
-                let format = handle.format();
-                if format != wgpu::TextureFormat::Rgba8UnormSrgb {
-                    return Err(RectangleError::UnsupportedColormapTextureFormat(format));
-                }
-                handle.handle()
-            } else {
-                ctx.texture_manager_2d.zeroed_texture_float().handle
-            };
+            let colormap_texture =
+                if let ColorMapper::Texture(handle) = &rectangle.colormapped_texture.color_mapper {
+                    let format = handle.format();
+                    if format != wgpu::TextureFormat::Rgba8UnormSrgb {
+                        return Err(RectangleError::UnsupportedColormapTextureFormat(format));
+                    }
+                    handle.handle()
+                } else {
+                    ctx.texture_manager_2d.zeroed_texture_float().handle
+                };
 
             instances.push(RectangleInstance {
                 bind_group: ctx.gpu_resources.bind_groups.alloc(

crates/re_renderer/src/resource_managers/texture_manager.rs

@@ -16,6 +16,12 @@ use crate::{
 #[derive(Clone)]
 pub struct GpuTexture2D(GpuTexture);
 
+impl std::fmt::Debug for GpuTexture2D {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("GpuTexture2D").field(&self.0.handle).finish()
+    }
+}
+
 impl GpuTexture2D {
     #[inline]
     pub fn handle(&self) -> crate::wgpu_resources::GpuTextureHandle {

crates/re_space_view_spatial/src/parts/images.rs

@@ -96,7 +96,7 @@ fn to_textured_rect(
             // so it's usually safer to turn off.
             // The best heuristic is this: if there is a color map being applied, use nearest.
             // TODO(emilk): apply filtering _after_ the color map?
-            let texture_filter_minification = if colormapped_texture.color_mapper.is_some() {
+            let texture_filter_minification = if colormapped_texture.color_mapper.is_on() {
                 re_renderer::renderer::TextureFilterMin::Nearest
             } else {
                 re_renderer::renderer::TextureFilterMin::Linear

crates/re_space_view_tensor/src/tensor_slice_to_gpu.rs

@@ -48,9 +48,7 @@ pub fn colormapped_texture(
         decode_srgb: false,
         multiply_rgb_with_alpha: false,
         gamma: color_mapping.gamma,
-        color_mapper: Some(re_renderer::renderer::ColorMapper::Function(
-            color_mapping.map,
-        )),
+        color_mapper: re_renderer::renderer::ColorMapper::Function(color_mapping.map),
         shader_decoding: match &tensor.buffer {
             &TensorBuffer::Nv12(_) => Some(ShaderDecoding::Nv12),
             _ => None,

crates/re_viewer_context/src/gpu_bridge/colormap.rs

@@ -43,7 +43,7 @@ fn colormap_preview_ui(
         multiply_rgb_with_alpha: false,
         gamma: 1.0,
         shader_decoding: None,
-        color_mapper: Some(re_renderer::renderer::ColorMapper::Function(colormap)),
+        color_mapper: re_renderer::renderer::ColorMapper::Function(colormap),
     };
 
     let debug_name = format!("colormap_{colormap}");

crates/re_viewer_context/src/gpu_bridge/tensor_to_gpu.rs

@@ -86,6 +86,7 @@ pub fn color_tensor_to_gpu(
     tensor_stats: &TensorStats,
 ) -> anyhow::Result<ColormappedTexture> {
     re_tracing::profile_function!();
+
     let texture_key = hash(tensor_data_row_id);
     let [height, width, depth] = texture_height_width_channels(tensor)?;
 
@@ -148,18 +149,36 @@ pub fn color_tensor_to_gpu(
         super::tensor_data_range_heuristic(tensor_stats, tensor.dtype())?
     };
 
-    let color_mapper = if shader_decoding.is_none() && texture_format.components() == 1 {
-        // Single-channel images = luminance = grayscale
-        Some(ColorMapper::Function(re_renderer::Colormap::Grayscale))
-    } else {
-        None
+    let color_mapper = match shader_decoding {
+        None => {
+            if texture_format.components() == 1 {
+                if decode_srgb {
+                    // Leave grayscale images unmolested - don't apply a colormap to them.
+                    ColorMapper::OffGrayscale
+                } else {
+                    // This is something like a uint16 image, or a float image
+                    // with a range outside of 0-255 (see tensor_decode_srgb_gamma_heuristic).
+                    // `tensor_data_range_heuristic` will make sure we map this to a 0-1
+                    // range, and then we apply a gray colormap to it.
+                    ColorMapper::Function(re_renderer::Colormap::Grayscale)
+                }
+            } else {
+                ColorMapper::OffRGB
+            }
+        }
+
+        Some(ShaderDecoding::Nv12) => ColorMapper::OffRGB,
     };
 
     // TODO(wumpf): There should be a way to specify whether a texture uses pre-multiplied alpha or not.
     // Assume that the texture is not pre-multiplied if it has an alpha channel.
     let multiply_rgb_with_alpha = depth == 4;
     let gamma = 1.0;
 
+    re_log::trace_once!(
+        "color_tensor_to_gpu {debug_name:?}, range: {range:?}, decode_srgb: {decode_srgb:?}, multiply_rgb_with_alpha: {multiply_rgb_with_alpha:?}, gamma: {gamma:?}, color_mapper: {color_mapper:?}, shader_decoding: {shader_decoding:?}",
+    );
+
     Ok(ColormappedTexture {
         texture: texture_handle,
         range,
@@ -242,7 +261,7 @@ pub fn class_id_tensor_to_gpu(
         decode_srgb: false, // Setting this to true would affect the class ids, not the color they resolve to.
         multiply_rgb_with_alpha: false, // already premultiplied!
         gamma: 1.0,
-        color_mapper: Some(ColorMapper::Texture(colormap_texture_handle)),
+        color_mapper: ColorMapper::Texture(colormap_texture_handle),
         shader_decoding: None,
     })
 }
@@ -279,7 +298,7 @@ pub fn depth_tensor_to_gpu(
         decode_srgb: false,
         multiply_rgb_with_alpha: false,
         gamma: 1.0,
-        color_mapper: Some(ColorMapper::Function(re_renderer::Colormap::Turbo)),
+        color_mapper: ColorMapper::Function(re_renderer::Colormap::Turbo),
         shader_decoding: None,
     })
 }

tests/python/nv12image/main.py

@@ -8,7 +8,7 @@
 
 import cv2
 import numpy as np
-import rerun
+import rerun as rr
 
 
 def bgra2nv12(bgra: Any) -> np.ndarray:
@@ -21,28 +21,30 @@ def bgra2nv12(bgra: Any) -> np.ndarray:
 
 def main() -> None:
     parser = argparse.ArgumentParser(description="Displaying NV12 encoded images.")
-    rerun.script_add_args(parser)
+    rr.script_add_args(parser)
     args = parser.parse_args()
 
-    rerun.script_setup(args, "rerun_test_nv12image")
+    rr.script_setup(args, "rerun_test_nv12image")
 
     # Make sure you use a colorful image!
     dir_path = os.path.dirname(os.path.realpath(__file__))
     img_path = f"{dir_path}/../../../crates/re_ui/data/logo_dark_mode.png"
     img_bgra = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
 
     img_rgb = cv2.cvtColor(img_bgra, cv2.COLOR_BGRA2RGB)
-    rerun.log("img_reference", rerun.Image(img_rgb))
+    rr.log("img_reference", rr.Image(img_rgb))
 
-    rerun.log(
+    rr.log(
         "img_nv12",
-        rerun.ImageEncoded(
+        rr.ImageEncoded(
             contents=bytes(bgra2nv12(img_bgra)),
-            format=rerun.ImageFormat.NV12((img_bgra.shape[0], img_bgra.shape[1])),
+            format=rr.ImageFormat.NV12((img_bgra.shape[0], img_bgra.shape[1])),
         ),
     )
 
-    rerun.script_teardown(args)
+    rr.log("expectation", rr.TextDocument("The images should look the same, except for some chroma artifacts."))
+
+    rr.script_teardown(args)
 
 
 if __name__ == "__main__":

tests/python/test_api/main.py

@@ -415,6 +415,25 @@ def run_extension_component(experimental_api: bool) -> None:
     )
 
 
+def run_gradient_image(experimental_api: bool) -> None:
+    rr.log(
+        "gradient_explain",
+        rr.TextLog("gradients should look the same, and external color picket should show 128 in the middle"),
+    )
+
+    x, _ = np.meshgrid(np.arange(0, 256), np.arange(0, 64))
+    im = x.astype(np.uint8)
+    rr.log("gradient_u8", rr.Image(im))
+
+    x, _ = np.meshgrid(np.arange(0, 255), np.arange(0, 64))
+    im = (x / 255.0).astype(np.float32)
+    rr.log("gradient_f32", rr.Image(im))
+
+    x, _ = np.meshgrid(np.arange(0, 256), np.arange(0, 64))
+    im = (x * 4).astype(np.uint16)
+    rr.log("gradient_u16_0_1020", rr.Image(im))
+
+
 def run_image_tensors(experimental_api: bool) -> None:
     # Make sure you use a colorful image with alpha!
     dir_path = os.path.dirname(os.path.realpath(__file__))
@@ -460,6 +479,7 @@ def main() -> None:
         "3d_points": run_3d_points,
         "bbox": run_bounding_box,
         "extension_components": run_extension_component,
+        "gradient_image": run_gradient_image,
         "image_tensors": run_image_tensors,
         "log_cleared": run_log_cleared,
         "raw_mesh": raw_mesh,