Make it possible to log segmentation images directly as floats

crates/re_types/definitions/rerun/archetypes/segmentation_image.fbs

@@ -9,7 +9,10 @@ namespace rerun.archetypes;
 /// An image made up of integer class-ids.
 ///
 /// The shape of the `TensorData` must be mappable to an `HxW` tensor.
-/// Each pixel corresponds to a depth value in units specified by meter.
+/// Each pixel corresponds to a class-id that will be mapped to a color based on annotation context.
+///
+/// In the case of floating point images, the label will be looked up based on rounding to the nearest
+/// integer value.
 ///
 /// Leading and trailing unit-dimensions are ignored, so that
 /// `1x640x480x1` is treated as a `640x480` image.

crates/re_viewer_context/src/gpu_bridge/tensor_to_gpu.rs

@@ -210,17 +210,14 @@ pub fn class_id_tensor_to_gpu(
         "Cannot apply annotations to tensor of shape {:?}",
         tensor.shape
     );
-    anyhow::ensure!(
-        tensor.dtype().is_integer(),
-        "Only integer tensors can be annotated"
-    );
 
-    let (min, max) = tensor_stats
+    let (_, mut max) = tensor_stats
         .range
         .ok_or_else(|| anyhow::anyhow!("compressed_tensor!?"))?;
-    anyhow::ensure!(0.0 <= min, "Negative class id");
 
-    anyhow::ensure!(max <= 65535.0, "Too many class ids"); // we only support u8 and u16 tensors
+    // We only support u8 and u16 class ids.
+    // Any values greater than this will be unmapped in the segmentation image.
+    max = max.min(65535.0);
 
     // We pack the colormap into a 2D texture so we don't go over the max texture size.
     // We only support u8 and u16 class ids, so 256^2 is the biggest texture we need.

rerun_py/rerun_sdk/rerun/archetypes/segmentation_image_ext.py

@@ -6,8 +6,6 @@
 import pyarrow as pa
 
 from .._validators import find_non_empty_dim_indices
-from ..datatypes import TensorBufferType
-from ..datatypes.tensor_data_ext import _build_buffer_array
 from ..error_utils import _send_warning_or_raise, catch_and_log_exceptions
 
 if TYPE_CHECKING:
@@ -18,9 +16,6 @@
 class SegmentationImageExt:
     """Extension for [SegmentationImage][rerun.archetypes.SegmentationImage]."""
 
-    U8_TYPE_ID = list(f.name for f in TensorBufferType().storage_type).index("U8")
-    U16_TYPE_ID = list(f.name for f in TensorBufferType().storage_type).index("U16")
-
     @staticmethod
     @catch_and_log_exceptions("SegmentationImage converter")
     def data__field_converter_override(data: TensorDataArrayLike) -> TensorDataBatch:
@@ -71,12 +66,6 @@ def data__field_converter_override(data: TensorDataArrayLike) -> TensorDataBatch
 
         buffer = tensor_data_arrow.storage.field(1)
 
-        # The viewer only supports u8 and u16 segmentation images at the moment:
-        # TODO(#3609): handle this in the viewer instead
-        if buffer[0].type_code not in (SegmentationImageExt.U8_TYPE_ID, SegmentationImageExt.U16_TYPE_ID):
-            np_buffer = np.require(buffer[0].value.values.to_numpy(), np.uint16)
-            buffer = _build_buffer_array(np_buffer)
-
         return TensorDataBatch(
             pa.StructArray.from_arrays(
                 [