fix image preview for images other than M x N x C and M x N

crates/re_data_ui/src/image.rs

@@ -550,17 +550,15 @@ fn tensor_pixel_value_ui(
         }
     });
 
-    let text = match tensor.num_dim() {
-        2 => tensor.get(&[y, x]).map(|v| format!("Val: {v}")),
-        3 => match tensor.shape()[2].size {
-            0 => Some("Cannot preview 0-size channel".to_owned()),
-            1 => tensor.get(&[y, x, 0]).map(|v| format!("Val: {v}")),
+    let text = if let Some([_, _, channel]) = tensor.image_height_width_channels() {
+        match channel {
+            1 => tensor.get_image(x, y, 0).map(|v| format!("Val: {v}")),
             3 => {
                 // TODO(jleibs): Track RGB ordering somehow -- don't just assume it
                 if let (Some(r), Some(g), Some(b)) = (
-                    tensor.get(&[y, x, 0]),
-                    tensor.get(&[y, x, 1]),
-                    tensor.get(&[y, x, 2]),
+                    tensor.get_image(x, y, 0),
+                    tensor.get_image(x, y, 1),
+                    tensor.get_image(x, y, 2),
                 ) {
                     match (r, g, b) {
                         (TensorElement::U8(r), TensorElement::U8(g), TensorElement::U8(b)) => {
@@ -575,10 +573,10 @@ fn tensor_pixel_value_ui(
             4 => {
                 // TODO(jleibs): Track RGB ordering somehow -- don't just assume it
                 if let (Some(r), Some(g), Some(b), Some(a)) = (
-                    tensor.get(&[y, x, 0]),
-                    tensor.get(&[y, x, 1]),
-                    tensor.get(&[y, x, 2]),
-                    tensor.get(&[y, x, 3]),
+                    tensor.get_image(x, y, 0),
+                    tensor.get_image(x, y, 1),
+                    tensor.get_image(x, y, 2),
+                    tensor.get_image(x, y, 3),
                 ) {
                     match (r, g, b, a) {
                         (
@@ -595,9 +593,13 @@ fn tensor_pixel_value_ui(
                     None
                 }
             }
-            channels => Some(format!("Cannot preview {channels}-channel image")),
-        },
-        dims => Some(format!("Cannot preview {dims}-dimensional image")),
+            channel => Some(format!("Cannot preview {channel}-size channel image")),
+        }
+    } else {
+        Some(format!(
+            "Cannot preview tensors with a shape of {:?}",
+            tensor.shape()
+        ))
     };
 
     if let Some(text) = text {

crates/re_log_types/src/component_types/tensor.rs

@@ -472,6 +472,38 @@ impl Tensor {
         self.meaning
     }
 
+    /// Query with x, y, channel indices.
+    ///
+    /// Allows to query values for any image like tensor even if it has more or less dimensions than 3.
+    /// (useful for sampling e.g. `N x M x C x 1` tensor which is a valid image)
+    #[inline]
+    pub fn get_image(&self, x: u64, y: u64, channel: u64) -> Option<TensorElement> {
+        match self.shape.len() {
+            1 => {
+                if channel != 0 {
+                    None
+                } else {
+                    self.get(&[x])
+                }
+            }
+            2 => {
+                if channel != 0 {
+                    None
+                } else {
+                    self.get(&[x, y])
+                }
+            }
+            3 => self.get(&[x, y, channel]),
+            4 => self.get(&[x, y, channel, 0]), // Optimization for common case, next case handles this too.
+            dim => self.get(
+                &[x, y, channel]
+                    .into_iter()
+                    .chain(std::iter::repeat(0).take(dim - 3))
+                    .collect::<Vec<u64>>(),
+            ),
+        }
+    }
+
     pub fn get(&self, index: &[u64]) -> Option<TensorElement> {
         let mut stride: usize = 1;
         let mut offset: usize = 0;