Apply color maps to all types of depth tensors (#1686)

* Refactor how tensor image cache turns tensors into images

* simplify the code

* Refactor how we color map depth images

* Apply colormaps to all types of depth images

* Add a comment

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 {

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()
     }
 }
 

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 {

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)>,
 }