Port SegmentationImage to the new image archetype style

CHANGELOG.md

@@ -7,7 +7,7 @@
 * 3D transform APIs: Previously, the transform component was represented as one of several variants (an Arrow union, `enum` in Rust) depending on how the transform was expressed. Instead, there are now several components for translation/scale/rotation/matrices that can live side-by-side in the 3D transform archetype.
 * Python: `NV12/YUY2` are now logged with the new `ImageChromaDownsampled`
 * [`ImageEncoded`](https://rerun.io/docs/reference/types/archetypes/image_encoded?speculative-link):s `format` parameter has been replaced with `media_type` (MIME)
-* [`DepthImage`](https://rerun.io/docs/reference/types/archetypes/depth_image) is no longer encoded as a tensor, and expects its shape in `[width, height]` order
+* [`DepthImage`](https://rerun.io/docs/reference/types/archetypes/depth_image) and [`SegmentationImage`](https://rerun.io/docs/reference/types/archetypes/segmentation_image) are no longer encoded as a tensors, and expects its shape in `[width, height]` order
 
 🧳 Migration guide: http://rerun.io/docs/reference/migration/migration-0-18?speculative-link
 

crates/store/re_types/definitions/rerun/archetypes/depth_image.fbs

@@ -1,10 +1,9 @@
 namespace rerun.archetypes;
 
 
-/// A depth image.
+/// A depth image, i.e. as captured by a depth camera.
 ///
-/// The shape of the [components.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 depth value in units specified by [components.DepthMeter].
 ///
 /// \cpp Since the underlying `rerun::datatypes::TensorData` uses `rerun::Collection` internally,
 /// \cpp data can be passed in without a copy from raw pointers or by reference from `std::vector`/`std::array`/c-arrays.

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

@@ -3,15 +3,11 @@ namespace rerun.archetypes;
 
 /// An image made up of integer [components.ClassId]s.
 ///
-/// The shape of the [components.TensorData] must be mappable to an `HxW` tensor.
 /// Each pixel corresponds to a [components.ClassId] 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.
-///
 /// See also [archetypes.AnnotationContext] to associate each class with a color and a label.
 ///
 /// \cpp Since the underlying `rerun::datatypes::TensorData` uses `rerun::Collection` internally,
@@ -27,8 +23,14 @@ table SegmentationImage (
 ) {
   // --- Required ---
 
-  /// The image data. Should always be a 2-dimensional tensor.
-  data: rerun.components.TensorData ("attr.rerun.component_required", order: 1000);
+  /// The raw image data.
+  data: rerun.components.Blob ("attr.rerun.component_required", order: 1000);
+
+  /// The size of the image.
+  resolution: rerun.components.Resolution2D ("attr.rerun.component_required", order: 1500);
+
+  /// The data type of the segmentation image data (U16, U32, …).
+  data_type: rerun.components.ChannelDataType ("attr.rerun.component_required", order: 2000);
 
   // --- Optional ---
 

crates/store/re_types/definitions/rerun/components/color_model.fbs

@@ -5,7 +5,9 @@ namespace rerun.components;
 /// Specified what color components are present in an [archetypes.Image].
 ///
 /// This combined with [components.ChannelDataType] determines the pixel format of an image.
-enum ColorModel: byte {
+enum ColorModel: byte (
+    "attr.docs.unreleased"
+) {
     /// Grayscale luminance intencity/brightness/value, sometimes called `Y`
     L (default),
 

crates/store/re_types/definitions/rerun/components/resolution2d.fbs

@@ -2,6 +2,7 @@ namespace rerun.components;
 
 /// The width and height of a 2D image.
 struct Resolution2D (
+  "attr.docs.unreleased",
   "attr.python.aliases": "npt.NDArray[np.int], Sequence[int], Tuple[int, int]",
   "attr.python.array_aliases": "npt.NDArray[np.int], Sequence[int]",
   "attr.rust.derive": "Default, Copy, PartialEq, Eq, Hash, bytemuck::Pod, bytemuck::Zeroable",

crates/store/re_types/src/archetypes/segmentation_image_ext.rs

@@ -1,5 +1,6 @@
 use crate::{
-    datatypes::TensorData,
+    components::{ChannelDataType, Resolution2D},
+    datatypes::{Blob, TensorBuffer, TensorData},
     image::{find_non_empty_dim_indices, ImageConstructionError},
 };
 
@@ -16,62 +17,47 @@ impl SegmentationImage {
     where
         <T as TryInto<TensorData>>::Error: std::error::Error,
     {
-        let mut data: TensorData = data
+        let tensor_data: TensorData = data
             .try_into()
             .map_err(ImageConstructionError::TensorDataConversion)?;
 
-        let non_empty_dim_inds = find_non_empty_dim_indices(&data.shape);
+        let non_empty_dim_inds = find_non_empty_dim_indices(&tensor_data.shape);
 
-        match non_empty_dim_inds.len() {
-            2 => {
-                assign_if_none(&mut data.shape[non_empty_dim_inds[0]].name, "height");
-                assign_if_none(&mut data.shape[non_empty_dim_inds[1]].name, "width");
+        if non_empty_dim_inds.len() != 2 {
+            return Err(ImageConstructionError::BadImageShape(tensor_data.shape));
+        }
+
+        let (blob, data_type) = match tensor_data.buffer {
+            TensorBuffer::U8(buffer) => (Blob(buffer), ChannelDataType::U8),
+            TensorBuffer::U16(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::U16),
+            TensorBuffer::U32(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::U32),
+            TensorBuffer::U64(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::U64),
+            TensorBuffer::I8(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::I8),
+            TensorBuffer::I16(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::I16),
+            TensorBuffer::I32(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::I32),
+            TensorBuffer::I64(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::I64),
+            TensorBuffer::F16(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::F16),
+            TensorBuffer::F32(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::F32),
+            TensorBuffer::F64(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::F64),
+            TensorBuffer::Nv12(_) | TensorBuffer::Yuy2(_) => {
+                return Err(ImageConstructionError::ChromaDownsamplingNotSupported);
             }
-            _ => return Err(ImageConstructionError::BadImageShape(data.shape)),
         };
 
+        let (height, width) = (
+            &tensor_data.shape[non_empty_dim_inds[0]],
+            &tensor_data.shape[non_empty_dim_inds[1]],
+        );
+        let height = height.size as u32;
+        let width = width.size as u32;
+        let resolution = Resolution2D::from([width, height]);
+
         Ok(Self {
-            data: data.into(),
+            data: blob.into(),
+            resolution,
+            data_type,
             draw_order: None,
             opacity: None,
         })
     }
 }
-
-fn assign_if_none(name: &mut Option<::re_types_core::ArrowString>, new_name: &str) {
-    if name.is_none() {
-        *name = Some(new_name.into());
-    }
-}
-
-// ----------------------------------------------------------------------------
-// Make it possible to create an ArrayView directly from an Image.
-
-macro_rules! forward_array_views {
-    ($type:ty, $alias:ty) => {
-        impl<'a> TryFrom<&'a $alias> for ::ndarray::ArrayViewD<'a, $type> {
-            type Error = crate::tensor_data::TensorCastError;
-
-            #[inline]
-            fn try_from(value: &'a $alias) -> Result<Self, Self::Error> {
-                (&value.data.0).try_into()
-            }
-        }
-    };
-}
-
-forward_array_views!(u8, SegmentationImage);
-forward_array_views!(u16, SegmentationImage);
-forward_array_views!(u32, SegmentationImage);
-forward_array_views!(u64, SegmentationImage);
-
-forward_array_views!(i8, SegmentationImage);
-forward_array_views!(i16, SegmentationImage);
-forward_array_views!(i32, SegmentationImage);
-forward_array_views!(i64, SegmentationImage);
-
-forward_array_views!(half::f16, SegmentationImage);
-forward_array_views!(f32, SegmentationImage);
-forward_array_views!(f64, SegmentationImage);
-
-// ----------------------------------------------------------------------------