Use GPU colormapping when showing images in the GUI (#1865)

* Cleanup: move Default close to the struct definition

* Simplify code: use if-let-else-return

* Simplify code: no need for Arc

* Add EntityDataUi so that the Tensor ui function knows entity path

* Better naming: selection -> item

* Simplify code: no optional tensor stats

* Less use of anyhow

* Use GPU colormapping when showing tensors in GUI

* Link to issue

* Optimize pad_to_four_elements for debug builds

* Refactor: simpler arguments to show_zoomed_image_region_area_outline

* Fix missing meter argument

* Refactor: break up long function

* Less use of Arc

* Pipe annotation context to the hover preview

* Simplify `AnnotationMap::find`

* Use new GPU colormapper for the hover-zoom-in tooltip

* Refactor

* Add helper function for turning a Tensor into an image::DynamicImage

* Fix warning on web builds

* Add helper function `Tensor::could_be_dynamic_image`

* Implement click-to-copy and click-to-save for tensors without egui

* Convert histogram to the new system

* Remove the TensorImageCache

* Fix TODO formatting

* bug fixes and cleanups

* Rename some stuff

* Build-fix

* Simplify some code

crates/re_data_store/src/entity_properties.rs

@@ -68,6 +68,22 @@ pub struct EntityProperties {
     pub backproject_radius_scale: EditableAutoValue<f32>,
 }
 
+#[cfg(feature = "serde")]
+impl Default for EntityProperties {
+    fn default() -> Self {
+        Self {
+            visible: true,
+            visible_history: ExtraQueryHistory::default(),
+            interactive: true,
+            color_mapper: EditableAutoValue::default(),
+            pinhole_image_plane_distance: EditableAutoValue::default(),
+            backproject_depth: EditableAutoValue::Auto(true),
+            depth_from_world_scale: EditableAutoValue::default(),
+            backproject_radius_scale: EditableAutoValue::Auto(1.0),
+        }
+    }
+}
+
 #[cfg(feature = "serde")]
 impl EntityProperties {
     /// Multiply/and these together.
@@ -97,22 +113,6 @@ impl EntityProperties {
     }
 }
 
-#[cfg(feature = "serde")]
-impl Default for EntityProperties {
-    fn default() -> Self {
-        Self {
-            visible: true,
-            visible_history: ExtraQueryHistory::default(),
-            interactive: true,
-            color_mapper: EditableAutoValue::default(),
-            pinhole_image_plane_distance: EditableAutoValue::default(),
-            backproject_depth: EditableAutoValue::Auto(true),
-            depth_from_world_scale: EditableAutoValue::default(),
-            backproject_radius_scale: EditableAutoValue::Auto(1.0),
-        }
-    }
-}
-
 // ----------------------------------------------------------------------------
 
 /// When showing an entity in the history view, add this much history to it.

crates/re_log_types/Cargo.toml

@@ -30,7 +30,7 @@ ecolor = ["dep:ecolor"]
 glam = ["dep:glam", "dep:macaw"]
 
 ## Integration with the [`image`](https://crates.io/crates/image/) crate.
-image = ["dep:image"]
+image = ["dep:ecolor", "dep:image"]
 
 ## Enable (de)serialization using serde.
 serde = [

crates/re_log_types/src/component_types/mod.rs

@@ -58,11 +58,11 @@ pub use radius::Radius;
 pub use rect::Rect2D;
 pub use scalar::{Scalar, ScalarPlotProps};
 pub use size::Size3D;
-#[cfg(feature = "image")]
-pub use tensor::TensorImageError;
 pub use tensor::{
     Tensor, TensorCastError, TensorData, TensorDataMeaning, TensorDimension, TensorId,
 };
+#[cfg(feature = "image")]
+pub use tensor::{TensorImageLoadError, TensorImageSaveError};
 pub use text_entry::TextEntry;
 pub use transform::{Pinhole, Rigid3, Transform};
 pub use vec::{Vec2D, Vec3D, Vec4D};

crates/re_log_types/src/component_types/tensor.rs

@@ -196,6 +196,23 @@ impl TensorData {
     pub fn is_empty(&self) -> bool {
         self.size_in_bytes() == 0
     }
+
+    pub fn is_compressed_image(&self) -> bool {
+        match self {
+            Self::U8(_)
+            | Self::U16(_)
+            | Self::U32(_)
+            | Self::U64(_)
+            | Self::I8(_)
+            | Self::I16(_)
+            | Self::I32(_)
+            | Self::I64(_)
+            | Self::F32(_)
+            | Self::F64(_) => false,
+
+            Self::JPEG(_) => true,
+        }
+    }
 }
 
 impl std::fmt::Debug for TensorData {
@@ -588,9 +605,10 @@ impl<'a> TryFrom<&'a Tensor> for ::ndarray::ArrayViewD<'a, half::f16> {
 
 // ----------------------------------------------------------------------------
 
+/// Errors when loading [`Tensor`] from the [`image`] crate.
 #[cfg(feature = "image")]
 #[derive(thiserror::Error, Debug)]
-pub enum TensorImageError {
+pub enum TensorImageLoadError {
     #[error(transparent)]
     Image(#[from] image::ImageError),
 
@@ -604,6 +622,20 @@ pub enum TensorImageError {
     ReadError(#[from] std::io::Error),
 }
 
+/// Errors when converting [`Tensor`] to [`image`] images.
+#[cfg(feature = "image")]
+#[derive(thiserror::Error, Debug)]
+pub enum TensorImageSaveError {
+    #[error("Expected image-shaped tensor, got {0:?}")]
+    ShapeNotAnImage(Vec<TensorDimension>),
+
+    #[error("Cannot convert tensor with {0} channels and datatype {1} to an image")]
+    UnsupportedChannelsDtype(u64, TensorDataType),
+
+    #[error("The tensor data did not match tensor dimensions")]
+    BadData,
+}
+
 impl Tensor {
     pub fn new(
         tensor_id: TensorId,
@@ -630,20 +662,20 @@ impl Tensor {
     #[cfg(not(target_arch = "wasm32"))]
     pub fn tensor_from_jpeg_file(
         image_path: impl AsRef<std::path::Path>,
-    ) -> Result<Self, TensorImageError> {
+    ) -> Result<Self, TensorImageLoadError> {
         let jpeg_bytes = std::fs::read(image_path)?;
         Self::tensor_from_jpeg_bytes(jpeg_bytes)
     }
 
     /// Construct a tensor from the contents of a JPEG file.
     ///
     /// Requires the `image` feature.
-    pub fn tensor_from_jpeg_bytes(jpeg_bytes: Vec<u8>) -> Result<Self, TensorImageError> {
+    pub fn tensor_from_jpeg_bytes(jpeg_bytes: Vec<u8>) -> Result<Self, TensorImageLoadError> {
         use image::ImageDecoder as _;
         let jpeg = image::codecs::jpeg::JpegDecoder::new(std::io::Cursor::new(&jpeg_bytes))?;
         if jpeg.color_type() != image::ColorType::Rgb8 {
             // TODO(emilk): support gray-scale jpeg as well
-            return Err(TensorImageError::UnsupportedJpegColorType(
+            return Err(TensorImageLoadError::UnsupportedJpegColorType(
                 jpeg.color_type(),
             ));
         }
@@ -665,14 +697,14 @@ impl Tensor {
     /// Construct a tensor from something that can be turned into a [`image::DynamicImage`].
     ///
     /// Requires the `image` feature.
-    pub fn from_image(image: impl Into<image::DynamicImage>) -> Result<Self, TensorImageError> {
+    pub fn from_image(image: impl Into<image::DynamicImage>) -> Result<Self, TensorImageLoadError> {
         Self::from_dynamic_image(image.into())
     }
 
     /// Construct a tensor from [`image::DynamicImage`].
     ///
     /// Requires the `image` feature.
-    pub fn from_dynamic_image(image: image::DynamicImage) -> Result<Self, TensorImageError> {
+    pub fn from_dynamic_image(image: image::DynamicImage) -> Result<Self, TensorImageLoadError> {
         let (w, h) = (image.width(), image.height());
 
         let (depth, data) = match image {
@@ -706,7 +738,9 @@ impl Tensor {
             }
             _ => {
                 // It is very annoying that DynamicImage is #[non_exhaustive]
-                return Err(TensorImageError::UnsupportedImageColorType(image.color()));
+                return Err(TensorImageLoadError::UnsupportedImageColorType(
+                    image.color(),
+                ));
             }
         };
 
@@ -722,6 +756,119 @@ impl Tensor {
             meter: None,
         })
     }
+
+    /// Predicts if [`Self::to_dynamic_image`] is likely to succeed, without doing anything expensive
+    pub fn could_be_dynamic_image(&self) -> bool {
+        self.is_shaped_like_an_image()
+            && matches!(
+                self.dtype(),
+                TensorDataType::U8
+                    | TensorDataType::U16
+                    | TensorDataType::F16
+                    | TensorDataType::F32
+                    | TensorDataType::F64
+            )
+    }
+
+    /// Try to convert an image-like tensor into an [`image::DynamicImage`].
+    pub fn to_dynamic_image(&self) -> Result<image::DynamicImage, TensorImageSaveError> {
+        use ecolor::{gamma_u8_from_linear_f32, linear_u8_from_linear_f32};
+        use image::{DynamicImage, GrayImage, RgbImage, RgbaImage};
+
+        type Rgb16Image = image::ImageBuffer<image::Rgb<u16>, Vec<u16>>;
+        type Rgba16Image = image::ImageBuffer<image::Rgba<u16>, Vec<u16>>;
+        type Gray16Image = image::ImageBuffer<image::Luma<u16>, Vec<u16>>;
+
+        let [h, w, channels] = self
+            .image_height_width_channels()
+            .ok_or_else(|| TensorImageSaveError::ShapeNotAnImage(self.shape.clone()))?;
+        let w = w as u32;
+        let h = h as u32;
+
+        let dyn_img_result =
+            match (channels, &self.data) {
+                (1, TensorData::U8(buf)) => {
+                    GrayImage::from_raw(w, h, buf.as_slice().to_vec()).map(DynamicImage::ImageLuma8)
+                }
+                (1, TensorData::U16(buf)) => Gray16Image::from_raw(w, h, buf.as_slice().to_vec())
+                    .map(DynamicImage::ImageLuma16),
+                // TODO(emilk) f16
+                (1, TensorData::F32(buf)) => {
+                    let pixels = buf
+                        .iter()
+                        .map(|pixel| gamma_u8_from_linear_f32(*pixel))
+                        .collect();
+                    GrayImage::from_raw(w, h, pixels).map(DynamicImage::ImageLuma8)
+                }
+                (1, TensorData::F64(buf)) => {
+                    let pixels = buf
+                        .iter()
+                        .map(|&pixel| gamma_u8_from_linear_f32(pixel as f32))
+                        .collect();
+                    GrayImage::from_raw(w, h, pixels).map(DynamicImage::ImageLuma8)
+                }
+
+                (3, TensorData::U8(buf)) => {
+                    RgbImage::from_raw(w, h, buf.as_slice().to_vec()).map(DynamicImage::ImageRgb8)
+                }
+                (3, TensorData::U16(buf)) => Rgb16Image::from_raw(w, h, buf.as_slice().to_vec())
+                    .map(DynamicImage::ImageRgb16),
+                (3, TensorData::F32(buf)) => {
+                    let pixels = buf.iter().copied().map(gamma_u8_from_linear_f32).collect();
+                    RgbImage::from_raw(w, h, pixels).map(DynamicImage::ImageRgb8)
+                }
+                (3, TensorData::F64(buf)) => {
+                    let pixels = buf
+                        .iter()
+                        .map(|&comp| gamma_u8_from_linear_f32(comp as f32))
+                        .collect();
+                    RgbImage::from_raw(w, h, pixels).map(DynamicImage::ImageRgb8)
+                }
+
+                (4, TensorData::U8(buf)) => {
+                    RgbaImage::from_raw(w, h, buf.as_slice().to_vec()).map(DynamicImage::ImageRgba8)
+                }
+                (4, TensorData::U16(buf)) => Rgba16Image::from_raw(w, h, buf.as_slice().to_vec())
+                    .map(DynamicImage::ImageRgba16),
+                (4, TensorData::F32(buf)) => {
+                    let rgba: &[[f32; 4]] = bytemuck::cast_slice(buf.as_slice());
+                    let pixels: Vec<u8> = rgba
+                        .iter()
+                        .flat_map(|&[r, g, b, a]| {
+                            let r = gamma_u8_from_linear_f32(r);
+                            let g = gamma_u8_from_linear_f32(g);
+                            let b = gamma_u8_from_linear_f32(b);
+                            let a = linear_u8_from_linear_f32(a);
+                            [r, g, b, a]
+                        })
+                        .collect();
+                    RgbaImage::from_raw(w, h, pixels).map(DynamicImage::ImageRgba8)
+                }
+                (4, TensorData::F64(buf)) => {
+                    let rgba: &[[f64; 4]] = bytemuck::cast_slice(buf.as_slice());
+                    let pixels: Vec<u8> = rgba
+                        .iter()
+                        .flat_map(|&[r, g, b, a]| {
+                            let r = gamma_u8_from_linear_f32(r as _);
+                            let g = gamma_u8_from_linear_f32(g as _);
+                            let b = gamma_u8_from_linear_f32(b as _);
+                            let a = linear_u8_from_linear_f32(a as _);
+                            [r, g, b, a]
+                        })
+                        .collect();
+                    RgbaImage::from_raw(w, h, pixels).map(DynamicImage::ImageRgba8)
+                }
+
+                (_, _) => {
+                    return Err(TensorImageSaveError::UnsupportedChannelsDtype(
+                        channels,
+                        self.data.dtype(),
+                    ))
+                }
+            };
+
+        dyn_img_result.ok_or(TensorImageSaveError::BadData)
+    }
 }
 
 // ----------------------------------------------------------------------------

crates/re_renderer/src/renderer/rectangles.rs

@@ -49,6 +49,7 @@ pub enum TextureFilterMin {
 }
 
 /// Describes a texture and how to map it to a color.
+#[derive(Clone)]
 pub struct ColormappedTexture {
     pub texture: GpuTexture2DHandle,
 
@@ -71,6 +72,7 @@ pub struct ColormappedTexture {
 }
 
 /// How to map the normalized `.r` component to a color.
+#[derive(Clone)]
 pub enum ColorMapper {
     /// Apply the given function.
     Function(Colormap),

crates/re_renderer/src/resource_managers/texture_manager.rs

@@ -31,6 +31,11 @@ impl GpuTexture2DHandle {
     pub fn height(&self) -> u32 {
         self.0.as_ref().map_or(1, |t| t.texture.height())
     }
+
+    /// Width and height of the texture.
+    pub fn width_height(&self) -> [u32; 2] {
+        [self.width(), self.height()]
+    }
 }
 
 /// Data required to create a texture 2d resource.

crates/re_renderer/src/view_builder.rs

@@ -10,7 +10,7 @@ use crate::{
     },
     global_bindings::FrameUniformBuffer,
     renderer::{CompositorDrawData, DebugOverlayDrawData, DrawData, Renderer},
-    wgpu_resources::{GpuBindGroup, GpuTexture, TextureDesc},
+    wgpu_resources::{GpuBindGroup, GpuTexture, PoolError, TextureDesc},
     DebugLabel, IntRect, Rgba, Size,
 };
 
@@ -498,7 +498,7 @@ impl ViewBuilder {
         &mut self,
         ctx: &RenderContext,
         clear_color: Rgba,
-    ) -> anyhow::Result<wgpu::CommandBuffer> {
+    ) -> Result<wgpu::CommandBuffer, PoolError> {
         crate::profile_function!();
 
         let setup = &self.setup;

crates/re_viewer/src/gpu_bridge/tensor_to_gpu.rs

@@ -414,9 +414,22 @@ fn narrow_f64_to_f32s(slice: &[f64]) -> Cow<'static, [u8]> {
 
 fn pad_to_four_elements<T: Copy>(data: &[T], pad: T) -> Vec<T> {
     crate::profile_function!();
-    data.chunks_exact(3)
-        .flat_map(|chunk| [chunk[0], chunk[1], chunk[2], pad])
-        .collect::<Vec<T>>()
+    if cfg!(debug_assertions) {
+        // fastest version in debug builds.
+        // 5x faster in debug builds, but 2x slower in release
+        let mut padded = vec![pad; data.len() / 3 * 4];
+        for i in 0..(data.len() / 3) {
+            padded[4 * i] = data[3 * i];
+            padded[4 * i + 1] = data[3 * i + 1];
+            padded[4 * i + 2] = data[3 * i + 2];
+        }
+        padded
+    } else {
+        // fastest version in optimized builds
+        data.chunks_exact(3)
+            .flat_map(|chunk| [chunk[0], chunk[1], chunk[2], pad])
+            .collect()
+    }
 }
 
 fn pad_and_cast<T: Copy + Pod>(data: &[T], pad: T) -> Cow<'static, [u8]> {