Automatically determine image/tensor color mapping & need for srgb decoding (#2342)

<!--
Open the PR up as a draft until you feel it is ready for a proper
review.

Do not make PR:s from your own `main` branch, as that makes it difficult
for reviewers to add their own fixes.

Add any improvements to the branch as new commits to make it easier for
reviewers to follow the progress. All commits will be squashed to a
single commit once the PR is merged into `main`.

Make sure you mention any issues that this PR closes in the description,
as well as any other related issues.

To get an auto-generated PR description you can put "copilot:summary" or
"copilot:walkthrough" anywhere.
-->

Fixes #2274
* #2274

Will need to be further worked on via:
* #2341

Introduces 3 new things actually:
* finite range, determined if the range of a tensor was in fact not
finite
* heuristic for color mapping range
* heuristic for srgb
   * ⚠️ does this break things?

Range & sRGB:

This
```
import rerun as rr
import numpy as np
import matplotlib.pyplot as plt

rr.init("image color", spawn=True)
img = np.random.random((300, 300, 3)) * 0.2
rr.log_image("image", img)
print(img.min(), img.max())  # to check they are in range [0,1]

plt.imshow(img)
plt.show()
```
Gives now the expected image:

![image](https://github.com/rerun-io/rerun/assets/1220815/9fcda125-f762-4fd4-ae24-ce6dcaaa496a)

When setting a pixel to inf, we show the "finite range" as well, and
behave correctly:
<img width="1026" alt="image"
src="https://github.com/rerun-io/rerun/assets/1220815/7d3037e8-f562-4b6c-8f08-ec4f5627c21a">

* [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)

<!-- This line will get updated when the PR build summary job finishes.
-->
PR Build Summary: https://build.rerun.io/pr/2342

<!-- pr-link-docs:start -->
Docs preview: https://rerun.io/preview/93697e6/docs
Examples preview: https://rerun.io/preview/93697e6/examples
<!-- pr-link-docs:end -->

Author: Wumpf

crates/re_data_ui/src/image.rs

@@ -282,7 +282,10 @@ pub fn tensor_summary_ui_grid_contents(
         }
     }
 
-    let TensorStats { range } = tensor_stats;
+    let TensorStats {
+        range,
+        finite_range,
+    } = tensor_stats;
 
     if let Some((min, max)) = range {
         ui.label("Data range")
@@ -294,6 +297,17 @@ pub fn tensor_summary_ui_grid_contents(
         ));
         ui.end_row();
     }
+    // Show finite range only if it is different from the actual range.
+    if let (true, Some((min, max))) = (range != finite_range, finite_range) {
+        ui.label("Finite data range")
+            .on_hover_text("The finite values (ignoring all NaN & -Inf/+Inf) of the tensor range within these bounds.");
+        ui.monospace(format!(
+            "[{} - {}]",
+            re_format::format_f64(*min),
+            re_format::format_f64(*max)
+        ));
+        ui.end_row();
+    }
 }
 
 pub fn tensor_summary_ui(

crates/re_log_types/src/component_types/tensor.rs

@@ -566,7 +566,7 @@ impl Component for Tensor {
     }
 }
 
-#[derive(thiserror::Error, Debug, PartialEq)]
+#[derive(thiserror::Error, Debug, PartialEq, Clone)]
 pub enum TensorCastError {
     #[error("ndarray type mismatch with tensor storage")]
     TypeMismatch,

crates/re_renderer/shader/rectangle.wgsl

@@ -53,6 +53,9 @@ struct UniformBuffer {
 
     /// Boolean: decode 0-1 sRGB gamma to linear space before filtering?
     decode_srgb: u32,
+
+    /// Boolean: multiply RGB with alpha before filtering
+    multiply_rgb_with_alpha: u32,
 };
 
 @group(1) @binding(0)

crates/re_renderer/shader/rectangle_fs.wgsl

@@ -14,78 +14,85 @@ fn tex_filter(pixel_coord: Vec2) -> u32 {
     }
 }
 
-fn decode_color(rgba_arg: Vec4) -> Vec4 {
-    var rgba = rgba_arg;
+fn normalize_range(sampled_value: Vec4) -> Vec4 {
+    let range = rect_info.range_min_max;
+    return (sampled_value - range.x) / (range.y - range.x);
+}
+
+fn decode_color(sampled_value: Vec4) -> Vec4 {
+    // Normalize the value first, otherwise premultiplying alpha and linear space conversion won't make sense.
+    var rgba = normalize_range(sampled_value);
 
-    // Convert to linear space:
+    // Convert to linear space
     if rect_info.decode_srgb != 0u {
-        rgba = linear_from_srgba(rgba);
+        if all(0.0 <= rgba.rgb) && all(rgba.rgb <= 1.0) {
+            rgba = linear_from_srgba(rgba);
+        } else {
+            rgba = ERROR_RGBA; // out of range
+        }
     }
 
-    // Premultiply alpha:
-    rgba = vec4(rgba.xyz * rgba.a, rgba.a);
+    // Premultiply alpha.
+    if rect_info.multiply_rgb_with_alpha != 0u {
+        rgba = vec4(rgba.xyz * rgba.a, rgba.a);
+    }
 
     return rgba;
 }
 
+fn filter_bilinear(coord: Vec2, v00: Vec4, v01: Vec4, v10: Vec4, v11: Vec4) -> Vec4 {
+    let top = mix(v00, v10, fract(coord.x - 0.5));
+    let bottom = mix(v01, v11, fract(coord.x - 0.5));
+    return mix(top, bottom, fract(coord.y - 0.5));
+}
+
 @fragment
 fn fs_main(in: VertexOut) -> @location(0) Vec4 {
     // Sample the main texture:
-    var sampled_value: Vec4;
+    var normalized_value: Vec4;
     if rect_info.sample_type == SAMPLE_TYPE_FLOAT {
         let coord = in.texcoord * Vec2(textureDimensions(texture_float).xy);
         if tex_filter(coord) == FILTER_NEAREST {
             // nearest
-            sampled_value = decode_color(textureLoad(texture_float, IVec2(coord), 0));
+            normalized_value = decode_color(textureLoad(texture_float, IVec2(coord), 0));
         } else {
             // bilinear
             let v00 = decode_color(textureLoad(texture_float, IVec2(coord + vec2(-0.5, -0.5)), 0));
             let v01 = decode_color(textureLoad(texture_float, IVec2(coord + vec2(-0.5,  0.5)), 0));
             let v10 = decode_color(textureLoad(texture_float, IVec2(coord + vec2( 0.5, -0.5)), 0));
             let v11 = decode_color(textureLoad(texture_float, IVec2(coord + vec2( 0.5,  0.5)), 0));
-            let top = mix(v00, v10, fract(coord.x - 0.5));
-            let bottom = mix(v01, v11, fract(coord.x - 0.5));
-            sampled_value = mix(top, bottom, fract(coord.y - 0.5));
+            normalized_value = filter_bilinear(coord, v00, v01, v10, v11);
         }
     } else if rect_info.sample_type == SAMPLE_TYPE_SINT {
         let coord = in.texcoord * Vec2(textureDimensions(texture_sint).xy);
         if tex_filter(coord) == FILTER_NEAREST {
             // nearest
-            sampled_value = Vec4(textureLoad(texture_sint, IVec2(coord), 0));
+            normalized_value = decode_color(Vec4(textureLoad(texture_sint, IVec2(coord), 0)));
         } else {
             // bilinear
-            let v00 = Vec4(textureLoad(texture_sint, IVec2(coord + vec2(-0.5, -0.5)), 0));
-            let v01 = Vec4(textureLoad(texture_sint, IVec2(coord + vec2(-0.5,  0.5)), 0));
-            let v10 = Vec4(textureLoad(texture_sint, IVec2(coord + vec2( 0.5, -0.5)), 0));
-            let v11 = Vec4(textureLoad(texture_sint, IVec2(coord + vec2( 0.5,  0.5)), 0));
-            let top = mix(v00, v10, fract(coord.x - 0.5));
-            let bottom = mix(v01, v11, fract(coord.x - 0.5));
-            sampled_value = mix(top, bottom, fract(coord.y - 0.5));
+            let v00 = decode_color(Vec4(textureLoad(texture_sint, IVec2(coord + vec2(-0.5, -0.5)), 0)));
+            let v01 = decode_color(Vec4(textureLoad(texture_sint, IVec2(coord + vec2(-0.5,  0.5)), 0)));
+            let v10 = decode_color(Vec4(textureLoad(texture_sint, IVec2(coord + vec2( 0.5, -0.5)), 0)));
+            let v11 = decode_color(Vec4(textureLoad(texture_sint, IVec2(coord + vec2( 0.5,  0.5)), 0)));
+            normalized_value = filter_bilinear(coord, v00, v01, v10, v11);
         }
     } else if rect_info.sample_type == SAMPLE_TYPE_UINT {
-        // TODO(emilk): support premultiplying alpha on this path. Requires knowing the alpha range (255, 65535, …).
         let coord = in.texcoord * Vec2(textureDimensions(texture_uint).xy);
         if tex_filter(coord) == FILTER_NEAREST {
             // nearest
-            sampled_value = Vec4(textureLoad(texture_uint, IVec2(coord), 0));
+            normalized_value = decode_color(Vec4(textureLoad(texture_uint, IVec2(coord), 0)));
         } else {
             // bilinear
-            let v00 = Vec4(textureLoad(texture_uint, IVec2(coord + vec2(-0.5, -0.5)), 0));
-            let v01 = Vec4(textureLoad(texture_uint, IVec2(coord + vec2(-0.5,  0.5)), 0));
-            let v10 = Vec4(textureLoad(texture_uint, IVec2(coord + vec2( 0.5, -0.5)), 0));
-            let v11 = Vec4(textureLoad(texture_uint, IVec2(coord + vec2( 0.5,  0.5)), 0));
-            let top = mix(v00, v10, fract(coord.x - 0.5));
-            let bottom = mix(v01, v11, fract(coord.x - 0.5));
-            sampled_value = mix(top, bottom, fract(coord.y - 0.5));
+            let v00 = decode_color(Vec4(textureLoad(texture_uint, IVec2(coord + vec2(-0.5, -0.5)), 0)));
+            let v01 = decode_color(Vec4(textureLoad(texture_uint, IVec2(coord + vec2(-0.5,  0.5)), 0)));
+            let v10 = decode_color(Vec4(textureLoad(texture_uint, IVec2(coord + vec2( 0.5, -0.5)), 0)));
+            let v11 = decode_color(Vec4(textureLoad(texture_uint, IVec2(coord + vec2( 0.5,  0.5)), 0)));
+            normalized_value = filter_bilinear(coord, v00, v01, v10, v11);
         }
     } 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);
 

crates/re_renderer/src/renderer/rectangles.rs

@@ -53,14 +53,23 @@ pub enum TextureFilterMin {
 pub struct ColormappedTexture {
     pub texture: GpuTexture2D,
 
+    /// Min/max range of the values in the texture.
+    ///
+    /// Used to normalize the input values (squash them to the 0-1 range).
+    /// The normalization is applied before sRGB gamma decoding and alpha pre-multiplication
+    /// (this transformation is also applied to alpha!).
+    pub range: [f32; 2],
+
     /// Decode 0-1 sRGB gamma values to linear space before filtering?
     ///
     /// Only applies to [`wgpu::TextureFormat::Rgba8Unorm`] and float textures.
     pub decode_srgb: bool,
 
-    /// Min/max range of the values in the texture.
-    /// Used to normalize the input values (squash them to the 0-1 range).
-    pub range: [f32; 2],
+    /// Multiply color channels with the alpha channel before filtering?
+    ///
+    /// Set this to false for textures that don't have an alpha channel or are already pre-multiplied.
+    /// Applied after range normalization and srgb decoding, before filtering.
+    pub multiply_rgb_with_alpha: bool,
 
     /// Raise the normalized values to this power (before any color mapping).
     /// Acts like an inverse brightness.
@@ -102,6 +111,7 @@ impl ColormappedTexture {
             decode_srgb,
             range: [0.0, 1.0],
             gamma: 1.0,
+            multiply_rgb_with_alpha: true,
             color_mapper: None,
         }
     }
@@ -224,7 +234,8 @@ mod gpu_data {
         magnification_filter: u32,
 
         decode_srgb: u32,
-        _row_padding: [u32; 3],
+        multiply_rgb_with_alpha: u32,
+        _row_padding: [u32; 2],
 
         _end_padding: [wgpu_buffer_types::PaddingRow; 16 - 7],
     }
@@ -251,6 +262,7 @@ mod gpu_data {
                 range,
                 gamma,
                 color_mapper,
+                multiply_rgb_with_alpha,
             } = colormapped_texture;
 
             let super::RectangleOptions {
@@ -322,6 +334,7 @@ mod gpu_data {
                 minification_filter,
                 magnification_filter,
                 decode_srgb: *decode_srgb as _,
+                multiply_rgb_with_alpha: *multiply_rgb_with_alpha as _,
                 _row_padding: Default::default(),
                 _end_padding: Default::default(),
             })

crates/re_viewer/src/ui/view_tensor/tensor_slice_to_gpu.rs

@@ -4,7 +4,7 @@ use re_renderer::{
     resource_managers::{GpuTexture2D, Texture2DCreationDesc, TextureManager2DError},
 };
 use re_viewer_context::{
-    gpu_bridge::{self, range, RangeError},
+    gpu_bridge::{self, tensor_data_range_heuristic, RangeError},
     TensorStats,
 };
 
@@ -33,16 +33,17 @@ pub fn colormapped_texture(
 ) -> Result<ColormappedTexture, TextureManager2DError<TensorUploadError>> {
     crate::profile_function!();
 
-    let range =
-        range(tensor_stats).map_err(|err| TextureManager2DError::DataCreation(err.into()))?;
+    let range = tensor_data_range_heuristic(tensor_stats, tensor.dtype())
+        .map_err(|err| TextureManager2DError::DataCreation(err.into()))?;
     let texture = upload_texture_slice_to_gpu(render_ctx, tensor, state.slice())?;
 
     let color_mapping = state.color_mapping();
 
     Ok(ColormappedTexture {
         texture,
-        decode_srgb: false,
         range,
+        decode_srgb: false,
+        multiply_rgb_with_alpha: false,
         gamma: color_mapping.gamma,
         color_mapper: Some(re_renderer::renderer::ColorMapper::Function(
             color_mapping.map,

crates/re_viewer/src/ui/view_tensor/ui.rs

@@ -326,8 +326,9 @@ fn paint_colormap_gradient(
 
     let colormapped_texture = re_renderer::renderer::ColormappedTexture {
         texture: horizontal_gradient,
-        decode_srgb: false,
         range: [0.0, 1.0],
+        decode_srgb: false,
+        multiply_rgb_with_alpha: false,
         gamma: 1.0,
         color_mapper: Some(re_renderer::renderer::ColorMapper::Function(colormap)),
     };

crates/re_viewer_context/src/gpu_bridge/mod.rs

@@ -25,28 +25,59 @@ pub enum RangeError {
     /// This is weird. Should only happen with JPEGs, and those should have been decoded already
     #[error("Missing a range.")]
     MissingRange,
-
-    #[error("Non-finite range of values")]
-    NonfiniteRange,
 }
 
 /// Get a valid, finite range for the gpu to use.
-pub fn range(tensor_stats: &TensorStats) -> Result<[f32; 2], RangeError> {
-    let (min, max) = tensor_stats.range.ok_or(RangeError::MissingRange)?;
+pub fn tensor_data_range_heuristic(
+    tensor_stats: &TensorStats,
+    data_type: re_log_types::TensorDataType,
+) -> Result<[f32; 2], RangeError> {
+    let (min, max) = tensor_stats.finite_range.ok_or(RangeError::MissingRange)?;
 
     let min = min as f32;
     let max = max as f32;
 
-    if !min.is_finite() || !max.is_finite() {
-        Err(RangeError::NonfiniteRange)
+    // Apply heuristic for ranges that are typically expected depending on the data type and the finite (!) range.
+    // (we ignore NaN/Inf values heres, since they are usually there by accident!)
+    if data_type.is_float() && 0.0 <= min && max <= 1.0 {
+        // Float values that are all between 0 and 1, assume that this is the range.
+        Ok([0.0, 1.0])
+    } else if 0.0 <= min && max <= 255.0 {
+        // If all values are between 0 and 255, assume this is the range.
+        // (This is very common, independent of the data type)
+        Ok([0.0, 255.0])
     } else if min == max {
         // uniform range. This can explode the colormapping, so let's map all colors to the middle:
         Ok([min - 1.0, max + 1.0])
     } else {
+        // Use range as is if nothing matches.
         Ok([min, max])
     }
 }
 
+/// Return whether a tensor should be assumed to be encoded in sRGB color space ("gamma space", no EOTF applied).
+pub fn tensor_decode_srgb_gamma_heuristic(
+    tensor_stats: &TensorStats,
+    data_type: re_log_types::TensorDataType,
+    channels: u32,
+) -> Result<bool, RangeError> {
+    if matches!(channels, 1 | 3 | 4) {
+        let (min, max) = tensor_stats.finite_range.ok_or(RangeError::MissingRange)?;
+        #[allow(clippy::if_same_then_else)]
+        if 0.0 <= min && max <= 255.0 {
+            // If the range is suspiciously reminding us of a "regular image", assume sRGB.
+            Ok(true)
+        } else if data_type.is_float() && 0.0 <= min && max <= 1.0 {
+            // Floating point images between 0 and 1 are often sRGB as well.
+            Ok(true)
+        } else {
+            Ok(false)
+        }
+    } else {
+        Ok(false)
+    }
+}
+
 // ----------------------------------------------------------------------------
 
 pub fn viewport_resolution_in_pixels(clip_rect: egui::Rect, pixels_from_point: f32) -> [u32; 2] {