Remember to decode JPEGs

crates/re_log_types/src/component_types/tensor.rs

@@ -330,7 +330,9 @@ pub enum TensorDataMeaning {
     Depth,
 }
 
-/// A Multi-dimensional Tensor
+/// A Multi-dimensional Tensor.
+///
+/// All clones are shallow.
 ///
 /// ## Examples
 ///

crates/re_viewer/src/gpu_bridge/tensor_to_gpu.rs

@@ -13,7 +13,7 @@ use re_renderer::{
     RenderContext,
 };
 
-use crate::{gpu_bridge::get_or_create_texture, misc::caches::TensorStats};
+use crate::{gpu_bridge::get_or_create_texture, misc::caches::TensorStats, DecodedTensor};
 
 use super::try_get_or_create_texture;
 
@@ -28,7 +28,7 @@ use super::try_get_or_create_texture;
 pub fn tensor_to_gpu(
     render_ctx: &mut RenderContext,
     debug_name: &str,
-    tensor: &Tensor,
+    tensor: &DecodedTensor,
     tensor_stats: &TensorStats,
     annotations: &crate::ui::Annotations,
 ) -> anyhow::Result<ColormappedTexture> {
@@ -60,7 +60,7 @@ pub fn tensor_to_gpu(
 fn color_tensor_to_gpu(
     render_ctx: &mut RenderContext,
     debug_name: &str,
-    tensor: &Tensor,
+    tensor: &DecodedTensor,
     tensor_stats: &TensorStats,
 ) -> anyhow::Result<ColormappedTexture> {
     let texture_handle = try_get_or_create_texture(render_ctx, hash(tensor.id()), || {
@@ -132,7 +132,7 @@ fn color_tensor_to_gpu(
 fn class_id_tensor_to_gpu(
     render_ctx: &mut RenderContext,
     debug_name: &str,
-    tensor: &Tensor,
+    tensor: &DecodedTensor,
     tensor_stats: &TensorStats,
     annotations: &crate::ui::Annotations,
 ) -> anyhow::Result<ColormappedTexture> {
@@ -200,7 +200,7 @@ fn class_id_tensor_to_gpu(
 fn depth_tensor_to_gpu(
     render_ctx: &mut RenderContext,
     debug_name: &str,
-    tensor: &Tensor,
+    tensor: &DecodedTensor,
     tensor_stats: &TensorStats,
 ) -> anyhow::Result<ColormappedTexture> {
     let [_height, _width, depth] = height_width_depth(tensor)?;
@@ -223,7 +223,10 @@ fn depth_tensor_to_gpu(
     })
 }
 
-fn depth_tensor_range(tensor: &Tensor, tensor_stats: &TensorStats) -> anyhow::Result<(f64, f64)> {
+fn depth_tensor_range(
+    tensor: &DecodedTensor,
+    tensor_stats: &TensorStats,
+) -> anyhow::Result<(f64, f64)> {
     let range = tensor_stats.range.ok_or(anyhow::anyhow!(
         "Tensor has no range!? Was this compressed?"
     ))?;
@@ -255,7 +258,7 @@ fn depth_tensor_range(tensor: &Tensor, tensor_stats: &TensorStats) -> anyhow::Re
 /// Uses no `Unorm/Snorm` formats.
 fn general_texture_creation_desc_from_tensor<'a>(
     debug_name: &str,
-    tensor: &'a Tensor,
+    tensor: &'a DecodedTensor,
 ) -> anyhow::Result<Texture2DCreationDesc<'a>> {
     let [height, width, depth] = height_width_depth(tensor)?;
 
@@ -277,7 +280,7 @@ fn general_texture_creation_desc_from_tensor<'a>(
                 TensorData::F64(buf) => (narrow_f64_to_f32s(buf), TextureFormat::R32Float), // narrowing to f32!
 
                 TensorData::JPEG(_) => {
-                    anyhow::bail!("JPEGs should have been decoded at this point")
+                    unreachable!("DecodedTensor cannot contain a JPEG")
                 }
             }
         }
@@ -299,7 +302,7 @@ fn general_texture_creation_desc_from_tensor<'a>(
                 TensorData::F64(buf) => (narrow_f64_to_f32s(buf), TextureFormat::Rg32Float), // narrowing to f32!
 
                 TensorData::JPEG(_) => {
-                    anyhow::bail!("JPEGs should have been decoded at this point")
+                    unreachable!("DecodedTensor cannot contain a JPEG")
                 }
             }
         }
@@ -336,7 +339,7 @@ fn general_texture_creation_desc_from_tensor<'a>(
                 ),
 
                 TensorData::JPEG(_) => {
-                    anyhow::bail!("JPEGs should have been decoded at this point")
+                    unreachable!("DecodedTensor cannot contain a JPEG")
                 }
             }
         }
@@ -360,7 +363,7 @@ fn general_texture_creation_desc_from_tensor<'a>(
                 TensorData::F64(buf) => (narrow_f64_to_f32s(buf), TextureFormat::Rgba32Float), // narrowing to f32!
 
                 TensorData::JPEG(_) => {
-                    anyhow::bail!("JPEGs should have been decoded at this point")
+                    unreachable!("DecodedTensor cannot contain a JPEG")
                 }
             }
         }

crates/re_viewer/src/lib.rs

@@ -16,6 +16,8 @@ pub(crate) use misc::{mesh_loader, Item, TimeControl, TimeView, ViewerContext};
 use re_log_types::PythonVersion;
 pub(crate) use ui::{memory_panel, selection_panel, time_panel, UiVerbosity};
 
+pub(crate) use misc::caches::DecodedTensor;
+
 pub use app::{App, StartupOptions};
 pub use remote_viewer_app::RemoteViewerApp;
 

crates/re_viewer/src/misc/caches/mod.rs

@@ -3,6 +3,8 @@ mod tensor_decode_cache;
 
 use re_log_types::component_types::{self};
 
+pub use tensor_decode_cache::DecodedTensor;
+
 /// Does memoization of different things for the immediate mode UI.
 #[derive(Default)]
 pub struct Caches {

crates/re_viewer/src/misc/caches/tensor_decode_cache.rs

@@ -1,5 +1,45 @@
 use re_log_types::component_types::{Tensor, TensorDimension, TensorId};
 
+// ----------------------------------------------------------------------------
+
+/// A thin wrapper around a [`Tensor`] that is guaranteed to not be compressed (never a jpeg).
+///
+/// All clones are shallow, like for [`Tensor`].
+#[derive(Clone)]
+pub struct DecodedTensor(Tensor);
+
+impl DecodedTensor {
+    #[inline(always)]
+    pub fn inner(&self) -> &Tensor {
+        &self.0
+    }
+}
+
+impl AsRef<Tensor> for DecodedTensor {
+    #[inline(always)]
+    fn as_ref(&self) -> &Tensor {
+        &self.0
+    }
+}
+
+impl std::ops::Deref for DecodedTensor {
+    type Target = Tensor;
+
+    #[inline(always)]
+    fn deref(&self) -> &Tensor {
+        &self.0
+    }
+}
+
+impl std::borrow::Borrow<Tensor> for DecodedTensor {
+    #[inline(always)]
+    fn borrow(&self) -> &Tensor {
+        &self.0
+    }
+}
+
+// ----------------------------------------------------------------------------
+
 #[derive(thiserror::Error, Clone, Debug)]
 pub enum TensorDecodeError {
     // TODO(jleibs): It would be nice to just transparently wrap
@@ -18,10 +58,9 @@ pub enum TensorDecodeError {
     },
 }
 
-#[derive(Clone)]
-struct DecodedTensor {
+struct DecodedTensorResult {
     /// Cached `Result` from decoding the `Tensor`
-    tensor: Result<Tensor, TensorDecodeError>,
+    tensor: Result<DecodedTensor, TensorDecodeError>,
 
     /// Total memory used by this `Tensor`.
     memory_used: u64,
@@ -33,7 +72,7 @@ struct DecodedTensor {
 /// A cache of decoded [`Tensor`] entities, indexed by `TensorId`.
 #[derive(Default)]
 pub struct DecodeCache {
-    images: nohash_hasher::IntMap<TensorId, DecodedTensor>,
+    images: nohash_hasher::IntMap<TensorId, DecodedTensorResult>,
     memory_used: u64,
     generation: u64,
 }
@@ -48,7 +87,7 @@ impl DecodeCache {
     pub fn try_decode_tensor_if_necessary(
         &mut self,
         maybe_encoded_tensor: Tensor,
-    ) -> Result<Tensor, TensorDecodeError> {
+    ) -> Result<DecodedTensor, TensorDecodeError> {
         crate::profile_function!();
         match &maybe_encoded_tensor.data {
             re_log_types::component_types::TensorData::JPEG(buf) => {
@@ -67,7 +106,7 @@ impl DecodeCache {
                             Ok(img) => match Tensor::from_image(img) {
                                 Ok(tensor) => {
                                     if tensor.shape() == maybe_encoded_tensor.shape() {
-                                        Ok(tensor)
+                                        Ok(DecodedTensor(tensor))
                                     } else {
                                         Err(TensorDecodeError::InvalidMetaData {
                                             expected: maybe_encoded_tensor.shape().into(),
@@ -86,7 +125,7 @@ impl DecodeCache {
                         };
                         self.memory_used += memory_used;
                         let last_use_generation = 0;
-                        DecodedTensor {
+                        DecodedTensorResult {
                             tensor,
                             memory_used,
                             last_use_generation,
@@ -96,7 +135,7 @@ impl DecodeCache {
 
                 lookup.tensor.clone()
             }
-            _ => Ok(maybe_encoded_tensor),
+            _ => Ok(DecodedTensor(maybe_encoded_tensor)),
         }
     }
 

crates/re_viewer/src/ui/data_ui/image.rs

@@ -11,6 +11,7 @@ use re_ui::ReUi;
 use crate::{
     misc::{caches::TensorStats, ViewerContext},
     ui::annotations::AnnotationMap,
+    DecodedTensor,
 };
 
 use super::{EntityDataUi, UiVerbosity};
@@ -54,7 +55,7 @@ fn tensor_ui(
     entity_path: &re_data_store::EntityPath,
     query: &re_arrow_store::LatestAtQuery,
     _encoded_tensor: &Tensor,
-    tensor: &Tensor,
+    tensor: &DecodedTensor,
 ) {
     // See if we can convert the tensor to a GPU texture.
     // Even if not, we will show info about the tensor.
@@ -328,7 +329,7 @@ fn show_zoomed_image_region_tooltip(
     render_ctx: &mut re_renderer::RenderContext,
     parent_ui: &mut egui::Ui,
     response: egui::Response,
-    tensor: &Tensor,
+    tensor: &DecodedTensor,
     tensor_stats: &TensorStats,
     annotations: &crate::ui::Annotations,
     meter: Option<f32>,
@@ -408,7 +409,7 @@ pub fn show_zoomed_image_region_area_outline(
 pub fn show_zoomed_image_region(
     render_ctx: &mut re_renderer::RenderContext,
     ui: &mut egui::Ui,
-    tensor: &Tensor,
+    tensor: &DecodedTensor,
     tensor_stats: &TensorStats,
     annotations: &crate::ui::Annotations,
     meter: Option<f32>,
@@ -434,7 +435,7 @@ pub fn show_zoomed_image_region(
 fn try_show_zoomed_image_region(
     render_ctx: &mut re_renderer::RenderContext,
     ui: &mut egui::Ui,
-    tensor: &Tensor,
+    tensor: &DecodedTensor,
     tensor_stats: &TensorStats,
     annotations: &crate::ui::Annotations,
     meter: Option<f32>,

crates/re_viewer/src/ui/view_spatial/scene/mod.rs

@@ -1,22 +1,25 @@
 use std::sync::Arc;
 
 use ahash::HashMap;
+
 use re_data_store::{EntityPath, InstancePathHash};
 use re_log_types::{
-    component_types::{ClassId, InstanceKey, KeypointId, Tensor},
+    component_types::{ClassId, InstanceKey, KeypointId},
     MeshId,
 };
 use re_renderer::{Color32, OutlineMaskPreference, Size};
 
-use super::{SpaceCamera3D, SpatialNavigationMode};
 use crate::{
     misc::{mesh_loader::LoadedMesh, SpaceViewHighlights, TransformCache, ViewerContext},
     ui::{
         annotations::{auto_color, AnnotationMap},
         Annotations, SceneQuery,
     },
+    DecodedTensor,
 };
 
+use super::{SpaceCamera3D, SpatialNavigationMode};
+
 mod picking;
 mod primitives;
 mod scene_part;
@@ -59,7 +62,7 @@ pub struct Image {
     /// Path to the image (note image instance ids would refer to pixels!)
     pub ent_path: EntityPath,
 
-    pub tensor: Tensor,
+    pub tensor: DecodedTensor,
 
     /// If this is a depth map, how long is a meter?
     ///