From 953c12bc5f658cbdcc5879b2c4a7671f4185cabc Mon Sep 17 00:00:00 2001
From: Jeremy Leibs <jeremy@rerun.io>
Date: Wed, 29 Nov 2023 15:21:03 +0100
Subject: [PATCH] Bugfix: show labels on segmentation images with trivial
 dimensions (#4368)

### What
- Resolves: https://github.com/rerun-io/rerun/issues/3824

### Checklist
* [x] I have read and agree to [Contributor
Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and
the [Code of
Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md)
* [x] I've included a screenshot or gif (if applicable)
* [x] I have tested [app.rerun.io](https://app.rerun.io/pr/4368) (if
applicable)
* [x] The PR title and labels are set such as to maximize their
usefulness for the next release's CHANGELOG

- [PR Build Summary](https://build.rerun.io/pr/4368)
- [Docs
preview](https://rerun.io/preview/dadad7fa99befa4d8b16919d8c9582868d1ccb86/docs)
<!--DOCS-PREVIEW-->
- [Examples
preview](https://rerun.io/preview/dadad7fa99befa4d8b16919d8c9582868d1ccb86/examples)
<!--EXAMPLES-PREVIEW-->
- [Recent benchmark results](https://build.rerun.io/graphs/crates.html)
- [Wasm size tracking](https://build.rerun.io/graphs/sizes.html)
---
 crates/re_data_ui/src/image.rs                   | 3 ++-
 crates/re_types/src/datatypes/tensor_data_ext.rs | 4 ++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/crates/re_data_ui/src/image.rs b/crates/re_data_ui/src/image.rs
index 4dd25e351a89..aac97e779128 100644
--- a/crates/re_data_ui/src/image.rs
+++ b/crates/re_data_ui/src/image.rs
@@ -664,7 +664,8 @@ fn tensor_pixel_value_ui(
         ui.label(format!("{x}, {y}"));
         ui.end_row();
 
-        if tensor.num_dim() == 2 {
+        // Check for annotations on any single-channel image
+        if let Some([_, _, 1]) = tensor.image_height_width_channels() {
             if let Some(raw_value) = tensor.get(&[y, x]) {
                 if let (TensorDataMeaning::ClassId, Some(u16_val)) =
                     (meaning, raw_value.try_as_u16())
diff --git a/crates/re_types/src/datatypes/tensor_data_ext.rs b/crates/re_types/src/datatypes/tensor_data_ext.rs
index 80b32eff3f68..c4db01cf74c3 100644
--- a/crates/re_types/src/datatypes/tensor_data_ext.rs
+++ b/crates/re_types/src/datatypes/tensor_data_ext.rs
@@ -80,8 +80,8 @@ impl TensorData {
                     2 => Some([shape_short[0].size, shape_short[1].size, 1]),
                     3 => {
                         let channels = shape_short[2].size;
-                        if matches!(channels, 3 | 4) {
-                            // rgb, rgba
+                        if matches!(channels, 1 | 3 | 4) {
+                            // mono, rgb, rgba
                             Some([shape_short[0].size, shape_short[1].size, channels])
                         } else {
                             None