Test and handle all tensor dtypes as images (#1840)

* Support i64 and u64 tensors

* api_demo: log all image types

* Don't even print out the contents of a tensor

* Handle unfilterable float textures

* fix typo

* py-format

* Simplify is_float_filterable

* Add a helper function pad_and_narrow_and_cast

* Reuse existing image

* Exclude image_tensors demo from default api_demo

* Still run all api demos in e2e test

* pyformat

crates/re_log_types/src/component_types/tensor.rs

@@ -142,7 +142,7 @@ impl ArrowDeserialize for TensorId {
 ///     ),
 /// );
 /// ```
-#[derive(Clone, Debug, PartialEq, ArrowField, ArrowSerialize, ArrowDeserialize)]
+#[derive(Clone, PartialEq, ArrowField, ArrowSerialize, ArrowDeserialize)]
 #[arrow_field(type = "dense")]
 pub enum TensorData {
     U8(BinaryBuffer),
@@ -198,6 +198,24 @@ impl TensorData {
     }
 }
 
+impl std::fmt::Debug for TensorData {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::U8(_) => write!(f, "U8({} bytes)", self.size_in_bytes()),
+            Self::U16(_) => write!(f, "U16({} bytes)", self.size_in_bytes()),
+            Self::U32(_) => write!(f, "U32({} bytes)", self.size_in_bytes()),
+            Self::U64(_) => write!(f, "U64({} bytes)", self.size_in_bytes()),
+            Self::I8(_) => write!(f, "I8({} bytes)", self.size_in_bytes()),
+            Self::I16(_) => write!(f, "I16({} bytes)", self.size_in_bytes()),
+            Self::I32(_) => write!(f, "I32({} bytes)", self.size_in_bytes()),
+            Self::I64(_) => write!(f, "I64({} bytes)", self.size_in_bytes()),
+            Self::F32(_) => write!(f, "F32({} bytes)", self.size_in_bytes()),
+            Self::F64(_) => write!(f, "F64({} bytes)", self.size_in_bytes()),
+            Self::JPEG(_) => write!(f, "JPEG({} bytes)", self.size_in_bytes()),
+        }
+    }
+}
+
 /// Flattened `Tensor` data payload
 ///
 /// ## Examples

crates/re_renderer/shader/rectangle.wgsl

@@ -6,9 +6,10 @@
 // Keep in sync with mirror in rectangle.rs
 
 // Which texture to read from?
-const SAMPLE_TYPE_FLOAT = 1u;
-const SAMPLE_TYPE_SINT  = 2u;
-const SAMPLE_TYPE_UINT  = 3u;
+const SAMPLE_TYPE_FLOAT_FILTER   = 1u;
+const SAMPLE_TYPE_FLOAT_NOFILTER = 2u;
+const SAMPLE_TYPE_SINT_NOFILTER  = 3u;
+const SAMPLE_TYPE_UINT_NOFILTER  = 4u;
 
 // How do we do colormapping?
 const COLOR_MAPPER_OFF      = 1u;
@@ -67,6 +68,9 @@ var texture_uint: texture_2d<u32>;
 @group(1) @binding(5)
 var colormap_texture: texture_2d<f32>;
 
+@group(1) @binding(6)
+var texture_float_filterable: texture_2d<f32>;
+
 
 struct VertexOut {
     @builtin(position) position: Vec4,
@@ -90,12 +94,16 @@ fn vs_main(@builtin(vertex_index) v_idx: u32) -> VertexOut {
 fn fs_main(in: VertexOut) -> @location(0) Vec4 {
     // Sample the main texture:
     var sampled_value: Vec4;
-    if rect_info.sample_type == SAMPLE_TYPE_FLOAT {
-        sampled_value = textureSampleLevel(texture_float, texture_sampler, in.texcoord, 0.0); // TODO(emilk): support mipmaps
-    } else if rect_info.sample_type == SAMPLE_TYPE_SINT {
+    if rect_info.sample_type == SAMPLE_TYPE_FLOAT_FILTER {
+        // TODO(emilk): support mipmaps
+        sampled_value = textureSampleLevel(texture_float_filterable, texture_sampler, in.texcoord, 0.0);
+    } else if rect_info.sample_type == SAMPLE_TYPE_FLOAT_NOFILTER {
+        let icoords = IVec2(in.texcoord * Vec2(textureDimensions(texture_float).xy));
+        sampled_value = Vec4(textureLoad(texture_float, icoords, 0));
+    } else if rect_info.sample_type == SAMPLE_TYPE_SINT_NOFILTER {
         let icoords = IVec2(in.texcoord * Vec2(textureDimensions(texture_sint).xy));
         sampled_value = Vec4(textureLoad(texture_sint, icoords, 0));
-    } else if rect_info.sample_type == SAMPLE_TYPE_UINT {
+    } else if rect_info.sample_type == SAMPLE_TYPE_UINT_NOFILTER {
         let icoords = IVec2(in.texcoord * Vec2(textureDimensions(texture_uint).xy));
         sampled_value = Vec4(textureLoad(texture_uint, icoords, 0));
     } else {

crates/re_renderer/src/renderer/rectangles.rs

@@ -182,9 +182,10 @@ mod gpu_data {
     // Keep in sync with mirror in rectangle.wgsl
 
     // Which texture to read from?
-    const SAMPLE_TYPE_FLOAT: u32 = 1;
-    const SAMPLE_TYPE_SINT: u32 = 2;
-    const SAMPLE_TYPE_UINT: u32 = 3;
+    const SAMPLE_TYPE_FLOAT_FILTER: u32 = 1;
+    const SAMPLE_TYPE_FLOAT_NOFILTER: u32 = 2;
+    const SAMPLE_TYPE_SINT_NOFILTER: u32 = 3;
+    const SAMPLE_TYPE_UINT_NOFILTER: u32 = 4;
 
     // How do we do colormapping?
     const COLOR_MAPPER_OFF: u32 = 1;
@@ -232,12 +233,18 @@ mod gpu_data {
             } = &rectangle.colormapped_texture;
 
             let sample_type = match texture_info.sample_type {
-                wgpu::TextureSampleType::Float { .. } => SAMPLE_TYPE_FLOAT,
+                wgpu::TextureSampleType::Float { .. } => {
+                    if super::is_float_filterable(texture_format) {
+                        SAMPLE_TYPE_FLOAT_FILTER
+                    } else {
+                        SAMPLE_TYPE_FLOAT_NOFILTER
+                    }
+                }
                 wgpu::TextureSampleType::Depth => {
                     return Err(RectangleError::DepthTexturesNotSupported);
                 }
-                wgpu::TextureSampleType::Sint => SAMPLE_TYPE_SINT,
-                wgpu::TextureSampleType::Uint => SAMPLE_TYPE_UINT,
+                wgpu::TextureSampleType::Sint => SAMPLE_TYPE_SINT_NOFILTER,
+                wgpu::TextureSampleType::Uint => SAMPLE_TYPE_UINT_NOFILTER,
             };
 
             let mut colormap_function = 0;
@@ -378,14 +385,19 @@ impl RectangleDrawData {
                 ));
             }
 
-            // We set up three texture sources, then instruct the shader to read from at most one of them.
-            let mut texture_float = ctx.texture_manager_2d.zeroed_texture_float().handle;
+            // We set up several texture sources, then instruct the shader to read from at most one of them.
+            let mut texture_float_filterable = ctx.texture_manager_2d.zeroed_texture_float().handle;
+            let mut texture_float_nofilter = ctx.texture_manager_2d.zeroed_texture_float().handle;
             let mut texture_sint = ctx.texture_manager_2d.zeroed_texture_sint().handle;
             let mut texture_uint = ctx.texture_manager_2d.zeroed_texture_uint().handle;
 
             match texture_description.sample_type {
                 wgpu::TextureSampleType::Float { .. } => {
-                    texture_float = texture.handle;
+                    if is_float_filterable(&texture_format) {
+                        texture_float_filterable = texture.handle;
+                    } else {
+                        texture_float_nofilter = texture.handle;
+                    }
                 }
                 wgpu::TextureSampleType::Depth => {
                     return Err(RectangleError::DepthTexturesNotSupported);
@@ -422,10 +434,11 @@ impl RectangleDrawData {
                         entries: smallvec![
                             uniform_buffer,
                             BindGroupEntry::Sampler(sampler),
-                            BindGroupEntry::DefaultTextureView(texture_float),
+                            BindGroupEntry::DefaultTextureView(texture_float_nofilter),
                             BindGroupEntry::DefaultTextureView(texture_sint),
                             BindGroupEntry::DefaultTextureView(texture_uint),
                             BindGroupEntry::DefaultTextureView(colormap_texture),
+                            BindGroupEntry::DefaultTextureView(texture_float_filterable),
                         ],
                         layout: rectangle_renderer.bind_group_layout,
                     },
@@ -483,12 +496,12 @@ impl Renderer for RectangleRenderer {
                         ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                         count: None,
                     },
-                    // float texture:
+                    // float textures without filtering (e.g. R32Float):
                     wgpu::BindGroupLayoutEntry {
                         binding: 2,
                         visibility: wgpu::ShaderStages::FRAGMENT,
                         ty: wgpu::BindingType::Texture {
-                            sample_type: wgpu::TextureSampleType::Float { filterable: true },
+                            sample_type: wgpu::TextureSampleType::Float { filterable: false },
                             view_dimension: wgpu::TextureViewDimension::D2,
                             multisampled: false,
                         },
@@ -527,6 +540,17 @@ impl Renderer for RectangleRenderer {
                         },
                         count: None,
                     },
+                    // float textures with filtering (e.g. Rgba8UnormSrgb):
+                    wgpu::BindGroupLayoutEntry {
+                        binding: 6,
+                        visibility: wgpu::ShaderStages::FRAGMENT,
+                        ty: wgpu::BindingType::Texture {
+                            sample_type: wgpu::TextureSampleType::Float { filterable: true },
+                            view_dimension: wgpu::TextureViewDimension::D2,
+                            multisampled: false,
+                        },
+                        count: None,
+                    },
                 ],
             },
         );
@@ -653,3 +677,11 @@ impl Renderer for RectangleRenderer {
         ]
     }
 }
+
+fn is_float_filterable(format: &wgpu::TextureFormat) -> bool {
+    format
+        .describe()
+        .guaranteed_format_features
+        .flags
+        .contains(wgpu::TextureFormatFeatureFlags::FILTERABLE)
+}

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

@@ -328,8 +328,12 @@ fn color_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> {
             Ok(ColorImage { size, pixels })
         }
 
-        (_depth, dtype) => {
-            anyhow::bail!("Don't know how to turn a tensor of shape={:?} and dtype={dtype:?} into a color image", tensor.shape)
+        (_depth, _dtype) => {
+            anyhow::bail!(
+                "Don't know how to turn a tensor of shape={:?} and dtype={:?} into a color image",
+                tensor.shape,
+                tensor.dtype()
+            )
         }
     }
 }