Make it possible to log segmentation images directly as floats (#4585)

### What
- Resolves: #3609

This appears to have already worked and was just being prevented by
obsolete checks. Also we weren't doing the user-space conversions in C++
or Rust anyways.

### 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 the web demo (if applicable):
* Using newly built examples:
[app.rerun.io](https://app.rerun.io/pr/4585/index.html)
* Using examples from latest `main` build:
[app.rerun.io](https://app.rerun.io/pr/4585/index.html?manifest_url=https://app.rerun.io/version/main/examples_manifest.json)
* Using full set of examples from `nightly` build:
[app.rerun.io](https://app.rerun.io/pr/4585/index.html?manifest_url=https://app.rerun.io/version/nightly/examples_manifest.json)
* [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/4585)
- [Docs
preview](https://rerun.io/preview/312d196768a5d04efc92e07fca391f53c141b0b4/docs)
<!--DOCS-PREVIEW-->
- [Examples
preview](https://rerun.io/preview/312d196768a5d04efc92e07fca391f53c141b0b4/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_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(
                 [

rerun_py/tests/unit/test_segmentation_image.py

@@ -6,7 +6,7 @@
 import pytest
 import rerun as rr
 import torch
-from rerun.datatypes import TensorBuffer, TensorBufferType, TensorData, TensorDataLike, TensorDimension
+from rerun.datatypes import TensorBuffer, TensorData, TensorDataLike, TensorDimension
 
 rng = np.random.default_rng(12345)
 RANDOM_IMAGE_SOURCE = rng.integers(0, 255, size=(10, 20))
@@ -74,14 +74,3 @@ def test_segmentation_image_shapes() -> None:
     for img in BAD_IMAGE_INPUTS:
         with pytest.raises(ValueError):
             rr.DepthImage(img)
-
-
-def test_segmentation_coercion() -> None:
-    # TODO(#3609): pass segmentation images unmolested to the viewer
-    seg_img = np.require(RANDOM_IMAGE_SOURCE, np.float32)
-
-    seg = rr.SegmentationImage(seg_img)
-
-    U16_TYPE_ID = list(f.name for f in TensorBufferType().storage_type).index("U16")
-
-    assert seg.data.as_arrow_array().storage.field(1)[0].type_code == U16_TYPE_ID