Simplify ImageCache

crates/re_log_types/src/data.rs

@@ -255,3 +255,21 @@ impl TensorElement {
         }
     }
 }
+
+impl std::fmt::Display for TensorElement {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            TensorElement::U8(elem) => std::fmt::Display::fmt(elem, f),
+            TensorElement::U16(elem) => std::fmt::Display::fmt(elem, f),
+            TensorElement::U32(elem) => std::fmt::Display::fmt(elem, f),
+            TensorElement::U64(elem) => std::fmt::Display::fmt(elem, f),
+            TensorElement::I8(elem) => std::fmt::Display::fmt(elem, f),
+            TensorElement::I16(elem) => std::fmt::Display::fmt(elem, f),
+            TensorElement::I32(elem) => std::fmt::Display::fmt(elem, f),
+            TensorElement::I64(elem) => std::fmt::Display::fmt(elem, f),
+            TensorElement::F16(elem) => std::fmt::Display::fmt(elem, f),
+            TensorElement::F32(elem) => std::fmt::Display::fmt(elem, f),
+            TensorElement::F64(elem) => std::fmt::Display::fmt(elem, f),
+        }
+    }
+}

crates/re_renderer/src/resource_managers/texture_manager.rs

@@ -1,5 +1,7 @@
 use std::{num::NonZeroU32, sync::Arc};
 
+use ahash::{HashMap, HashSet};
+
 use crate::{
     wgpu_resources::{GpuTexture, GpuTexturePool, TextureDesc},
     DebugLabel,
@@ -62,6 +64,16 @@ pub struct TextureManager2D {
     // For convenience to reduce amount of times we need to pass them around
     device: Arc<wgpu::Device>,
     queue: Arc<wgpu::Queue>,
+
+    // Cache the texture using a user-provided u64 id. This is expected
+    // to be derived from the `TensorId` which isn't available here for
+    // dependency reasons.
+    // TODO(jleibs): Introduce a proper key here.
+    //
+    // Any texture which wasn't accessed on the previous frame
+    // is ejected from the cache during [`begin_frame`].
+    texture_cache: HashMap<u64, GpuTexture2DHandle>,
+    accessed_textures: HashSet<u64>,
 }
 
 impl TextureManager2D {
@@ -87,11 +99,13 @@ impl TextureManager2D {
             white_texture,
             device,
             queue,
+            texture_cache: Default::default(),
+            accessed_textures: Default::default(),
         }
     }
 
     /// Creates a new 2D texture resource and schedules data upload to the GPU.
-    #[allow(clippy::unused_self)]
+    /// TODO(jleibs): All usages of this should be be replaced with `get_or_create`, which is strictly preferable
     pub fn create(
         &mut self,
         texture_pool: &mut GpuTexturePool,
@@ -116,6 +130,21 @@ impl TextureManager2D {
         Self::create_and_upload_texture(&self.device, &self.queue, texture_pool, creation_desc)
     }
 
+    /// Creates a new 2D texture resource and schedules data upload to the GPU if a texture
+    /// wasn't already created using the same key.
+    pub fn get_or_create(
+        &mut self,
+        key: u64,
+        texture_pool: &mut GpuTexturePool,
+        creation_desc: &Texture2DCreationDesc<'_>,
+    ) -> GpuTexture2DHandle {
+        let texture_handle = self.texture_cache.entry(key).or_insert_with(|| {
+            Self::create_and_upload_texture(&self.device, &self.queue, texture_pool, creation_desc)
+        });
+        self.accessed_textures.insert(key);
+        texture_handle.clone()
+    }
+
     /// Returns a single pixel white pixel.
     pub fn white_texture_handle(&self) -> &GpuTexture2DHandle {
         &self.white_texture
@@ -193,9 +222,10 @@ impl TextureManager2D {
             .map(|h| h.clone())
     }
 
-    #[allow(clippy::unused_self)]
     pub(crate) fn begin_frame(&mut self, _frame_index: u64) {
-        // no-op.
-        // In the future we might add handling of background processing or introduce frame-lived textures.
+        // Drop any textures that weren't accessed in the last frame
+        self.texture_cache
+            .retain(|k, _| self.accessed_textures.contains(k));
+        self.accessed_textures.clear();
     }
 }

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

@@ -3,7 +3,7 @@ mod tensor_decode_cache;
 mod tensor_image_cache;
 
 use re_log_types::component_types::{self, TensorTrait};
-pub use tensor_image_cache::{AsDynamicImage, TensorImageView};
+pub use tensor_image_cache::ColoredTensorView;
 
 /// Does memoization of different things for the immediate mode UI.
 #[derive(Default)]