From 0f99f9bd83a357294e45adaabb942875e6bab56f Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Fri, 1 Sep 2023 01:40:03 +0200 Subject: [PATCH 01/12] WIP: Initial support for python tensor --- .../rerun/datatypes/tensor_data.fbs | 1 + .../definitions/rerun/datatypes/tensor_id.fbs | 1 + crates/re_types/source_hash.txt | 2 +- crates/re_types_builder/src/codegen/python.rs | 1 + .../_rerun2/datatypes/_overrides/__init__.py | 1 + .../datatypes/_overrides/tensor_data.py | 130 ++++++++++++++++++ .../rerun/_rerun2/datatypes/tensor_data.py | 10 +- .../rerun/_rerun2/datatypes/tensor_id.py | 9 +- rerun_py/rerun_sdk/rerun/experimental.py | 2 + rerun_py/tests/unit/test_tensor.py | 39 ++++++ 10 files changed, 188 insertions(+), 8 deletions(-) create mode 100644 rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py create mode 100644 rerun_py/tests/unit/test_tensor.py diff --git a/crates/re_types/definitions/rerun/datatypes/tensor_data.fbs b/crates/re_types/definitions/rerun/datatypes/tensor_data.fbs index 61f069aa2c18..23f70b9dcf6c 100644 --- a/crates/re_types/definitions/rerun/datatypes/tensor_data.fbs +++ b/crates/re_types/definitions/rerun/datatypes/tensor_data.fbs @@ -19,6 +19,7 @@ namespace rerun.datatypes; /// which stores a contiguous array of typed values. table TensorData ( order: 100, + "attr.python.array_aliases": "npt.NDArray[np.uint8]", "attr.rust.derive": "PartialEq" ) { id: rerun.datatypes.TensorId (order: 100); diff --git a/crates/re_types/definitions/rerun/datatypes/tensor_id.fbs b/crates/re_types/definitions/rerun/datatypes/tensor_id.fbs index 5059df1c0406..d90f34443bef 100644 --- a/crates/re_types/definitions/rerun/datatypes/tensor_id.fbs +++ b/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" ) { diff --git a/crates/re_types/source_hash.txt b/crates/re_types/source_hash.txt index ba1e02006196..a76e67820620 100644 --- a/crates/re_types/source_hash.txt +++ b/crates/re_types/source_hash.txt @@ -1,4 +1,4 @@ # This is a sha256 hash for all direct and indirect dependencies of this crate's build script. # It can be safely removed at anytime to force the build script to run again. # Check out build.rs to see how it's computed. -02e355095c59bc89dac399a4bb4364ff9aa1a906b0162dcc5f4c663266c4220a +4fdaa4bd554c4ab9196b183f63d3e337b3ca0e08b549a7c21a2e708447e21902 diff --git a/crates/re_types_builder/src/codegen/python.rs b/crates/re_types_builder/src/codegen/python.rs index 6a3572b559c9..014d2fe91a57 100644 --- a/crates/re_types_builder/src/codegen/python.rs +++ b/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, diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/__init__.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/__init__.py index e976479c9ccb..f1970cb330c8 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/__init__.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/__init__.py @@ -20,6 +20,7 @@ from .rotation3d import rotation3d_inner_converter from .rotation_axis_angle import rotationaxisangle_angle_converter from .scale3d import scale3d_inner_converter +from .tensor_data import tensordata_native_to_pa_array from .transform3d import transform3d_native_to_pa_array from .translation_and_mat3x3 import translationandmat3x3_init from .translation_rotation_scale3d import translationrotationscale3d_init diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py new file mode 100644 index 000000000000..7c5aab39308b --- /dev/null +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py @@ -0,0 +1,130 @@ +from __future__ import annotations + +import uuid +from typing import TYPE_CHECKING, Final, Sequence, Union, cast + +import numpy as np +import numpy.typing as npt +import pyarrow as pa + +if TYPE_CHECKING: + from .. import TensorDataArrayLike, TensorDataLike, TensorBufferLike + + +# TODO(jleibs): Move this somewhere common +def _build_dense_union(data_type: pa.DenseUnionType, discriminant: str, child: pa.Array) -> pa.Array: + """ + Build a dense UnionArray given the `data_type`, a discriminant, and the child value array. + + If the discriminant string doesn't match any possible value, a `ValueError` is raised. + """ + try: + idx = [f.name for f in list(data_type)].index(discriminant) + type_ids = pa.array([idx] * len(child), type=pa.int8()) + value_offsets = pa.array(range(len(child)), type=pa.int32()) + + children = [pa.nulls(0, type=f.type) for f in list(data_type)] + try: + children[idx] = child.cast(data_type[idx].type, safe=False) + except pa.ArrowInvalid: + # Since we're having issues with nullability in union types (see below), + # the cast sometimes fails but can be skipped. + children[idx] = child + + return pa.Array.from_buffers( + type=data_type, + length=len(child), + buffers=[None, type_ids.buffers()[1], value_offsets.buffers()[1]], + children=children, + ) + + except ValueError as e: + raise ValueError(e.args) + + +def _build_tensorid(id: uuid.UUID) -> pa.Array: + from .. import TensorIdType + + data_type = TensorIdType().storage_type + + array = np.asarray(list(id.bytes), dtype=np.uint8).flatten() + return pa.FixedSizeListArray.from_arrays(array, type=data_type) + + +def _build_shape_array(dims: Sequence[int]) -> pa.Array: + from .. import TensorDimensionType + + data_type = TensorDimensionType().storage_type + + array = np.asarray(dims, dtype=np.uint64).flatten() + names = pa.array(["" for d in dims], mask=[True for d in dims], type=data_type.field("name").type) + + return pa.ListArray.from_arrays( + offsets=[0, len(array)], + values=pa.StructArray.from_arrays( + [ + array, + names, + ], + fields=[data_type.field("size"), data_type.field("name")], + ), + ) + + +DTYPE_MAP: Final[dict[npt.DTypeLike, str]] = { + np.uint8: "U8", + np.uint16: "U16", + np.uint32: "U32", + np.uint64: "U64", + np.int8: "I8", + np.int16: "I16", + np.int32: "I32", + np.int64: "I64", + np.float16: "F16", + np.float32: "F32", + np.float64: "F64", +} + + +def _build_buffer_array(buffer: TensorBufferLike) -> pa.Array: + from .. import TensorBuffer, TensorBufferType + + data_type = TensorBufferType().storage_type + + if isinstance(buffer, TensorBuffer): + buffer = buffer.inner + + buffer = buffer.flatten() + + data_inner = pa.ListArray.from_arrays(pa.array([0, len(buffer)]), buffer) + + return _build_dense_union( + data_type, + discriminant=DTYPE_MAP[buffer.dtype.type], + child=data_inner, + ) + + +def tensordata_native_to_pa_array(data: TensorDataArrayLike, data_type: pa.DataType) -> pa.Array: + from .. import TensorData + + if isinstance(data, TensorData): + data = data.buffer.inner + + tensor_id = _build_tensorid(uuid.uuid4()) + shape = _build_shape_array(data.shape).cast(data_type.field("shape").type) + buffer = _build_buffer_array(data) + + storage = pa.StructArray.from_arrays( + [ + tensor_id, + shape, + buffer, + ], + fields=[data_type.field("id"), data_type.field("shape"), data_type.field("buffer")], + ).cast(data_type) + + storage.validate(full=True) + # TODO(john) enable extension type wrapper + # return cast(TensorArray, pa.ExtensionArray.from_storage(TensorType(), storage)) + return storage # type: ignore[no-any-return] diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py index 332ef9ff107f..9b8866858c73 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py @@ -4,6 +4,8 @@ from typing import Sequence, Union +import numpy as np +import numpy.typing as npt import pyarrow as pa from attrs import define, field @@ -12,6 +14,7 @@ BaseExtensionArray, BaseExtensionType, ) +from ._overrides import tensordata_native_to_pa_array # noqa: F401 __all__ = ["TensorData", "TensorDataArray", "TensorDataArrayLike", "TensorDataLike", "TensorDataType"] @@ -49,10 +52,7 @@ class TensorData: TensorDataLike = TensorData -TensorDataArrayLike = Union[ - TensorData, - Sequence[TensorDataLike], -] +TensorDataArrayLike = Union[TensorData, Sequence[TensorDataLike], npt.NDArray[np.uint8]] # --- Arrow support --- @@ -176,7 +176,7 @@ class TensorDataArray(BaseExtensionArray[TensorDataArrayLike]): @staticmethod def _native_to_pa_array(data: TensorDataArrayLike, data_type: pa.DataType) -> pa.Array: - raise NotImplementedError + return tensordata_native_to_pa_array(data, data_type) TensorDataType._ARRAY_TYPE = TensorDataArray diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_id.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_id.py index de409fd4900c..dc5b7b2c6cef 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_id.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_id.py @@ -2,7 +2,8 @@ from __future__ import annotations -from typing import Any, Sequence, Union +import uuid +from typing import TYPE_CHECKING, Any, Sequence, Union import numpy as np import numpy.typing as npt @@ -28,7 +29,11 @@ def __array__(self, dtype: npt.DTypeLike = None) -> npt.NDArray[Any]: return np.asarray(self.uuid, dtype=dtype) -TensorIdLike = TensorId +if TYPE_CHECKING: + TensorIdLike = Union[TensorId, uuid.UUID] +else: + TensorIdLike = Any + TensorIdArrayLike = Union[ TensorId, Sequence[TensorIdLike], diff --git a/rerun_py/rerun_sdk/rerun/experimental.py b/rerun_py/rerun_sdk/rerun/experimental.py index 05660021206b..0febb4f06138 100644 --- a/rerun_py/rerun_sdk/rerun/experimental.py +++ b/rerun_py/rerun_sdk/rerun/experimental.py @@ -17,6 +17,7 @@ "LineStrips3D", "Points2D", "Points3D", + "Tensor", "Transform3D", "add_space_view", "arch", @@ -42,6 +43,7 @@ LineStrips3D, Points2D, Points3D, + Tensor, Transform3D, ) from ._rerun2.log import log diff --git a/rerun_py/tests/unit/test_tensor.py b/rerun_py/tests/unit/test_tensor.py new file mode 100644 index 000000000000..da1a807476e2 --- /dev/null +++ b/rerun_py/tests/unit/test_tensor.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import itertools +from typing import Any, Optional, cast + +import numpy as np +import pytest +import rerun.experimental as rr2 +from numpy.random import default_rng +from rerun.experimental import cmp as rrc +from rerun.experimental import dt as rrd + + +rng = default_rng(12345) +tensor_data = rng.uniform(0.0, 1.0, (8, 6, 3, 5)) + + +def tensor_data_expected() -> Any: + return rrc.TensorDataArray.from_similar(tensor_data) + + +rng = default_rng(12345) +tensor_data_array: list[rrd.TensorDataArrayLike | None] = [tensor_data] + + +def compare_tensors(left, right): + assert left.storage.field(1) == right.storage.field(1) + assert left.storage.field(2) == right.storage.field(2) + + +def test_tensor() -> None: + for data in tensor_data_array: + arch = rr2.Tensor( + data, + ) + + expected = tensor_data_expected() + + compare_tensors(arch.data, expected) From 7a50aea6e96e699f601e7dc51a9e0b5a9c045fb9 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Fri, 1 Sep 2023 20:38:51 +0200 Subject: [PATCH 02/12] typing and edge cases --- .../rerun/datatypes/tensor_data.fbs | 2 +- .../datatypes/_overrides/tensor_data.py | 48 +++++++++++-------- .../rerun/_rerun2/datatypes/tensor_data.py | 15 +++++- rerun_py/tests/unit/test_tensor.py | 46 +++++++++++------- 4 files changed, 72 insertions(+), 39 deletions(-) diff --git a/crates/re_types/definitions/rerun/datatypes/tensor_data.fbs b/crates/re_types/definitions/rerun/datatypes/tensor_data.fbs index 23f70b9dcf6c..988a19b768f7 100644 --- a/crates/re_types/definitions/rerun/datatypes/tensor_data.fbs +++ b/crates/re_types/definitions/rerun/datatypes/tensor_data.fbs @@ -19,7 +19,7 @@ namespace rerun.datatypes; /// which stores a contiguous array of typed values. table TensorData ( order: 100, - "attr.python.array_aliases": "npt.NDArray[np.uint8]", + "attr.python.array_aliases": "npt.NDArray[np.uint8] | npt.NDArray[np.uint16] | npt.NDArray[np.uint32] | npt.NDArray[np.uint64] | npt.NDArray[np.int8] | npt.NDArray[np.int16] | npt.NDArray[np.int32] | npt.NDArray[np.int64] | npt.NDArray[np.float32] | npt.NDArray[np.float64]", "attr.rust.derive": "PartialEq" ) { id: rerun.datatypes.TensorId (order: 100); diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py index 7c5aab39308b..15b918095103 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py @@ -1,14 +1,14 @@ from __future__ import annotations import uuid -from typing import TYPE_CHECKING, Final, Sequence, Union, cast +from typing import TYPE_CHECKING, Final import numpy as np import numpy.typing as npt import pyarrow as pa if TYPE_CHECKING: - from .. import TensorDataArrayLike, TensorDataLike, TensorBufferLike + from .. import TensorBufferLike, TensorDataArrayLike, TensorDimension, TensorIdLike # TODO(jleibs): Move this somewhere common @@ -42,22 +42,28 @@ def _build_dense_union(data_type: pa.DenseUnionType, discriminant: str, child: p raise ValueError(e.args) -def _build_tensorid(id: uuid.UUID) -> pa.Array: - from .. import TensorIdType +def _build_tensorid(id: TensorIdLike) -> pa.Array: + from .. import TensorId, TensorIdType + + if isinstance(id, uuid.UUID): + array = np.asarray(list(id.bytes), dtype=np.uint8) + elif isinstance(id, TensorId): + array = id.uuid + else: + raise ValueError("Unsupported TensorId input") data_type = TensorIdType().storage_type - array = np.asarray(list(id.bytes), dtype=np.uint8).flatten() return pa.FixedSizeListArray.from_arrays(array, type=data_type) -def _build_shape_array(dims: Sequence[int]) -> pa.Array: +def _build_shape_array(dims: list[TensorDimension]) -> pa.Array: from .. import TensorDimensionType data_type = TensorDimensionType().storage_type - array = np.asarray(dims, dtype=np.uint64).flatten() - names = pa.array(["" for d in dims], mask=[True for d in dims], type=data_type.field("name").type) + array = np.asarray([d.size for d in dims], dtype=np.uint64).flatten() + names = pa.array([d.name for d in dims], mask=[d is None for d in dims], type=data_type.field("name").type) return pa.ListArray.from_arrays( offsets=[0, len(array)], @@ -106,16 +112,23 @@ def _build_buffer_array(buffer: TensorBufferLike) -> pa.Array: def tensordata_native_to_pa_array(data: TensorDataArrayLike, data_type: pa.DataType) -> pa.Array: - from .. import TensorData + from .. import TensorData, TensorDimension + + if isinstance(data, np.ndarray): + tensor_id = _build_tensorid(uuid.uuid4()) + shape = [TensorDimension(d) for d in data.shape] + shape = _build_shape_array(shape).cast(data_type.field("shape").type) + buffer = _build_buffer_array(data) - if isinstance(data, TensorData): - data = data.buffer.inner + elif isinstance(data, TensorData): + tensor_id = _build_tensorid(data.id) + shape = _build_shape_array(data.shape).cast(data_type.field("shape").type) + buffer = _build_buffer_array(data.buffer) - tensor_id = _build_tensorid(uuid.uuid4()) - shape = _build_shape_array(data.shape).cast(data_type.field("shape").type) - buffer = _build_buffer_array(data) + else: + raise ValueError("Unsupported TensorData source") - storage = pa.StructArray.from_arrays( + return pa.StructArray.from_arrays( [ tensor_id, shape, @@ -123,8 +136,3 @@ def tensordata_native_to_pa_array(data: TensorDataArrayLike, data_type: pa.DataT ], fields=[data_type.field("id"), data_type.field("shape"), data_type.field("buffer")], ).cast(data_type) - - storage.validate(full=True) - # TODO(john) enable extension type wrapper - # return cast(TensorArray, pa.ExtensionArray.from_storage(TensorType(), storage)) - return storage # type: ignore[no-any-return] diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py index 9b8866858c73..1bb84014f831 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py @@ -52,7 +52,20 @@ class TensorData: TensorDataLike = TensorData -TensorDataArrayLike = Union[TensorData, Sequence[TensorDataLike], npt.NDArray[np.uint8]] +TensorDataArrayLike = Union[ + TensorData, + Sequence[TensorDataLike], + npt.NDArray[np.uint8] + | npt.NDArray[np.uint16] + | npt.NDArray[np.uint32] + | npt.NDArray[np.uint64] + | npt.NDArray[np.int8] + | npt.NDArray[np.int16] + | npt.NDArray[np.int32] + | npt.NDArray[np.int64] + | npt.NDArray[np.float32] + | npt.NDArray[np.float64], +] # --- Arrow support --- diff --git a/rerun_py/tests/unit/test_tensor.py b/rerun_py/tests/unit/test_tensor.py index da1a807476e2..8a0de3171c4b 100644 --- a/rerun_py/tests/unit/test_tensor.py +++ b/rerun_py/tests/unit/test_tensor.py @@ -1,39 +1,51 @@ from __future__ import annotations -import itertools -from typing import Any, Optional, cast +from typing import Any import numpy as np -import pytest import rerun.experimental as rr2 -from numpy.random import default_rng from rerun.experimental import cmp as rrc from rerun.experimental import dt as rrd +rng = np.random.default_rng(12345) +RANDOM_TENSOR_SOURCE = rng.uniform(0.0, 1.0, (8, 6, 3, 5)) -rng = default_rng(12345) -tensor_data = rng.uniform(0.0, 1.0, (8, 6, 3, 5)) +TENSOR_DATA_INPUTS: list[rrd.TensorDataArrayLike | None] = [ + # Full explicit construction + rrd.TensorData( + id=rrd.TensorId(uuid=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]), + shape=[ + rrd.TensorDimension(8), + rrd.TensorDimension(6), + rrd.TensorDimension(3), + rrd.TensorDimension(5), + ], + buffer=rrd.TensorBuffer(RANDOM_TENSOR_SOURCE), + ), + # Implicit construction from ndarray + RANDOM_TENSOR_SOURCE, +] -def tensor_data_expected() -> Any: - return rrc.TensorDataArray.from_similar(tensor_data) +CHECK_TENSOR_ID: list[bool] = [True, False] -rng = default_rng(12345) -tensor_data_array: list[rrd.TensorDataArrayLike | None] = [tensor_data] +def tensor_data_expected() -> Any: + return rrc.TensorDataArray.from_similar(TENSOR_DATA_INPUTS[0]) -def compare_tensors(left, right): +def compare_tensors(left: Any, right: Any, check_id: bool) -> None: + # Skip tensor_id + if check_id: + assert left.storage.field(0) == right.storage.field(0) assert left.storage.field(1) == right.storage.field(1) assert left.storage.field(2) == right.storage.field(2) def test_tensor() -> None: - for data in tensor_data_array: - arch = rr2.Tensor( - data, - ) + expected = tensor_data_expected() - expected = tensor_data_expected() + for input, check_id in zip(TENSOR_DATA_INPUTS, CHECK_TENSOR_ID): + arch = rr2.Tensor(data=input) - compare_tensors(arch.data, expected) + compare_tensors(arch.data, expected, check_id) From c391a847bd1c816d1235b655f296ed162aa3462d Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Sat, 2 Sep 2023 00:28:05 +0200 Subject: [PATCH 03/12] Beginning image validation and more construction options --- crates/re_types/source_hash.txt | 2 +- crates/re_types_builder/src/codegen/python.rs | 28 ++-- .../_rerun2/archetypes/_overrides/__init__.py | 2 + .../_rerun2/archetypes/_overrides/image.py | 18 +++ .../rerun/_rerun2/archetypes/image.py | 6 +- .../_rerun2/datatypes/_overrides/__init__.py | 3 +- .../datatypes/_overrides/tensor_data.py | 148 ++++++++++++++---- .../_rerun2/datatypes/_overrides/tensor_id.py | 21 +++ .../rerun/_rerun2/datatypes/tensor_data.py | 7 +- .../rerun/_rerun2/datatypes/tensor_id.py | 6 +- rerun_py/rerun_sdk/rerun/experimental.py | 2 + rerun_py/tests/unit/test_image.py | 50 ++++++ rerun_py/tests/unit/test_tensor.py | 39 +++-- 13 files changed, 263 insertions(+), 69 deletions(-) create mode 100644 rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/image.py create mode 100644 rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_id.py create mode 100644 rerun_py/tests/unit/test_image.py diff --git a/crates/re_types/source_hash.txt b/crates/re_types/source_hash.txt index a76e67820620..e21e526bee61 100644 --- a/crates/re_types/source_hash.txt +++ b/crates/re_types/source_hash.txt @@ -1,4 +1,4 @@ # This is a sha256 hash for all direct and indirect dependencies of this crate's build script. # It can be safely removed at anytime to force the build script to run again. # Check out build.rs to see how it's computed. -4fdaa4bd554c4ab9196b183f63d3e337b3ca0e08b549a7c21a2e708447e21902 +b9d21f0105824f605c7148277477eeda0904f3480ac511712db94b74067219a5 diff --git a/crates/re_types_builder/src/codegen/python.rs b/crates/re_types_builder/src/codegen/python.rs index 014d2fe91a57..c4ad0b02d63c 100644 --- a/crates/re_types_builder/src/codegen/python.rs +++ b/crates/re_types_builder/src/codegen/python.rs @@ -437,25 +437,23 @@ impl QuotedObject { let (default_converter, converter_function) = quote_field_converter_from_field(obj, objects, field); - let converter = if *kind == ObjectKind::Archetype { + // components and datatypes have converters only if manually provided + 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 { 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); } diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/__init__.py b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/__init__.py index 9d48db4f9f85..2d22b434d76c 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/__init__.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/__init__.py @@ -1 +1,3 @@ from __future__ import annotations + +from .image import image_data_converter diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/image.py b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/image.py new file mode 100644 index 000000000000..1e7441f77bec --- /dev/null +++ b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/image.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +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. + dimensions = tensor_data[0].value["shape"].values.field(0).to_numpy() + + return tensor_data diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/image.py b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/image.py index e22c1e27510b..adb8d4b586dc 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/image.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/image.py @@ -8,6 +8,7 @@ from .._baseclasses import ( Archetype, ) +from ._overrides import image_data_converter # noqa: F401 __all__ = ["Image"] @@ -25,10 +26,7 @@ class Image(Archetype): The viewer has limited support for ignoring extra empty dimensions. """ - data: components.TensorDataArray = field( - metadata={"component": "primary"}, - converter=components.TensorDataArray.from_similar, # type: ignore[misc] - ) + data: components.TensorDataArray = field(metadata={"component": "primary"}, converter=image_data_converter) """ The image data. Should always be a rank-2 or rank-3 tensor. """ diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/__init__.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/__init__.py index f1970cb330c8..9aad8671d5d0 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/__init__.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/__init__.py @@ -20,7 +20,8 @@ from .rotation3d import rotation3d_inner_converter from .rotation_axis_angle import rotationaxisangle_angle_converter from .scale3d import scale3d_inner_converter -from .tensor_data import tensordata_native_to_pa_array +from .tensor_data import tensordata_init, tensordata_native_to_pa_array +from .tensor_id import tensorid_uuid_converter from .transform3d import transform3d_native_to_pa_array from .translation_and_mat3x3 import translationandmat3x3_init from .translation_rotation_scale3d import translationrotationscale3d_init diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py index 15b918095103..010be4ce8858 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py @@ -1,14 +1,131 @@ from __future__ import annotations import uuid -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Sequence import numpy as np import numpy.typing as npt import pyarrow as pa +from rerun.log.error_utils import _send_warning + if TYPE_CHECKING: - from .. import TensorBufferLike, TensorDataArrayLike, TensorDimension, TensorIdLike + from .. import TensorBufferLike, TensorData, TensorDataArrayLike, TensorDimension, TensorDimensionLike, TensorIdLike + +################################################################################ +# Init overrides +################################################################################ + + +def tensordata_init( + self: TensorData, + *, + id: TensorIdLike | None = None, + shape: Sequence[TensorDimensionLike] | None = None, + buffer: TensorBufferLike | None = None, + array: npt.NDArray[np.float32] + | npt.NDArray[np.float64] + | npt.NDArray[np.int16] + | npt.NDArray[np.int32] + | npt.NDArray[np.int64] + | npt.NDArray[np.int8] + | npt.NDArray[np.uint16] + | npt.NDArray[np.uint32] + | npt.NDArray[np.uint64] + | npt.NDArray[np.uint8] + | None = None, + names: Sequence[str] | None = None, +) -> None: + if array is None and buffer is None: + raise ValueError("Must provide one of 'array' or 'buffer'") + if array is not None and buffer is not None: + raise ValueError("Can only provide one of 'array' or 'buffer'") + if buffer is not None and shape is None: + raise ValueError("If 'buffer' is provided, 'shape' is also required") + if shape is not None and names is not None: + raise ValueError("Can only provide one of 'shape' or 'names'") + + from .. import TensorBuffer, TensorDimension + from ..tensor_data import _tensordata_buffer_converter, _tensordata_id_converter + + # Assign an id if one wasn't provided + if id: + self.id = _tensordata_id_converter(id) + else: + self.id = _tensordata_id_converter(uuid.uuid4()) + + if shape: + resolved_shape = list(shape) + else: + resolved_shape = None + + # Figure out the shape + if array is not None: + # If a shape we provided, it must match the array + if resolved_shape: + shape_tuple = tuple(d.size for d in resolved_shape) + if shape_tuple != array.shape: + raise ValueError(f"Provided array ({array.shape}) does not match shape argument ({shape_tuple}).") + elif names: + if len(array.shape) != len(names): + _send_warning( + ( + f"len(array.shape) = {len(array.shape)} != " + + f"len(names) = {len(names)}. Dropping tensor dimension names." + ), + 2, + ) + resolved_shape = [TensorDimension(size, name) for size, name in zip(array.shape, names)] + else: + resolved_shape = [TensorDimension(size) for size in array.shape] + + if resolved_shape is not None: + self.shape = resolved_shape + else: + # This shouldn't be possible but typing can't figure it out + raise ValueError("No shape provided.") + + if buffer is not None: + self.buffer = _tensordata_buffer_converter(buffer) + elif array is not None: + self.buffer = TensorBuffer(array.flatten()) + + +################################################################################ +# Arrow converters +################################################################################ + + +def tensordata_native_to_pa_array(data: TensorDataArrayLike, data_type: pa.DataType) -> pa.Array: + from .. import TensorData, TensorDimension + + if isinstance(data, np.ndarray): + tensor_id = _build_tensorid(uuid.uuid4()) + shape = [TensorDimension(d) for d in data.shape] + shape = _build_shape_array(shape).cast(data_type.field("shape").type) + buffer = _build_buffer_array(data) + + elif isinstance(data, TensorData): + tensor_id = _build_tensorid(data.id) + shape = _build_shape_array(data.shape).cast(data_type.field("shape").type) + buffer = _build_buffer_array(data.buffer) + + else: + raise ValueError("Unsupported TensorData source") + + return pa.StructArray.from_arrays( + [ + tensor_id, + shape, + buffer, + ], + fields=[data_type.field("id"), data_type.field("shape"), data_type.field("buffer")], + ).cast(data_type) + + +################################################################################ +# Internal construction helpers +################################################################################ # TODO(jleibs): Move this somewhere common @@ -109,30 +226,3 @@ def _build_buffer_array(buffer: TensorBufferLike) -> pa.Array: discriminant=DTYPE_MAP[buffer.dtype.type], child=data_inner, ) - - -def tensordata_native_to_pa_array(data: TensorDataArrayLike, data_type: pa.DataType) -> pa.Array: - from .. import TensorData, TensorDimension - - if isinstance(data, np.ndarray): - tensor_id = _build_tensorid(uuid.uuid4()) - shape = [TensorDimension(d) for d in data.shape] - shape = _build_shape_array(shape).cast(data_type.field("shape").type) - buffer = _build_buffer_array(data) - - elif isinstance(data, TensorData): - tensor_id = _build_tensorid(data.id) - shape = _build_shape_array(data.shape).cast(data_type.field("shape").type) - buffer = _build_buffer_array(data.buffer) - - else: - raise ValueError("Unsupported TensorData source") - - return pa.StructArray.from_arrays( - [ - tensor_id, - shape, - buffer, - ], - fields=[data_type.field("id"), data_type.field("shape"), data_type.field("buffer")], - ).cast(data_type) diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_id.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_id.py new file mode 100644 index 000000000000..45b7afeba55e --- /dev/null +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_id.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +import uuid + +import numpy as np +import numpy.typing as npt + + +def tensorid_uuid_converter(id: bytes | uuid.UUID | npt.ArrayLike) -> npt.NDArray[np.uint8]: + if isinstance(id, uuid.UUID): + id = id.bytes + + if isinstance(id, bytes): + id = list(id) + + id = np.asarray(id, dtype=np.uint8) + + if len(id) != 16: + raise ValueError("TensorId must be 16 bytes") + + return id diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py index 1bb84014f831..eb5c848d33cf 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py @@ -14,7 +14,7 @@ BaseExtensionArray, BaseExtensionType, ) -from ._overrides import tensordata_native_to_pa_array # noqa: F401 +from ._overrides import tensordata_init, tensordata_native_to_pa_array # noqa: F401 __all__ = ["TensorData", "TensorDataArray", "TensorDataArrayLike", "TensorDataLike", "TensorDataType"] @@ -33,7 +33,7 @@ def _tensordata_buffer_converter(x: datatypes.TensorBufferLike) -> datatypes.Ten return datatypes.TensorBuffer(x) -@define +@define(init=False) class TensorData: """ A multi-dimensional `Tensor` of data. @@ -46,6 +46,9 @@ class TensorData: which stores a contiguous array of typed values. """ + def __init__(self, *args, **kwargs): # type: ignore[no-untyped-def] + tensordata_init(self, *args, **kwargs) + id: datatypes.TensorId = field(converter=_tensordata_id_converter) shape: list[datatypes.TensorDimension] = field() buffer: datatypes.TensorBuffer = field(converter=_tensordata_buffer_converter) diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_id.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_id.py index dc5b7b2c6cef..25b7d42f67c8 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_id.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_id.py @@ -14,16 +14,14 @@ BaseExtensionArray, BaseExtensionType, ) -from .._converters import ( - to_np_uint8, -) +from ._overrides import tensorid_uuid_converter # noqa: F401 __all__ = ["TensorId", "TensorIdArray", "TensorIdArrayLike", "TensorIdLike", "TensorIdType"] @define class TensorId: - uuid: npt.NDArray[np.uint8] = field(converter=to_np_uint8) + uuid: npt.NDArray[np.uint8] = field(converter=tensorid_uuid_converter) def __array__(self, dtype: npt.DTypeLike = None) -> npt.NDArray[Any]: return np.asarray(self.uuid, dtype=dtype) diff --git a/rerun_py/rerun_sdk/rerun/experimental.py b/rerun_py/rerun_sdk/rerun/experimental.py index 0febb4f06138..6f82a979b14a 100644 --- a/rerun_py/rerun_sdk/rerun/experimental.py +++ b/rerun_py/rerun_sdk/rerun/experimental.py @@ -13,6 +13,7 @@ "AnnotationContext", "Arrows3D", "DisconnectedSpace", + "Image", "LineStrips2D", "LineStrips3D", "Points2D", @@ -39,6 +40,7 @@ AnnotationContext, Arrows3D, DisconnectedSpace, + Image, LineStrips2D, LineStrips3D, Points2D, diff --git a/rerun_py/tests/unit/test_image.py b/rerun_py/tests/unit/test_image.py new file mode 100644 index 000000000000..ed13da98a9cb --- /dev/null +++ b/rerun_py/tests/unit/test_image.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +import rerun.experimental as rr2 +from rerun.experimental import cmp as rrc +from rerun.experimental import dt as rrd + +rng = np.random.default_rng(12345) +RANDOM_IMAGE_SOURCE = rng.uniform(0.0, 1.0, (8, 6, 3, 5)) + + +IMAGE_DATA_INPUTS: list[rrd.TensorDataArrayLike | None] = [ + # Full explicit construction + rrd.TensorData( + id=rrd.TensorId(uuid=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]), + shape=[ + rrd.TensorDimension(8), + rrd.TensorDimension(6), + rrd.TensorDimension(3), + rrd.TensorDimension(5), + ], + buffer=rrd.TensorBuffer(RANDOM_IMAGE_SOURCE), + ), + # Implicit construction from ndarray + RANDOM_IMAGE_SOURCE, +] + +CHECK_IMAGE_ID: list[bool] = [True, False] + + +def tensor_data_expected() -> Any: + return rrc.TensorDataArray.from_similar(IMAGE_DATA_INPUTS[0]) + + +def compare_tensors(left: Any, right: Any, check_id: bool) -> None: + if check_id: + assert left.storage.field(0) == right.storage.field(0) + assert left.storage.field(1) == right.storage.field(1) + assert left.storage.field(2) == right.storage.field(2) + + +def test_image() -> None: + expected = tensor_data_expected() + + for input, check_id in zip(IMAGE_DATA_INPUTS, CHECK_IMAGE_ID): + arch = rr2.Image(data=input) + + compare_tensors(arch.data, expected, check_id) diff --git a/rerun_py/tests/unit/test_tensor.py b/rerun_py/tests/unit/test_tensor.py index 8a0de3171c4b..d4317caffb25 100644 --- a/rerun_py/tests/unit/test_tensor.py +++ b/rerun_py/tests/unit/test_tensor.py @@ -1,5 +1,6 @@ from __future__ import annotations +import uuid from typing import Any import numpy as np @@ -16,36 +17,48 @@ rrd.TensorData( id=rrd.TensorId(uuid=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]), shape=[ - rrd.TensorDimension(8), - rrd.TensorDimension(6), - rrd.TensorDimension(3), - rrd.TensorDimension(5), + rrd.TensorDimension(8, name="a"), + rrd.TensorDimension(6, name="b"), + rrd.TensorDimension(3, name="c"), + rrd.TensorDimension(5, name="d"), ], buffer=rrd.TensorBuffer(RANDOM_TENSOR_SOURCE), ), # Implicit construction from ndarray RANDOM_TENSOR_SOURCE, + # Explicit construction from array + rrd.TensorData(array=RANDOM_TENSOR_SOURCE), + # Explicit construction from array + rrd.TensorData(array=RANDOM_TENSOR_SOURCE, names=["a", "b", "c", "d"]), + # Explicit construction from array + rrd.TensorData(id=uuid.uuid4(), array=RANDOM_TENSOR_SOURCE, names=["a", "b", "c", "d"]), ] -CHECK_TENSOR_ID: list[bool] = [True, False] +# 0 = id +# 1 = shape +# 2 = buffer +CHECK_FIELDS: list[list[int]] = [ + [0, 1, 2], + [2], + [2], + [1, 2], + [1, 2], +] def tensor_data_expected() -> Any: return rrc.TensorDataArray.from_similar(TENSOR_DATA_INPUTS[0]) -def compare_tensors(left: Any, right: Any, check_id: bool) -> None: - # Skip tensor_id - if check_id: - assert left.storage.field(0) == right.storage.field(0) - assert left.storage.field(1) == right.storage.field(1) - assert left.storage.field(2) == right.storage.field(2) +def compare_tensors(left: Any, right: Any, check_fields: list[int]) -> None: + for field in check_fields: + assert left.storage.field(field) == right.storage.field(field) def test_tensor() -> None: expected = tensor_data_expected() - for input, check_id in zip(TENSOR_DATA_INPUTS, CHECK_TENSOR_ID): + for input, check_fields in zip(TENSOR_DATA_INPUTS, CHECK_FIELDS): arch = rr2.Tensor(data=input) - compare_tensors(arch.data, expected, check_id) + compare_tensors(arch.data, expected, check_fields) From 8545df633626db316920b87963cf901d9b9ca6f8 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Sat, 2 Sep 2023 01:51:01 +0200 Subject: [PATCH 04/12] Assorted validation --- .../_rerun2/archetypes/_overrides/image.py | 24 +++- .../_rerun2/datatypes/_overrides/__init__.py | 1 + .../datatypes/_overrides/tensor_buffer.py | 10 ++ .../datatypes/_overrides/tensor_data.py | 112 ++++++++++++------ .../rerun/_rerun2/datatypes/tensor_buffer.py | 5 +- rerun_py/tests/unit/test_image.py | 69 ++++++++--- rerun_py/tests/unit/test_tensor.py | 57 +++++++++ 7 files changed, 227 insertions(+), 51 deletions(-) create mode 100644 rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_buffer.py diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/image.py b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/image.py index 1e7441f77bec..a3c5919135d5 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/image.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/image.py @@ -2,6 +2,8 @@ from typing import TYPE_CHECKING +from rerun.log.error_utils import _send_warning + from ...datatypes import TensorDataArray if TYPE_CHECKING: @@ -13,6 +15,26 @@ def image_data_converter(data: TensorDataArrayLike) -> TensorDataArray: # TODO(jleibs): Doing this on raw arrow data is not great. Clean this up # once we coerce to a canonical non-arrow type. - dimensions = tensor_data[0].value["shape"].values.field(0).to_numpy() + 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) how do we send warnings form down inside these converters? + 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 diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/__init__.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/__init__.py index 9aad8671d5d0..b8f41eb750be 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/__init__.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/__init__.py @@ -20,6 +20,7 @@ from .rotation3d import rotation3d_inner_converter from .rotation_axis_angle import rotationaxisangle_angle_converter from .scale3d import scale3d_inner_converter +from .tensor_buffer import tensorbuffer_inner_converter from .tensor_data import tensordata_init, tensordata_native_to_pa_array from .tensor_id import tensorid_uuid_converter from .transform3d import transform3d_native_to_pa_array diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_buffer.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_buffer.py new file mode 100644 index 000000000000..882a9a8475b2 --- /dev/null +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_buffer.py @@ -0,0 +1,10 @@ +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]: + return np.asarray(inner).flatten() diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py index 010be4ce8858..36027d6ffa08 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py @@ -1,7 +1,9 @@ from __future__ import annotations +import collections import uuid -from typing import TYPE_CHECKING, Final, Sequence +from math import prod +from typing import TYPE_CHECKING, Any, Final, Protocol, Sequence, Union import numpy as np import numpy.typing as npt @@ -12,6 +14,35 @@ if TYPE_CHECKING: from .. import TensorBufferLike, TensorData, TensorDataArrayLike, TensorDimension, TensorDimensionLike, TensorIdLike + +################################################################################ +# Torch-like array converters +################################################################################ + + +class TorchTensorLike(Protocol): + """Describes what is need from a Torch Tensor to be loggable to Rerun.""" + + def numpy(self, force: bool) -> npt.NDArray[Any]: + ... + + +Tensor = Union[npt.ArrayLike, TorchTensorLike] +"""Type helper for a tensor-like object that can be logged to Rerun.""" + + +def _to_numpy(tensor: Tensor) -> npt.NDArray[Any]: + # isinstance is 4x faster than catching AttributeError + if isinstance(tensor, np.ndarray): + return tensor + + try: + # Make available to the cpu + return tensor.numpy(force=True) # type: ignore[union-attr] + except AttributeError: + return np.array(tensor, copy=False) + + ################################################################################ # Init overrides ################################################################################ @@ -23,17 +54,7 @@ def tensordata_init( id: TensorIdLike | None = None, shape: Sequence[TensorDimensionLike] | None = None, buffer: TensorBufferLike | None = None, - array: npt.NDArray[np.float32] - | npt.NDArray[np.float64] - | npt.NDArray[np.int16] - | npt.NDArray[np.int32] - | npt.NDArray[np.int64] - | npt.NDArray[np.int8] - | npt.NDArray[np.uint16] - | npt.NDArray[np.uint32] - | npt.NDArray[np.uint64] - | npt.NDArray[np.uint8] - | None = None, + array: Tensor | None = None, names: Sequence[str] | None = None, ) -> None: if array is None and buffer is None: @@ -61,23 +82,34 @@ def tensordata_init( # Figure out the shape if array is not None: + array = _to_numpy(array) + # If a shape we provided, it must match the array if resolved_shape: shape_tuple = tuple(d.size for d in resolved_shape) if shape_tuple != array.shape: - raise ValueError(f"Provided array ({array.shape}) does not match shape argument ({shape_tuple}).") - elif names: - if len(array.shape) != len(names): _send_warning( ( - f"len(array.shape) = {len(array.shape)} != " - + f"len(names) = {len(names)}. Dropping tensor dimension names." + f"Provided array ({array.shape}) does not match shape argument ({shape_tuple}). " + + "Ignoring shape argument." ), 2, ) - resolved_shape = [TensorDimension(size, name) for size, name in zip(array.shape, names)] - else: - resolved_shape = [TensorDimension(size) for size in array.shape] + resolved_shape = None + + if resolved_shape is None: + if names: + if len(array.shape) != len(names): + _send_warning( + ( + f"len(array.shape) = {len(array.shape)} != " + + f"len(names) = {len(names)}. Dropping tensor dimension names." + ), + 2, + ) + resolved_shape = [TensorDimension(size, name) for size, name in zip(array.shape, names)] + else: + resolved_shape = [TensorDimension(size) for size in array.shape] if resolved_shape is not None: self.shape = resolved_shape @@ -90,6 +122,13 @@ def tensordata_init( elif array is not None: self.buffer = TensorBuffer(array.flatten()) + expected_buffer_size = prod(d.size for d in self.shape) + + if len(self.buffer.inner) != expected_buffer_size: + raise ValueError( + f"Shape and buffer size do not match. {len(self.buffer.inner)} {self.shape}->{expected_buffer_size}" + ) + ################################################################################ # Arrow converters @@ -97,21 +136,24 @@ def tensordata_init( def tensordata_native_to_pa_array(data: TensorDataArrayLike, data_type: pa.DataType) -> pa.Array: - from .. import TensorData, TensorDimension - - if isinstance(data, np.ndarray): - tensor_id = _build_tensorid(uuid.uuid4()) - shape = [TensorDimension(d) for d in data.shape] - shape = _build_shape_array(shape).cast(data_type.field("shape").type) - buffer = _build_buffer_array(data) - - elif isinstance(data, TensorData): - tensor_id = _build_tensorid(data.id) - shape = _build_shape_array(data.shape).cast(data_type.field("shape").type) - buffer = _build_buffer_array(data.buffer) - - else: - raise ValueError("Unsupported TensorData source") + from .. import TensorData + + # If it's a sequence, grab the first one + if isinstance(data, collections.abc.Sequence): + if len(data) != 1: + raise ValueError("Tensors do not support batches") + data = data[0] + + # If it's not a TensorData, it should be an NDArray-like. coerce it into TensorData with the + # constructor. + if not isinstance(data, TensorData): + array = _to_numpy(data) + data = TensorData(array=array) + + # Now build the actual arrow fields + tensor_id = _build_tensorid(data.id) + shape = _build_shape_array(data.shape).cast(data_type.field("shape").type) + buffer = _build_buffer_array(data.buffer) return pa.StructArray.from_arrays( [ diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_buffer.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_buffer.py index 16f9dc776b42..6d321fa215e7 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_buffer.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_buffer.py @@ -13,6 +13,7 @@ BaseExtensionArray, BaseExtensionType, ) +from ._overrides import tensorbuffer_inner_converter # noqa: F401 __all__ = ["TensorBuffer", "TensorBufferArray", "TensorBufferArrayLike", "TensorBufferLike", "TensorBufferType"] @@ -31,7 +32,9 @@ class TensorBuffer: np.uint64 ] | npt.NDArray[ np.uint8 - ] = field() + ] = field( + converter=tensorbuffer_inner_converter + ) """ U8 (npt.NDArray[np.uint8]): diff --git a/rerun_py/tests/unit/test_image.py b/rerun_py/tests/unit/test_image.py index ed13da98a9cb..b8559fb0ef19 100644 --- a/rerun_py/tests/unit/test_image.py +++ b/rerun_py/tests/unit/test_image.py @@ -3,23 +3,23 @@ from typing import Any import numpy as np +import pytest import rerun.experimental as rr2 from rerun.experimental import cmp as rrc from rerun.experimental import dt as rrd rng = np.random.default_rng(12345) -RANDOM_IMAGE_SOURCE = rng.uniform(0.0, 1.0, (8, 6, 3, 5)) +RANDOM_IMAGE_SOURCE = rng.uniform(0.0, 1.0, (10, 20, 3)) -IMAGE_DATA_INPUTS: list[rrd.TensorDataArrayLike | None] = [ +IMAGE_INPUTS: list[rrd.TensorDataArrayLike | None] = [ # Full explicit construction rrd.TensorData( id=rrd.TensorId(uuid=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]), shape=[ - rrd.TensorDimension(8), - rrd.TensorDimension(6), + rrd.TensorDimension(10), + rrd.TensorDimension(20), rrd.TensorDimension(3), - rrd.TensorDimension(5), ], buffer=rrd.TensorBuffer(RANDOM_IMAGE_SOURCE), ), @@ -27,24 +27,65 @@ RANDOM_IMAGE_SOURCE, ] -CHECK_IMAGE_ID: list[bool] = [True, False] +# 0 = id +# 1 = shape +# 2 = buffer +CHECK_FIELDS: list[list[int]] = [ + [0, 1, 2], + [2], +] def tensor_data_expected() -> Any: - return rrc.TensorDataArray.from_similar(IMAGE_DATA_INPUTS[0]) + return rrc.TensorDataArray.from_similar(IMAGE_INPUTS[0]) -def compare_tensors(left: Any, right: Any, check_id: bool) -> None: - if check_id: - assert left.storage.field(0) == right.storage.field(0) - assert left.storage.field(1) == right.storage.field(1) - assert left.storage.field(2) == right.storage.field(2) +def compare_images(left: Any, right: Any, check_fields: list[int]) -> None: + for field in check_fields: + assert left.storage.field(field) == right.storage.field(field) def test_image() -> None: expected = tensor_data_expected() - for input, check_id in zip(IMAGE_DATA_INPUTS, CHECK_IMAGE_ID): + for input, check_fields in zip(IMAGE_INPUTS, CHECK_FIELDS): arch = rr2.Image(data=input) - compare_tensors(arch.data, expected, check_id) + compare_images(arch.data, expected, check_fields) + + +GOOD_IMAGE_INPUTS: list[rrd.TensorDataArrayLike | None] = [ + # Mono + rng.uniform(0.0, 1.0, (10, 20)), + # RGB + rng.uniform(0.0, 1.0, (10, 20, 3)), + # RGBA + rng.uniform(0.0, 1.0, (10, 20, 4)), + # Assorted Extra Dimensions + rng.uniform(0.0, 1.0, (1, 10, 20)), + rng.uniform(0.0, 1.0, (1, 10, 20, 3)), + rng.uniform(0.0, 1.0, (1, 10, 20, 4)), + rng.uniform(0.0, 1.0, (10, 20, 1)), + rng.uniform(0.0, 1.0, (10, 20, 3, 1)), + rng.uniform(0.0, 1.0, (10, 20, 4, 1)), +] + +BAD_IMAGE_INPUTS: list[rrd.TensorDataArrayLike | None] = [ + rng.uniform(0.0, 1.0, (10,)), + rng.uniform(0.0, 1.0, (10, 20, 2)), + rng.uniform(0.0, 1.0, (10, 20, 5)), + rng.uniform(0.0, 1.0, (10, 20, 3, 2)), +] + + +def test_image_shapes() -> None: + import rerun as rr + + rr.set_strict_mode(True) + + for img in GOOD_IMAGE_INPUTS: + rr2.Image(img) + + for img in BAD_IMAGE_INPUTS: + with pytest.raises(TypeError): + rr2.Image(img) diff --git a/rerun_py/tests/unit/test_tensor.py b/rerun_py/tests/unit/test_tensor.py index d4317caffb25..416694b8bf6c 100644 --- a/rerun_py/tests/unit/test_tensor.py +++ b/rerun_py/tests/unit/test_tensor.py @@ -4,6 +4,7 @@ from typing import Any import numpy as np +import pytest import rerun.experimental as rr2 from rerun.experimental import cmp as rrc from rerun.experimental import dt as rrd @@ -62,3 +63,59 @@ def test_tensor() -> None: arch = rr2.Tensor(data=input) compare_tensors(arch.data, expected, check_fields) + + +def test_bad_tensors() -> None: + import rerun as rr + + rr.set_strict_mode(True) + + # No buffers + with pytest.raises(ValueError): + rrd.TensorData(), + + # Buffer with no indication of shape + with pytest.raises(ValueError): + rrd.TensorData( + buffer=RANDOM_TENSOR_SOURCE, + ), + + # Both array and buffer + with pytest.raises(ValueError): + rrd.TensorData( + array=RANDOM_TENSOR_SOURCE, + buffer=RANDOM_TENSOR_SOURCE, + ), + + # Wrong size buffer for dimensions + with pytest.raises(ValueError): + rrd.TensorData( + shape=[ + rrd.TensorDimension(8, name="a"), + rrd.TensorDimension(6, name="b"), + rrd.TensorDimension(3, name="c"), + rrd.TensorDimension(4, name="d"), + ], + buffer=RANDOM_TENSOR_SOURCE, + ), + + # TODO(jleibs) send_warning bottoms out in TypeError but these ought to be ValueErrors + + # Wrong number of names + with pytest.raises(TypeError): + rrd.TensorData( + names=["a", "b", "c"], + array=RANDOM_TENSOR_SOURCE, + ), + + # Shape disagrees with array + with pytest.raises(TypeError): + rrd.TensorData( + shape=[ + rrd.TensorDimension(8, name="a"), + rrd.TensorDimension(6, name="b"), + rrd.TensorDimension(5, name="c"), + rrd.TensorDimension(3, name="d"), + ], + array=RANDOM_TENSOR_SOURCE, + ), From f29498af56b652a0988250a10095f4d810645ccf Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Sat, 2 Sep 2023 01:54:11 +0200 Subject: [PATCH 05/12] Code-example for image API --- docs/code-examples/image_simple_v2.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 docs/code-examples/image_simple_v2.py diff --git a/docs/code-examples/image_simple_v2.py b/docs/code-examples/image_simple_v2.py new file mode 100644 index 000000000000..d4f7639be779 --- /dev/null +++ b/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 Pillow +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)) From d601fbe664f32f4c1dda0eb6a2e6ed0106bef4a0 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Sat, 2 Sep 2023 01:58:27 +0200 Subject: [PATCH 06/12] Fix codegen comment --- crates/re_types/source_hash.txt | 2 +- crates/re_types_builder/src/codegen/python.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/re_types/source_hash.txt b/crates/re_types/source_hash.txt index e21e526bee61..d5c329b415ea 100644 --- a/crates/re_types/source_hash.txt +++ b/crates/re_types/source_hash.txt @@ -1,4 +1,4 @@ # This is a sha256 hash for all direct and indirect dependencies of this crate's build script. # It can be safely removed at anytime to force the build script to run again. # Check out build.rs to see how it's computed. -b9d21f0105824f605c7148277477eeda0904f3480ac511712db94b74067219a5 +2659734e396eb3a29b890308e9c0bf91052d12df5453c23aceba57fef0dcda83 diff --git a/crates/re_types_builder/src/codegen/python.rs b/crates/re_types_builder/src/codegen/python.rs index c4ad0b02d63c..f298a7f87092 100644 --- a/crates/re_types_builder/src/codegen/python.rs +++ b/crates/re_types_builder/src/codegen/python.rs @@ -437,7 +437,6 @@ impl QuotedObject { let (default_converter, converter_function) = quote_field_converter_from_field(obj, objects, field); - // components and datatypes have converters only if manually provided let override_name = format!( "{}_{}_converter", name.to_lowercase(), @@ -446,6 +445,7 @@ impl QuotedObject { 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") From 9ecceca25825ed37f362253cb5da09aa8c17efda Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Sat, 2 Sep 2023 02:00:25 +0200 Subject: [PATCH 07/12] Clarify TODO for warnings --- rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/image.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/image.py b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/image.py index a3c5919135d5..60df2dcde718 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/image.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/_overrides/image.py @@ -19,7 +19,7 @@ def image_data_converter(data: TensorDataArrayLike) -> TensorDataArray: non_empty_dims = [d for d in shape if d != 1] num_non_empty_dims = len(non_empty_dims) - # TODO(jleibs) how do we send warnings form down inside these converters? + # 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) From 9e393d6f2c2dc5509ccaf8393ccbabbe8aa2036d Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Sat, 2 Sep 2023 02:38:46 +0200 Subject: [PATCH 08/12] More comments and cleanup --- .../definitions/rerun/archetypes/image.fbs | 7 +++ crates/re_types/source_hash.txt | 2 +- .../rerun/_rerun2/archetypes/image.py | 18 ++++++++ .../datatypes/_overrides/tensor_buffer.py | 1 + .../datatypes/_overrides/tensor_data.py | 44 ++++++++++++++++--- 5 files changed, 66 insertions(+), 6 deletions(-) diff --git a/crates/re_types/definitions/rerun/archetypes/image.fbs b/crates/re_types/definitions/rerun/archetypes/image.fbs index 9b3f228b43bd..5a719f2281b1 100644 --- a/crates/re_types/definitions/rerun/archetypes/image.fbs +++ b/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 diff --git a/crates/re_types/source_hash.txt b/crates/re_types/source_hash.txt index d5c329b415ea..6e569e2f6d28 100644 --- a/crates/re_types/source_hash.txt +++ b/crates/re_types/source_hash.txt @@ -1,4 +1,4 @@ # This is a sha256 hash for all direct and indirect dependencies of this crate's build script. # It can be safely removed at anytime to force the build script to run again. # Check out build.rs to see how it's computed. -2659734e396eb3a29b890308e9c0bf91052d12df5453c23aceba57fef0dcda83 +4f8f4b875d35d8123e972eb88bb46266519d389aa815aac1dec2861e16cb8781 diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/image.py b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/image.py index adb8d4b586dc..9bdbc0e80734 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/image.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/image.py @@ -24,6 +24,24 @@ class Image(Archetype): - A `HxWx4` tensor, treated as an RGBA image. The viewer has limited support for ignoring extra empty dimensions. + + Example + ------- + ```python + + import numpy as np + import rerun as rr + import rerun.experimental as rr2 + + # Create an image with Pillow + 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)) + ``` """ data: components.TensorDataArray = field(metadata={"component": "primary"}, converter=image_data_converter) diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_buffer.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_buffer.py index 882a9a8475b2..e641b5959a4f 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_buffer.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_buffer.py @@ -7,4 +7,5 @@ def tensorbuffer_inner_converter(inner: npt.ArrayLike) -> npt.NDArray[Any]: + # A tensor buffer is always a flat array return np.asarray(inner).flatten() diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py index 36027d6ffa08..f0aae2ae5ea8 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/_overrides/tensor_data.py @@ -47,6 +47,10 @@ def _to_numpy(tensor: Tensor) -> npt.NDArray[Any]: # Init overrides ################################################################################ +# TODO(jleibs): Should also provide custom converters for shape / buffer +# assignment that prevent the user from putting the TensorData into an +# inconsistent state. + def tensordata_init( self: TensorData, @@ -57,6 +61,34 @@ def tensordata_init( array: Tensor | None = None, names: Sequence[str] | None = None, ) -> None: + """ + Construct a `TensorData` object. + + The `TensorData` object is internally represented by three fields: `id`, `shape`, and `buffer`. + + This constructor provides additional arguments 'array', and 'names'. When passing in a + multi-dimensional array such as a `np.ndarray`, the `shape` and `buffer` fields will be + populated automagically. + + Parameters + ---------- + self: TensorData + The TensorData object to construct. + id: TensorIdLike | None + The id of the tensor. If None, a random id will be generated. + shape: Sequence[TensorDimensionLike] | None + The shape of the tensor. If None, and an array is proviced, the shape will be inferred + from the shape of the array. + buffer: TensorBufferLike | None + The buffer of the tensor. If None, and an array is provided, the buffer will be generated + from the array. + array: Tensor | None + A numpy array (or The array of the tensor. If None, the array will be inferred from the buffer. + names: Sequence[str] | None + The names of the tensor dimensions when generating the shape from an array. + """ + # TODO(jleibs): Need to figure out how to get the above docstring to show up in the TensorData class + # documentation. if array is None and buffer is None: raise ValueError("Must provide one of 'array' or 'buffer'") if array is not None and buffer is not None: @@ -138,16 +170,18 @@ def tensordata_init( def tensordata_native_to_pa_array(data: TensorDataArrayLike, data_type: pa.DataType) -> pa.Array: from .. import TensorData - # If it's a sequence, grab the first one + # If it's a sequence of a single TensorData, grab the first one if isinstance(data, collections.abc.Sequence): - if len(data) != 1: - raise ValueError("Tensors do not support batches") - data = data[0] + if len(data) > 0: + if isinstance(data[0], TensorData): + if len(data) > 1: + raise ValueError("Tensors do not support batches") + data = data[0] # If it's not a TensorData, it should be an NDArray-like. coerce it into TensorData with the # constructor. if not isinstance(data, TensorData): - array = _to_numpy(data) + array = _to_numpy(data) # type: ignore[arg-type] data = TensorData(array=array) # Now build the actual arrow fields From fa535b9675112d48f366cb52b476de7aa2a8beba Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Sat, 2 Sep 2023 02:46:30 +0200 Subject: [PATCH 09/12] Use npt.ArrayLike --- .../definitions/rerun/datatypes/tensor_data.fbs | 4 ++-- crates/re_types/source_hash.txt | 2 +- .../rerun/_rerun2/datatypes/tensor_data.py | 16 +--------------- 3 files changed, 4 insertions(+), 18 deletions(-) diff --git a/crates/re_types/definitions/rerun/datatypes/tensor_data.fbs b/crates/re_types/definitions/rerun/datatypes/tensor_data.fbs index 988a19b768f7..d0f88a8bf085 100644 --- a/crates/re_types/definitions/rerun/datatypes/tensor_data.fbs +++ b/crates/re_types/definitions/rerun/datatypes/tensor_data.fbs @@ -19,8 +19,8 @@ namespace rerun.datatypes; /// which stores a contiguous array of typed values. table TensorData ( order: 100, - "attr.python.array_aliases": "npt.NDArray[np.uint8] | npt.NDArray[np.uint16] | npt.NDArray[np.uint32] | npt.NDArray[np.uint64] | npt.NDArray[np.int8] | npt.NDArray[np.int16] | npt.NDArray[np.int32] | npt.NDArray[np.int64] | npt.NDArray[np.float32] | npt.NDArray[np.float64]", - "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); diff --git a/crates/re_types/source_hash.txt b/crates/re_types/source_hash.txt index 6e569e2f6d28..4d48905e11b4 100644 --- a/crates/re_types/source_hash.txt +++ b/crates/re_types/source_hash.txt @@ -1,4 +1,4 @@ # This is a sha256 hash for all direct and indirect dependencies of this crate's build script. # It can be safely removed at anytime to force the build script to run again. # Check out build.rs to see how it's computed. -4f8f4b875d35d8123e972eb88bb46266519d389aa815aac1dec2861e16cb8781 +e3f01db86ef5464b4d836d3de4ff38cb9a0e07d3a3bc7d8aec47479b86949c6a diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py index eb5c848d33cf..066aab73ec2a 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/tensor_data.py @@ -4,7 +4,6 @@ from typing import Sequence, Union -import numpy as np import numpy.typing as npt import pyarrow as pa from attrs import define, field @@ -55,20 +54,7 @@ def __init__(self, *args, **kwargs): # type: ignore[no-untyped-def] TensorDataLike = TensorData -TensorDataArrayLike = Union[ - TensorData, - Sequence[TensorDataLike], - npt.NDArray[np.uint8] - | npt.NDArray[np.uint16] - | npt.NDArray[np.uint32] - | npt.NDArray[np.uint64] - | npt.NDArray[np.int8] - | npt.NDArray[np.int16] - | npt.NDArray[np.int32] - | npt.NDArray[np.int64] - | npt.NDArray[np.float32] - | npt.NDArray[np.float64], -] +TensorDataArrayLike = Union[TensorData, Sequence[TensorDataLike], npt.ArrayLike] # --- Arrow support --- From 71e14f3ada3bf1713dc92a74a88d148f33f7017c Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Sat, 2 Sep 2023 02:47:50 +0200 Subject: [PATCH 10/12] Fix code example comment --- docs/code-examples/image_simple_v2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/code-examples/image_simple_v2.py b/docs/code-examples/image_simple_v2.py index d4f7639be779..d0d22bfc977c 100644 --- a/docs/code-examples/image_simple_v2.py +++ b/docs/code-examples/image_simple_v2.py @@ -4,7 +4,7 @@ import rerun as rr import rerun.experimental as rr2 -# Create an image with Pillow +# 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) From 85540951ca96991ec46bc5f4c7c83daafb79c938 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Sat, 2 Sep 2023 02:59:51 +0200 Subject: [PATCH 11/12] Mypy --- crates/re_types/source_hash.txt | 2 +- rerun_py/rerun_sdk/rerun/_rerun2/archetypes/image.py | 2 +- rerun_py/tests/unit/test_image.py | 2 +- rerun_py/tests/unit/test_tensor.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/re_types/source_hash.txt b/crates/re_types/source_hash.txt index 4d48905e11b4..6fa0b1285378 100644 --- a/crates/re_types/source_hash.txt +++ b/crates/re_types/source_hash.txt @@ -1,4 +1,4 @@ # This is a sha256 hash for all direct and indirect dependencies of this crate's build script. # It can be safely removed at anytime to force the build script to run again. # Check out build.rs to see how it's computed. -e3f01db86ef5464b4d836d3de4ff38cb9a0e07d3a3bc7d8aec47479b86949c6a +2140a5f267ddd2ea601f42726b866571353223b77a6aa001c8743644438081fb diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/image.py b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/image.py index 9bdbc0e80734..18865963cc24 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/image.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/image.py @@ -33,7 +33,7 @@ class Image(Archetype): import rerun as rr import rerun.experimental as rr2 - # Create an image with Pillow + # 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) diff --git a/rerun_py/tests/unit/test_image.py b/rerun_py/tests/unit/test_image.py index b8559fb0ef19..69e183d38a94 100644 --- a/rerun_py/tests/unit/test_image.py +++ b/rerun_py/tests/unit/test_image.py @@ -12,7 +12,7 @@ RANDOM_IMAGE_SOURCE = rng.uniform(0.0, 1.0, (10, 20, 3)) -IMAGE_INPUTS: list[rrd.TensorDataArrayLike | None] = [ +IMAGE_INPUTS: list[rrd.TensorDataArrayLike] = [ # Full explicit construction rrd.TensorData( id=rrd.TensorId(uuid=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]), diff --git a/rerun_py/tests/unit/test_tensor.py b/rerun_py/tests/unit/test_tensor.py index 416694b8bf6c..eec614424d24 100644 --- a/rerun_py/tests/unit/test_tensor.py +++ b/rerun_py/tests/unit/test_tensor.py @@ -13,7 +13,7 @@ RANDOM_TENSOR_SOURCE = rng.uniform(0.0, 1.0, (8, 6, 3, 5)) -TENSOR_DATA_INPUTS: list[rrd.TensorDataArrayLike | None] = [ +TENSOR_DATA_INPUTS: list[rrd.TensorDataArrayLike] = [ # Full explicit construction rrd.TensorData( id=rrd.TensorId(uuid=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]), From f2af9f5d3a4d2155104f3dffdd50315ea9644202 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Mon, 4 Sep 2023 14:38:46 +0200 Subject: [PATCH 12/12] More mypy --- rerun_py/tests/unit/test_image.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rerun_py/tests/unit/test_image.py b/rerun_py/tests/unit/test_image.py index 69e183d38a94..84e220393a10 100644 --- a/rerun_py/tests/unit/test_image.py +++ b/rerun_py/tests/unit/test_image.py @@ -54,7 +54,7 @@ def test_image() -> None: compare_images(arch.data, expected, check_fields) -GOOD_IMAGE_INPUTS: list[rrd.TensorDataArrayLike | None] = [ +GOOD_IMAGE_INPUTS: list[rrd.TensorDataArrayLike] = [ # Mono rng.uniform(0.0, 1.0, (10, 20)), # RGB @@ -70,7 +70,7 @@ def test_image() -> None: rng.uniform(0.0, 1.0, (10, 20, 4, 1)), ] -BAD_IMAGE_INPUTS: list[rrd.TensorDataArrayLike | None] = [ +BAD_IMAGE_INPUTS: list[rrd.TensorDataArrayLike] = [ rng.uniform(0.0, 1.0, (10,)), rng.uniform(0.0, 1.0, (10, 20, 2)), rng.uniform(0.0, 1.0, (10, 20, 5)),