Introduce IDL-based Tensor/Image APIs for python.

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

@@ -15,6 +15,13 @@ namespace rerun.archetypes;
 ///
 /// The viewer has limited support for ignoring extra empty dimensions.
 ///
+/// \py Example
+/// \py -------
+/// \py
+/// \py ```python
+/// \py \include:../../../../../docs/code-examples/image_simple_v2.py
+/// \py ```
+///
 /// \rs ## Example
 /// \rs
 /// \rs ```ignore

crates/re_types/definitions/rerun/datatypes/tensor_data.fbs

@@ -19,7 +19,8 @@ namespace rerun.datatypes;
 /// which stores a contiguous array of typed values.
 table TensorData (
   order: 100,
-  "attr.rust.derive": "PartialEq"
+  "attr.python.array_aliases": "npt.ArrayLike",
+  "attr.rust.derive": "PartialEq,"
 ) {
   id: rerun.datatypes.TensorId (order: 100);
   shape: [rerun.datatypes.TensorDimension] (order: 200);

crates/re_types/definitions/rerun/datatypes/tensor_id.fbs

@@ -8,6 +8,7 @@ namespace rerun.datatypes;
 // A unique id per [`Tensor`].
 struct TensorId (
   order: 100,
+  "attr.python.aliases": "uuid.UUID",
   "attr.arrow.transparent",
   "attr.rust.derive": "Copy, Default, Eq, PartialEq"
 ) {

crates/re_types_builder/src/codegen/python.rs

@@ -281,6 +281,7 @@ fn quote_objects(
             import numpy as np
             import numpy.typing as npt
             import pyarrow as pa
+            import uuid
 
             from .._baseclasses import (
                 Archetype,
@@ -436,25 +437,23 @@ impl QuotedObject {
                 let (default_converter, converter_function) =
                     quote_field_converter_from_field(obj, objects, field);
 
-                let converter = if *kind == ObjectKind::Archetype {
+                let override_name = format!(
+                    "{}_{}_converter",
+                    name.to_lowercase(),
+                    field.name.to_lowercase()
+                );
+                let converter = if overrides.contains(&override_name) {
+                    format!("converter={override_name}")
+                } else if *kind == ObjectKind::Archetype {
+                    // Archetypes default to using `from_similar` from the Component
                     let (typ_unwrapped, _) = quote_field_type_from_field(objects, field, true);
                     // archetype always delegate field init to the component array object
                     format!("converter={typ_unwrapped}Array.from_similar, # type: ignore[misc]\n")
+                } else if !default_converter.is_empty() {
+                    code.push_text(&converter_function, 1, 0);
+                    format!("converter={default_converter}")
                 } else {
-                    // components and datatypes have converters only if manually provided
-                    let override_name = format!(
-                        "{}_{}_converter",
-                        name.to_lowercase(),
-                        field.name.to_lowercase()
-                    );
-                    if overrides.contains(&override_name) {
-                        format!("converter={override_name}")
-                    } else if !default_converter.is_empty() {
-                        code.push_text(&converter_function, 1, 0);
-                        format!("converter={default_converter}")
-                    } else {
-                        String::new()
-                    }
+                    String::new()
                 };
                 field_converters.insert(field.fqname.clone(), converter);
             }

docs/code-examples/image_simple_v2.py

@@ -0,0 +1,14 @@
+"""Create and log an image."""
+
+import numpy as np
+import rerun as rr
+import rerun.experimental as rr2
+
+# Create an image with numpy
+image = np.zeros((200, 300, 3), dtype=np.uint8)
+image[:, :, 0] = 255
+image[50:150, 50:150] = (0, 255, 0)
+
+rr.init("rerun_example_images", spawn=True)
+
+rr2.log("simple", rr2.Image(image))

rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/image.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from rerun.log.error_utils import _send_warning
+
+from ...datatypes import TensorDataArray
+
+if TYPE_CHECKING:
+    from ...datatypes import TensorDataArrayLike
+
+
+def image_data_converter(data: TensorDataArrayLike) -> TensorDataArray:
+    tensor_data = TensorDataArray.from_similar(data)
+
+    # TODO(jleibs): Doing this on raw arrow data is not great. Clean this up
+    # once we coerce to a canonical non-arrow type.
+    shape = tensor_data[0].value["shape"].values.field(0).to_numpy()
+    non_empty_dims = [d for d in shape if d != 1]
+    num_non_empty_dims = len(non_empty_dims)
+
+    # TODO(jleibs): What `recording` should we be passing here? How should we be getting it?
+    if num_non_empty_dims < 2 or 3 < num_non_empty_dims:
+        _send_warning(f"Expected image, got array of shape {shape}", 1, recording=None)
+
+    if num_non_empty_dims == 3:
+        depth = shape[-1]
+        if depth not in (1, 3, 4):
+            _send_warning(
+                f"Expected image depth of 1 (gray), 3 (RGB) or 4 (RGBA). Instead got array of shape {shape}",
+                1,
+                recording=None,
+            )
+
+    # TODO(jleibs): The rust code labels the tensor dimensions as well. Would be nice to do something
+    # similar here if they are unnamed.
+
+    # TODO(jleibs): Should we enforce specific names on images? Specifically, what if the existing names are wrong.
+
+    return tensor_data

rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_buffer.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+
+
+def tensorbuffer_inner_converter(inner: npt.ArrayLike) -> npt.NDArray[Any]:
+    # A tensor buffer is always a flat array
+    return np.asarray(inner).flatten()