From 42bc02d20719cd0f47abe656f33ecfd84dd3109b Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 20 Dec 2023 17:01:34 +0100 Subject: [PATCH] Make it possible to log segmentation images directly as floats (#4585) ### What - Resolves: https://github.com/rerun-io/rerun/issues/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) - [Examples preview](https://rerun.io/preview/312d196768a5d04efc92e07fca391f53c141b0b4/examples) - [Recent benchmark results](https://build.rerun.io/graphs/crates.html) - [Wasm size tracking](https://build.rerun.io/graphs/sizes.html) --- .../rerun/archetypes/segmentation_image.fbs | 5 ++++- .../re_types/src/archetypes/segmentation_image.rs | 5 ++++- .../src/gpu_bridge/tensor_to_gpu.rs | 11 ++++------- .../types/archetypes/segmentation_image.md | 5 ++++- .../src/rerun/archetypes/segmentation_image.hpp | 5 ++++- .../rerun/archetypes/segmentation_image.py | 5 ++++- .../rerun/archetypes/segmentation_image_ext.py | 11 ----------- rerun_py/tests/unit/test_segmentation_image.py | 13 +------------ 8 files changed, 25 insertions(+), 35 deletions(-) diff --git a/crates/re_types/definitions/rerun/archetypes/segmentation_image.fbs b/crates/re_types/definitions/rerun/archetypes/segmentation_image.fbs index 6be79daada3c..6338b4310324 100644 --- a/crates/re_types/definitions/rerun/archetypes/segmentation_image.fbs +++ b/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. diff --git a/crates/re_types/src/archetypes/segmentation_image.rs b/crates/re_types/src/archetypes/segmentation_image.rs index e4a5afa176a0..edd7c7c7e8ca 100644 --- a/crates/re_types/src/archetypes/segmentation_image.rs +++ b/crates/re_types/src/archetypes/segmentation_image.rs @@ -24,7 +24,10 @@ use ::re_types_core::{DeserializationError, DeserializationResult}; /// **Archetype**: 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. diff --git a/crates/re_viewer_context/src/gpu_bridge/tensor_to_gpu.rs b/crates/re_viewer_context/src/gpu_bridge/tensor_to_gpu.rs index 5b56a5b46d93..92b1a0168a95 100644 --- a/crates/re_viewer_context/src/gpu_bridge/tensor_to_gpu.rs +++ b/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. diff --git a/docs/content/reference/types/archetypes/segmentation_image.md b/docs/content/reference/types/archetypes/segmentation_image.md index 213ed02a125f..f093fcb8d15d 100644 --- a/docs/content/reference/types/archetypes/segmentation_image.md +++ b/docs/content/reference/types/archetypes/segmentation_image.md @@ -5,7 +5,10 @@ title: "SegmentationImage" 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. diff --git a/rerun_cpp/src/rerun/archetypes/segmentation_image.hpp b/rerun_cpp/src/rerun/archetypes/segmentation_image.hpp index ee613b1c0ea1..83709b9c104c 100644 --- a/rerun_cpp/src/rerun/archetypes/segmentation_image.hpp +++ b/rerun_cpp/src/rerun/archetypes/segmentation_image.hpp @@ -20,7 +20,10 @@ namespace rerun::archetypes { /// **Archetype**: 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. diff --git a/rerun_py/rerun_sdk/rerun/archetypes/segmentation_image.py b/rerun_py/rerun_sdk/rerun/archetypes/segmentation_image.py index eed5176fbdb7..a93dd55ef4f2 100644 --- a/rerun_py/rerun_sdk/rerun/archetypes/segmentation_image.py +++ b/rerun_py/rerun_sdk/rerun/archetypes/segmentation_image.py @@ -23,7 +23,10 @@ class SegmentationImage(SegmentationImageExt, Archetype): **Archetype**: 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. diff --git a/rerun_py/rerun_sdk/rerun/archetypes/segmentation_image_ext.py b/rerun_py/rerun_sdk/rerun/archetypes/segmentation_image_ext.py index a69b92b2e07c..c69494f44ea4 100644 --- a/rerun_py/rerun_sdk/rerun/archetypes/segmentation_image_ext.py +++ b/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( [ diff --git a/rerun_py/tests/unit/test_segmentation_image.py b/rerun_py/tests/unit/test_segmentation_image.py index f1812d46052d..186ff64902e3 100644 --- a/rerun_py/tests/unit/test_segmentation_image.py +++ b/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