From 041b47401ba09276e3a724f0bfc1e0dbf707ae1c Mon Sep 17 00:00:00 2001
From: Emil Ernerfeldt <emil.ernerfeldt@gmail.com>
Date: Thu, 23 Mar 2023 12:46:08 +0100
Subject: [PATCH 1/5] Refactor how tensor image cache turns tensors into images

---
 .../src/component_types/tensor.rs             |  62 ++--
 .../src/misc/caches/tensor_image_cache.rs     | 274 +++++++++++-------
 2 files changed, 214 insertions(+), 122 deletions(-)

diff --git a/crates/re_log_types/src/component_types/tensor.rs b/crates/re_log_types/src/component_types/tensor.rs
index 4b678232d7d2..502ab2b73c4c 100644
--- a/crates/re_log_types/src/component_types/tensor.rs
+++ b/crates/re_log_types/src/component_types/tensor.rs
@@ -174,6 +174,42 @@ pub enum TensorData {
     JPEG(BinaryBuffer),
 }
 
+impl TensorData {
+    pub fn dtype(&self) -> TensorDataType {
+        match self {
+            Self::U8(_) | Self::JPEG(_) => TensorDataType::U8,
+            Self::U16(_) => TensorDataType::U16,
+            Self::U32(_) => TensorDataType::U32,
+            Self::U64(_) => TensorDataType::U64,
+            Self::I8(_) => TensorDataType::I8,
+            Self::I16(_) => TensorDataType::I16,
+            Self::I32(_) => TensorDataType::I32,
+            Self::I64(_) => TensorDataType::I64,
+            Self::F32(_) => TensorDataType::F32,
+            Self::F64(_) => TensorDataType::F64,
+        }
+    }
+
+    pub fn size_in_bytes(&self) -> usize {
+        match self {
+            Self::U8(buf) | Self::JPEG(buf) => buf.0.len(),
+            Self::U16(buf) => buf.len(),
+            Self::U32(buf) => buf.len(),
+            Self::U64(buf) => buf.len(),
+            Self::I8(buf) => buf.len(),
+            Self::I16(buf) => buf.len(),
+            Self::I32(buf) => buf.len(),
+            Self::I64(buf) => buf.len(),
+            Self::F32(buf) => buf.len(),
+            Self::F64(buf) => buf.len(),
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.size_in_bytes() == 0
+    }
+}
+
 /// Flattened `Tensor` data payload
 ///
 /// ## Examples
@@ -394,33 +430,11 @@ impl TensorTrait for Tensor {
     }
 
     fn dtype(&self) -> TensorDataType {
-        match &self.data {
-            TensorData::U8(_) | TensorData::JPEG(_) => TensorDataType::U8,
-            TensorData::U16(_) => TensorDataType::U16,
-            TensorData::U32(_) => TensorDataType::U32,
-            TensorData::U64(_) => TensorDataType::U64,
-            TensorData::I8(_) => TensorDataType::I8,
-            TensorData::I16(_) => TensorDataType::I16,
-            TensorData::I32(_) => TensorDataType::I32,
-            TensorData::I64(_) => TensorDataType::I64,
-            TensorData::F32(_) => TensorDataType::F32,
-            TensorData::F64(_) => TensorDataType::F64,
-        }
+        self.data.dtype()
     }
 
     fn size_in_bytes(&self) -> usize {
-        match &self.data {
-            TensorData::U8(buf) | TensorData::JPEG(buf) => buf.0.len(),
-            TensorData::U16(buf) => buf.len(),
-            TensorData::U32(buf) => buf.len(),
-            TensorData::U64(buf) => buf.len(),
-            TensorData::I8(buf) => buf.len(),
-            TensorData::I16(buf) => buf.len(),
-            TensorData::I32(buf) => buf.len(),
-            TensorData::I64(buf) => buf.len(),
-            TensorData::F32(buf) => buf.len(),
-            TensorData::F64(buf) => buf.len(),
-        }
+        self.data.size_in_bytes()
     }
 }
 
diff --git a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs
index fba14a2af73e..3a99774256c6 100644
--- a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs
+++ b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs
@@ -239,12 +239,20 @@ impl CachedImage {
 }
 
 fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::Result<ColorImage> {
+    match tensor.meaning {
+        TensorDataMeaning::Unknown => color_tensor_as_color_image(tensor),
+        TensorDataMeaning::ClassId => class_id_tensor_as_color_image(tensor, annotations),
+        TensorDataMeaning::Depth => depth_tensor_as_color_image(tensor),
+    }
+}
+
+fn color_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> {
     use anyhow::Context as _;
 
     crate::profile_function!(format!(
-        "dtype: {}, meaning: {:?}",
+        "dtype: {}, shape: {:?}",
         tensor.dtype(),
-        tensor.meaning
+        tensor.shape()
     ));
 
     let shape = &tensor.shape();
@@ -274,43 +282,8 @@ fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::R
 
     let size = [width as _, height as _];
 
-    match (depth, &tensor.data, tensor.meaning) {
-        (1, TensorData::U8(buf), TensorDataMeaning::ClassId) => {
-            // Apply annotation mapping to raw bytes interpreted as u8
-            let color_lookup: Vec<Color32> = (0..256)
-                .map(|id| {
-                    annotations
-                        .class_description(Some(ClassId(id)))
-                        .annotation_info()
-                        .color(None, DefaultColor::TransparentBlack)
-                })
-                .collect();
-            let pixels: Vec<Color32> = buf
-                .0
-                .iter()
-                .map(|p: &u8| color_lookup[*p as usize])
-                .collect();
-            crate::profile_scope!("from_raw");
-            Ok(ColorImage { size, pixels })
-        }
-        (1, TensorData::U16(buf), TensorDataMeaning::ClassId) => {
-            // Apply annotations mapping to bytes interpreted as u16
-            let mut color_lookup: ahash::HashMap<u16, Color32> = Default::default();
-            let pixels = buf
-                .iter()
-                .map(|id: &u16| {
-                    *color_lookup.entry(*id).or_insert_with(|| {
-                        annotations
-                            .class_description(Some(ClassId(*id)))
-                            .annotation_info()
-                            .color(None, DefaultColor::TransparentBlack)
-                    })
-                })
-                .collect();
-            crate::profile_scope!("from_raw");
-            Ok(ColorImage { size, pixels })
-        }
-        (1, TensorData::U8(buf), _) => {
+    match (depth, &tensor.data) {
+        (1, TensorData::U8(buf)) => {
             // TODO(emilk): we should read some meta-data to check if this is luminance or alpha.
             let pixels = buf
                 .0
@@ -319,7 +292,7 @@ fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::R
                 .collect();
             Ok(ColorImage { size, pixels })
         }
-        (1, TensorData::U16(buf), _) => {
+        (1, TensorData::U16(buf)) => {
             // TODO(emilk): we should read some meta-data to check if this is luminance or alpha.
             let pixels = buf
                 .iter()
@@ -328,50 +301,7 @@ fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::R
 
             Ok(ColorImage { size, pixels })
         }
-        (1, TensorData::F32(buf), TensorDataMeaning::Depth) => {
-            if buf.is_empty() {
-                Ok(ColorImage::default())
-            } else {
-                // Convert to u16 so we can put them in an image.
-                // TODO(emilk): Eventually we want a renderer that can show f32 images natively.
-                // One big downside of the approach below is that if we have two depth images
-                // in the same range, they cannot be visually compared with each other,
-                // because their individual max-depths will be scaled to 65535.
-
-                let mut min = f32::INFINITY;
-                let mut max = f32::NEG_INFINITY;
-                for float in buf.iter() {
-                    min = min.min(*float);
-                    max = max.max(*float);
-                }
-
-                anyhow::ensure!(
-                    min.is_finite() && max.is_finite(),
-                    "Depth image had non-finite values"
-                );
-
-                let ints: Vec<u16> = if min == max {
-                    // Uniform image. We can't remap it to a 0-1 range, so do whatever:
-                    buf.iter().map(|&float| float as u16).collect()
-                } else {
-                    buf.iter()
-                        .map(|&float| egui::remap(float, min..=max, 0.0..=65535.0) as u16)
-                        .collect()
-                };
-
-                let pixels = ints
-                    .iter()
-                    .map(|pixel| {
-                        let [r, g, b, _] =
-                            re_renderer::colormap_turbo_srgb((*pixel as f32) / (u16::MAX as f32));
-                        egui::Color32::from_rgb(r, g, b)
-                    })
-                    .collect();
-
-                Ok(ColorImage { size, pixels })
-            }
-        }
-        (1, TensorData::F32(buf), _) => {
+        (1, TensorData::F32(buf)) => {
             let pixels = buf
                 .iter()
                 .map(|pixel| Color32::from_gray(linear_u8_from_linear_f32(*pixel)))
@@ -379,13 +309,13 @@ fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::R
 
             Ok(ColorImage { size, pixels })
         }
-        (3, TensorData::U8(buf), _) => Ok(ColorImage::from_rgb(size, buf.0.as_slice())),
-        (3, TensorData::U16(buf), _) => {
+        (3, TensorData::U8(buf)) => Ok(ColorImage::from_rgb(size, buf.0.as_slice())),
+        (3, TensorData::U16(buf)) => {
             let u8_buf: Vec<u8> = buf.iter().map(|pixel| (*pixel / 256) as u8).collect();
 
             Ok(ColorImage::from_rgb(size, &u8_buf))
         }
-        (3, TensorData::F32(buf), _) => {
+        (3, TensorData::F32(buf)) => {
             let rgb: &[[f32; 3]] = bytemuck::cast_slice(buf.as_slice());
             let pixels: Vec<Color32> = rgb
                 .iter()
@@ -400,15 +330,13 @@ fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::R
             Ok(ColorImage { size, pixels })
         }
 
-        (4, TensorData::U8(buf), _) => {
-            Ok(ColorImage::from_rgba_unmultiplied(size, buf.0.as_slice()))
-        }
-        (4, TensorData::U16(buf), _) => {
+        (4, TensorData::U8(buf)) => Ok(ColorImage::from_rgba_unmultiplied(size, buf.0.as_slice())),
+        (4, TensorData::U16(buf)) => {
             let u8_buf: Vec<u8> = buf.iter().map(|pixel| (*pixel / 256) as u8).collect();
 
             Ok(ColorImage::from_rgba_unmultiplied(size, &u8_buf))
         }
-        (4, TensorData::F32(buf), _) => {
+        (4, TensorData::F32(buf)) => {
             let rgba: &[[f32; 4]] = bytemuck::cast_slice(buf.as_slice());
             let pixels: Vec<Color32> = rgba
                 .iter()
@@ -423,18 +351,168 @@ fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::R
 
             Ok(ColorImage { size, pixels })
         }
-        (_, TensorData::JPEG(_), _) => {
-            anyhow::bail!("JPEG tensor should have been decoded before using TensorImageCache")
+
+        (_depth, dtype) => {
+            anyhow::bail!("Don't know how to turn a tensor of shape={shape:?} and dtype={dtype:?} into a color image")
         }
+    }
+}
+
+fn class_id_tensor_as_color_image(
+    tensor: &Tensor,
+    annotations: &Annotations,
+) -> anyhow::Result<ColorImage> {
+    use anyhow::Context as _;
+
+    crate::profile_function!(format!(
+        "dtype: {}, shape: {:?}",
+        tensor.dtype(),
+        tensor.shape()
+    ));
+
+    let shape = &tensor.shape();
 
-        (_depth, dtype, meaning @ TensorDataMeaning::ClassId) => {
+    anyhow::ensure!(
+        shape.len() == 2 || shape.len() == 3,
+        "Expected a 2D or 3D tensor, got {shape:?}",
+    );
+
+    let [height, width] = [
+        u32::try_from(shape[0].size).context("tensor too large")?,
+        u32::try_from(shape[1].size).context("tensor too large")?,
+    ];
+    let depth = if shape.len() == 2 { 1 } else { shape[2].size };
+
+    anyhow::ensure!(
+        depth == 1,
+        "Cannot apply annotations to tensor of shape {shape:?}"
+    );
+    debug_assert!(
+        tensor.is_shaped_like_an_image(),
+        "We should make the same checks above, but with actual error messages"
+    );
+    let size = [width as _, height as _];
+
+    match &tensor.data {
+        TensorData::U8(buf) => {
+            // Apply annotation mapping to raw bytes interpreted as u8
+            let color_lookup: Vec<Color32> = (0..256)
+                .map(|id| {
+                    annotations
+                        .class_description(Some(ClassId(id)))
+                        .annotation_info()
+                        .color(None, DefaultColor::TransparentBlack)
+                })
+                .collect();
+            let pixels: Vec<Color32> = buf
+                .0
+                .iter()
+                .map(|p: &u8| color_lookup[*p as usize])
+                .collect();
+            crate::profile_scope!("from_raw");
+            Ok(ColorImage { size, pixels })
+        }
+        TensorData::U16(buf) => {
+            // Apply annotations mapping to bytes interpreted as u16
+            let mut color_lookup: ahash::HashMap<u16, Color32> = Default::default();
+            let pixels = buf
+                .iter()
+                .map(|id: &u16| {
+                    *color_lookup.entry(*id).or_insert_with(|| {
+                        annotations
+                            .class_description(Some(ClassId(*id)))
+                            .annotation_info()
+                            .color(None, DefaultColor::TransparentBlack)
+                    })
+                })
+                .collect();
+            crate::profile_scope!("from_raw");
+            Ok(ColorImage { size, pixels })
+        }
+        _ => {
             anyhow::bail!(
-                "Shape={shape:?} and dtype={dtype:?} is incompatible with meaning={meaning:?}"
+                "Cannot apply annotations to tensor of dtype {}",
+                tensor.dtype()
             )
         }
+    }
+}
+
+fn depth_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> {
+    if tensor.data.is_empty() {
+        return Ok(ColorImage::default());
+    }
+
+    use anyhow::Context as _;
+
+    crate::profile_function!(format!(
+        "dtype: {}, shape: {:?}",
+        tensor.dtype(),
+        tensor.shape()
+    ));
 
-        (_depth, dtype, _) => {
-            anyhow::bail!("Don't know how to turn a tensor of shape={shape:?} and dtype={dtype:?} into an image")
+    let shape = &tensor.shape();
+
+    anyhow::ensure!(
+        shape.len() == 2 || shape.len() == 3,
+        "Expected a 2D or 3D tensor, got {shape:?}",
+    );
+
+    let [height, width] = [
+        u32::try_from(shape[0].size).context("tensor too large")?,
+        u32::try_from(shape[1].size).context("tensor too large")?,
+    ];
+    let depth = if shape.len() == 2 { 1 } else { shape[2].size };
+
+    anyhow::ensure!(depth == 1, "Depth tensor of shape {shape:?}");
+    debug_assert!(
+        tensor.is_shaped_like_an_image(),
+        "We should make the same checks above, but with actual error messages"
+    );
+    let size = [width as _, height as _];
+
+    match &tensor.data {
+        TensorData::F32(buf) => {
+            // Convert to u16 so we can put them in an image.
+            // TODO(emilk): Eventually we want a renderer that can show f32 images natively.
+            // One big downside of the approach below is that if we have two depth images
+            // in the same range, they cannot be visually compared with each other,
+            // because their individual max-depths will be scaled to 65535.
+
+            let mut min = f32::INFINITY;
+            let mut max = f32::NEG_INFINITY;
+            for float in buf.iter() {
+                min = min.min(*float);
+                max = max.max(*float);
+            }
+
+            anyhow::ensure!(
+                min.is_finite() && max.is_finite(),
+                "Depth image had non-finite values"
+            );
+
+            let ints: Vec<u16> = if min == max {
+                // Uniform image. We can't remap it to a 0-1 range, so do whatever:
+                buf.iter().map(|&float| float as u16).collect()
+            } else {
+                buf.iter()
+                    .map(|&float| egui::remap(float, min..=max, 0.0..=65535.0) as u16)
+                    .collect()
+            };
+
+            let pixels = ints
+                .iter()
+                .map(|pixel| {
+                    let [r, g, b, _] =
+                        re_renderer::colormap_turbo_srgb((*pixel as f32) / (u16::MAX as f32));
+                    egui::Color32::from_rgb(r, g, b)
+                })
+                .collect();
+
+            Ok(ColorImage { size, pixels })
+        }
+        _ => {
+            color_tensor_as_color_image(tensor) // TODO(emilk): support more depth dtypes
         }
     }
 }

From 5fe43c7d2778b77d7a07cde3cb0cd4734b7e0403 Mon Sep 17 00:00:00 2001
From: Emil Ernerfeldt <emil.ernerfeldt@gmail.com>
Date: Thu, 23 Mar 2023 12:53:32 +0100
Subject: [PATCH 2/5] simplify the code

---
 .../src/misc/caches/tensor_image_cache.rs     | 70 ++++++-------------
 1 file changed, 20 insertions(+), 50 deletions(-)

diff --git a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs
index 3a99774256c6..29cfd93489b2 100644
--- a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs
+++ b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs
@@ -246,15 +246,9 @@ fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::R
     }
 }
 
-fn color_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> {
+fn height_width_depth(tensor: &Tensor) -> anyhow::Result<[u32; 3]> {
     use anyhow::Context as _;
 
-    crate::profile_function!(format!(
-        "dtype: {}, shape: {:?}",
-        tensor.dtype(),
-        tensor.shape()
-    ));
-
     let shape = &tensor.shape();
 
     anyhow::ensure!(
@@ -277,8 +271,19 @@ fn color_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> {
         "We should make the same checks above, but with actual error messages"
     );
 
-    use egui::epaint::ecolor::gamma_u8_from_linear_f32;
-    use egui::epaint::ecolor::linear_u8_from_linear_f32;
+    Ok([height, width, depth as u32])
+}
+
+fn color_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> {
+    crate::profile_function!(format!(
+        "dtype: {}, shape: {:?}",
+        tensor.dtype(),
+        tensor.shape()
+    ));
+
+    let [height, width, depth] = height_width_depth(tensor)?;
+
+    use egui::epaint::ecolor::{gamma_u8_from_linear_f32, linear_u8_from_linear_f32};
 
     let size = [width as _, height as _];
 
@@ -353,7 +358,7 @@ fn color_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> {
         }
 
         (_depth, dtype) => {
-            anyhow::bail!("Don't know how to turn a tensor of shape={shape:?} and dtype={dtype:?} into a color image")
+            anyhow::bail!("Don't know how to turn a tensor of shape={:?} and dtype={dtype:?} into a color image", tensor.shape)
         }
     }
 }
@@ -362,34 +367,17 @@ fn class_id_tensor_as_color_image(
     tensor: &Tensor,
     annotations: &Annotations,
 ) -> anyhow::Result<ColorImage> {
-    use anyhow::Context as _;
-
     crate::profile_function!(format!(
         "dtype: {}, shape: {:?}",
         tensor.dtype(),
         tensor.shape()
     ));
 
-    let shape = &tensor.shape();
-
-    anyhow::ensure!(
-        shape.len() == 2 || shape.len() == 3,
-        "Expected a 2D or 3D tensor, got {shape:?}",
-    );
-
-    let [height, width] = [
-        u32::try_from(shape[0].size).context("tensor too large")?,
-        u32::try_from(shape[1].size).context("tensor too large")?,
-    ];
-    let depth = if shape.len() == 2 { 1 } else { shape[2].size };
-
+    let [height, width, depth] = height_width_depth(tensor)?;
     anyhow::ensure!(
         depth == 1,
-        "Cannot apply annotations to tensor of shape {shape:?}"
-    );
-    debug_assert!(
-        tensor.is_shaped_like_an_image(),
-        "We should make the same checks above, but with actual error messages"
+        "Cannot apply annotations to tensor of shape {:?}",
+        tensor.shape
     );
     let size = [width as _, height as _];
 
@@ -443,32 +431,14 @@ fn depth_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> {
         return Ok(ColorImage::default());
     }
 
-    use anyhow::Context as _;
-
     crate::profile_function!(format!(
         "dtype: {}, shape: {:?}",
         tensor.dtype(),
         tensor.shape()
     ));
 
-    let shape = &tensor.shape();
-
-    anyhow::ensure!(
-        shape.len() == 2 || shape.len() == 3,
-        "Expected a 2D or 3D tensor, got {shape:?}",
-    );
-
-    let [height, width] = [
-        u32::try_from(shape[0].size).context("tensor too large")?,
-        u32::try_from(shape[1].size).context("tensor too large")?,
-    ];
-    let depth = if shape.len() == 2 { 1 } else { shape[2].size };
-
-    anyhow::ensure!(depth == 1, "Depth tensor of shape {shape:?}");
-    debug_assert!(
-        tensor.is_shaped_like_an_image(),
-        "We should make the same checks above, but with actual error messages"
-    );
+    let [height, width, depth] = height_width_depth(tensor)?;
+    anyhow::ensure!(depth == 1, "Depth tensor of shape {:?}", tensor.shape);
     let size = [width as _, height as _];
 
     match &tensor.data {

From 33227cefcfd505fa6b280ecf73af0e445e64b16b Mon Sep 17 00:00:00 2001
From: Emil Ernerfeldt <emil.ernerfeldt@gmail.com>
Date: Thu, 23 Mar 2023 14:44:52 +0100
Subject: [PATCH 3/5] Refactor how we color map depth images

---
 crates/re_log_types/src/data.rs               | 32 +++++++++
 .../src/misc/caches/tensor_image_cache.rs     | 72 ++++++++++---------
 2 files changed, 70 insertions(+), 34 deletions(-)

diff --git a/crates/re_log_types/src/data.rs b/crates/re_log_types/src/data.rs
index 7d59384cdac8..779adfff4cf6 100644
--- a/crates/re_log_types/src/data.rs
+++ b/crates/re_log_types/src/data.rs
@@ -93,6 +93,38 @@ impl TensorDataType {
             Self::F64 => std::mem::size_of::<f64>() as _,
         }
     }
+
+    pub fn is_float(&self) -> bool {
+        match self {
+            Self::U8
+            | Self::U16
+            | Self::U32
+            | Self::U64
+            | Self::I8
+            | Self::I16
+            | Self::I32
+            | Self::I64 => false,
+            Self::F16 | Self::F32 | Self::F64 => true,
+        }
+    }
+
+    pub fn max_value(&self) -> f64 {
+        match self {
+            Self::U8 => u8::MAX as _,
+            Self::U16 => u16::MAX as _,
+            Self::U32 => u32::MAX as _,
+            Self::U64 => u64::MAX as _,
+
+            Self::I8 => i8::MAX as _,
+            Self::I16 => i16::MAX as _,
+            Self::I32 => i32::MAX as _,
+            Self::I64 => i64::MAX as _,
+
+            Self::F16 => f16::MAX.into(),
+            Self::F32 => f32::MAX as _,
+            Self::F64 => f64::MAX,
+        }
+    }
 }
 
 impl std::fmt::Display for TensorDataType {
diff --git a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs
index 29cfd93489b2..2f3c97406f4b 100644
--- a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs
+++ b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs
@@ -12,7 +12,10 @@ use re_renderer::{
     RenderContext,
 };
 
-use crate::ui::{Annotations, DefaultColor, MISSING_ANNOTATIONS};
+use crate::{
+    misc::caches::TensorStats,
+    ui::{Annotations, DefaultColor, MISSING_ANNOTATIONS},
+};
 
 // ---
 
@@ -431,6 +434,12 @@ fn depth_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> {
         return Ok(ColorImage::default());
     }
 
+    // This function applies color mapping to a depth image.
+    // We are planning on moving this to the GPU: https://github.com/rerun-io/rerun/issues/1612
+    // One big downside of the approach below is that if we have two depth images
+    // in the same range, they cannot be visually compared with each other,
+    // because their individual max-depths will be scaled to 65535.
+
     crate::profile_function!(format!(
         "dtype: {}, shape: {:?}",
         tensor.dtype(),
@@ -441,42 +450,37 @@ fn depth_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> {
     anyhow::ensure!(depth == 1, "Depth tensor of shape {:?}", tensor.shape);
     let size = [width as _, height as _];
 
+    let range = TensorStats::new(tensor).range.ok_or(anyhow::anyhow!(
+        "Depth image had no range!? Was this compressed?"
+    ))?;
+
+    let (mut min, mut max) = range;
+
+    anyhow::ensure!(
+        min.is_finite() && max.is_finite(),
+        "Depth image had non-finite values"
+    );
+
+    if min == max {
+        // Uniform image. We can't remap it to a 0-1 range, so do whatever:
+        min = 0.0;
+        max = if tensor.dtype().is_float() {
+            1.0
+        } else {
+            tensor.dtype().max_value()
+        };
+    }
+
+    fn colormap(t: f32) -> egui::Color32 {
+        let [r, g, b, _] = re_renderer::colormap_turbo_srgb(t);
+        egui::Color32::from_rgb(r, g, b)
+    }
+
     match &tensor.data {
         TensorData::F32(buf) => {
-            // Convert to u16 so we can put them in an image.
-            // TODO(emilk): Eventually we want a renderer that can show f32 images natively.
-            // One big downside of the approach below is that if we have two depth images
-            // in the same range, they cannot be visually compared with each other,
-            // because their individual max-depths will be scaled to 65535.
-
-            let mut min = f32::INFINITY;
-            let mut max = f32::NEG_INFINITY;
-            for float in buf.iter() {
-                min = min.min(*float);
-                max = max.max(*float);
-            }
-
-            anyhow::ensure!(
-                min.is_finite() && max.is_finite(),
-                "Depth image had non-finite values"
-            );
-
-            let ints: Vec<u16> = if min == max {
-                // Uniform image. We can't remap it to a 0-1 range, so do whatever:
-                buf.iter().map(|&float| float as u16).collect()
-            } else {
-                buf.iter()
-                    .map(|&float| egui::remap(float, min..=max, 0.0..=65535.0) as u16)
-                    .collect()
-            };
-
-            let pixels = ints
+            let pixels = buf
                 .iter()
-                .map(|pixel| {
-                    let [r, g, b, _] =
-                        re_renderer::colormap_turbo_srgb((*pixel as f32) / (u16::MAX as f32));
-                    egui::Color32::from_rgb(r, g, b)
-                })
+                .map(|&float| colormap(egui::remap(float, min as f32..=max as f32, 0.0..=1.0)))
                 .collect();
 
             Ok(ColorImage { size, pixels })

From 59e3201d084c38dd44375a10e9be772fefff2e76 Mon Sep 17 00:00:00 2001
From: Emil Ernerfeldt <emil.ernerfeldt@gmail.com>
Date: Thu, 23 Mar 2023 14:51:35 +0100
Subject: [PATCH 4/5] Apply colormaps to all types of depth images

---
 .../component_types/arrow_convert_shims.rs    |  5 ++
 .../src/misc/caches/tensor_image_cache.rs     | 54 +++++++++++++++----
 2 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/crates/re_log_types/src/component_types/arrow_convert_shims.rs b/crates/re_log_types/src/component_types/arrow_convert_shims.rs
index 050575101229..8d196842169d 100644
--- a/crates/re_log_types/src/component_types/arrow_convert_shims.rs
+++ b/crates/re_log_types/src/component_types/arrow_convert_shims.rs
@@ -27,6 +27,11 @@ impl BinaryBuffer {
     pub fn as_slice(&self) -> &[u8] {
         self.0.as_slice()
     }
+
+    #[inline]
+    pub fn iter(&self) -> impl Iterator<Item = &u8> {
+        self.0.iter()
+    }
 }
 
 impl Index<usize> for BinaryBuffer {
diff --git a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs
index 2f3c97406f4b..7ecb0bb1b0d7 100644
--- a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs
+++ b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs
@@ -471,22 +471,58 @@ fn depth_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> {
         };
     }
 
-    fn colormap(t: f32) -> egui::Color32 {
+    let colormap = |value: f64| {
+        let t = egui::remap(value, min..=max, 0.0..=1.0) as f32;
         let [r, g, b, _] = re_renderer::colormap_turbo_srgb(t);
         egui::Color32::from_rgb(r, g, b)
-    }
+    };
 
     match &tensor.data {
-        TensorData::F32(buf) => {
-            let pixels = buf
-                .iter()
-                .map(|&float| colormap(egui::remap(float, min as f32..=max as f32, 0.0..=1.0)))
-                .collect();
+        TensorData::U8(buf) => {
+            let pixels = buf.iter().map(|&value| colormap(value as _)).collect();
+            Ok(ColorImage { size, pixels })
+        }
+        TensorData::U16(buf) => {
+            let pixels = buf.iter().map(|&value| colormap(value as _)).collect();
+            Ok(ColorImage { size, pixels })
+        }
+        TensorData::U32(buf) => {
+            let pixels = buf.iter().map(|&value| colormap(value as _)).collect();
+            Ok(ColorImage { size, pixels })
+        }
+        TensorData::U64(buf) => {
+            let pixels = buf.iter().map(|&value| colormap(value as _)).collect();
+            Ok(ColorImage { size, pixels })
+        }
 
+        TensorData::I8(buf) => {
+            let pixels = buf.iter().map(|&value| colormap(value as _)).collect();
             Ok(ColorImage { size, pixels })
         }
-        _ => {
-            color_tensor_as_color_image(tensor) // TODO(emilk): support more depth dtypes
+        TensorData::I16(buf) => {
+            let pixels = buf.iter().map(|&value| colormap(value as _)).collect();
+            Ok(ColorImage { size, pixels })
+        }
+        TensorData::I32(buf) => {
+            let pixels = buf.iter().map(|&value| colormap(value as _)).collect();
+            Ok(ColorImage { size, pixels })
+        }
+        TensorData::I64(buf) => {
+            let pixels = buf.iter().map(|&value| colormap(value as _)).collect();
+            Ok(ColorImage { size, pixels })
+        }
+
+        TensorData::F32(buf) => {
+            let pixels = buf.iter().map(|&value| colormap(value as _)).collect();
+            Ok(ColorImage { size, pixels })
+        }
+        TensorData::F64(buf) => {
+            let pixels = buf.iter().map(|&value| colormap(value)).collect();
+            Ok(ColorImage { size, pixels })
+        }
+
+        TensorData::JPEG(_) => {
+            anyhow::bail!("Cannot apply colormap to JPEG image")
         }
     }
 }

From df0c0314a5bb4ccc82cd38fc1b3b32d8247fb7ce Mon Sep 17 00:00:00 2001
From: Emil Ernerfeldt <emil.ernerfeldt@gmail.com>
Date: Thu, 23 Mar 2023 14:54:27 +0100
Subject: [PATCH 5/5] Add a comment

---
 crates/re_viewer/src/misc/caches/mod.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/crates/re_viewer/src/misc/caches/mod.rs b/crates/re_viewer/src/misc/caches/mod.rs
index f0363a68b5d8..db61a56109e2 100644
--- a/crates/re_viewer/src/misc/caches/mod.rs
+++ b/crates/re_viewer/src/misc/caches/mod.rs
@@ -53,6 +53,7 @@ impl Caches {
 }
 
 pub struct TensorStats {
+    /// This will currently only be `None` for jpeg-encoded tensors.
     pub range: Option<(f64, f64)>,
 }