From 9e9986004e50bdec2664f57216dcd3fa640d61f2 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 26 Jul 2023 01:18:25 +0200 Subject: [PATCH 1/8] Expose batch APIs for oriented bounding boxes --- rerun_py/rerun_sdk/rerun/__init__.py | 2 +- rerun_py/rerun_sdk/rerun/log/bounding_box.py | 130 ++++++++++++++++++- 2 files changed, 129 insertions(+), 3 deletions(-) diff --git a/rerun_py/rerun_sdk/rerun/__init__.py b/rerun_py/rerun_sdk/rerun/__init__.py index 0a93d705789a..8dcde3a38020 100644 --- a/rerun_py/rerun_sdk/rerun/__init__.py +++ b/rerun_py/rerun_sdk/rerun/__init__.py @@ -95,7 +95,7 @@ ) from .log.annotation import AnnotationInfo, ClassDescription, log_annotation_context from .log.arrow import log_arrow -from .log.bounding_box import log_obb +from .log.bounding_box import log_obb, log_obbs from .log.camera import log_pinhole from .log.clear import log_cleared from .log.extension_components import log_extension_components diff --git a/rerun_py/rerun_sdk/rerun/log/bounding_box.py b/rerun_py/rerun_sdk/rerun/log/bounding_box.py index b8d76e334aaa..41acdfd52e9a 100644 --- a/rerun_py/rerun_sdk/rerun/log/bounding_box.py +++ b/rerun_py/rerun_sdk/rerun/log/bounding_box.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any +from typing import Any, Sequence import numpy as np import numpy.typing as npt @@ -14,13 +14,22 @@ from rerun.components.quaternion import QuaternionArray from rerun.components.radius import RadiusArray from rerun.components.vec import Vec3DArray -from rerun.log import Color, _normalize_colors, _normalize_ids, _normalize_radii +from rerun.log import ( + Color, + Colors, + OptionalClassIds, + _normalize_colors, + _normalize_ids, + _normalize_labels, + _normalize_radii, +) from rerun.log.extension_components import _add_extension_components from rerun.log.log_decorator import log_decorator from rerun.recording_stream import RecordingStream __all__ = [ "log_obb", + "log_obbs", ] @@ -136,3 +145,120 @@ def log_obb( # Always the primary component last so range-based queries will include the other data. See(#1215) if instanced: bindings.log_arrow_msg(entity_path, components=instanced, timeless=timeless, recording=recording) + + +@log_decorator +def log_obbs( + entity_path: str, + *, + half_sizes: npt.ArrayLike | None, + positions: npt.ArrayLike | None = None, + rotations_q: npt.ArrayLike | None = None, + colors: Color | Colors | None = None, + stroke_widths: npt.ArrayLike | None = None, + labels: Sequence[str] | None = None, + class_ids: OptionalClassIds | None = None, + ext: dict[str, Any] | None = None, + timeless: bool = False, + recording: RecordingStream | None = None, +) -> None: + """ + Log a 3D Oriented Bounding Box, or OBB. + + Example: + -------- + ``` + rr.log_obb("my_obb", half_size=[1.0, 2.0, 3.0], position=[0, 0, 0], rotation_q=[0, 0, 0, 1]) + ``` + + Parameters + ---------- + entity_path: + The path to the oriented bounding box in the space hierarchy. + half_sizes: + Nx3 Array. Each row is the [x, y, z] half dimensions of an OBB. + positions: + Optional Nx3 array. Each row is [x, y, z] positions of an OBB in world space. + rotations_q: + Optional Nx3 array. Each row is quaternion coordinates [x, y, z, w] for the rotation from model to world space. + colors: + Optional Nx3 or Nx4 array. Each row is RGB or RGBA in sRGB gamma-space as either 0-1 floats or 0-255 integers, + with separate alpha. + stroke_widths: + Optional array of the width of the line edges. + labels: + Optional array of text labels placed at `position`. + class_ids: + Optional array of class id for the OBBs. The class id provides colors and labels if not specified explicitly. + ext: + Optional dictionary of extension components. See [rerun.log_extension_components][] + timeless: + If true, the bounding box will be timeless (default: False). + recording: + Specifies the [`rerun.RecordingStream`][] to use. + If left unspecified, defaults to the current active data recording, if there is one. + See also: [`rerun.init`][], [`rerun.set_global_data_recording`][]. + + """ + recording = RecordingStream.to_native(recording) + + colors = _normalize_colors(colors) + radii = _normalize_radii(stroke_widths) + radii / 2 + labels = _normalize_labels(labels) + class_ids = _normalize_ids(class_ids) + + # 0 = instanced, 1 = splat + comps = [{}, {}] # type: ignore[var-annotated] + + if half_sizes is not None: + half_sizes = np.require(half_sizes, dtype="float32") + + if half_sizes.shape[1] == 3: + comps[0]["rerun.box3d"] = Box3DArray.from_numpy(half_sizes) + else: + raise TypeError("half_size should be Nx3") + + if positions is not None: + positions = np.require(positions, dtype="float32") + + if positions.shape[1] == 3: + comps[0]["rerun.vec3d"] = Vec3DArray.from_numpy(positions) + else: + raise TypeError("position should be 1x3") + + if rotations_q is not None: + rotations_q = np.require(rotations_q, dtype="float32") + + if rotations_q.shape[1] == 4: + comps[0]["rerun.quaternion"] = QuaternionArray.from_numpy(rotations_q) + else: + raise TypeError("rotation should be 1x4") + + if len(colors): + is_splat = len(colors.shape) == 1 + if is_splat: + colors = colors.reshape(1, len(colors)) + comps[is_splat]["rerun.colorrgba"] = ColorRGBAArray.from_numpy(colors) + + if len(radii): + is_splat = len(radii) == 1 + comps[is_splat]["rerun.radius"] = RadiusArray.from_numpy(radii) + + if len(labels): + is_splat = len(labels) == 1 + comps[is_splat]["rerun.label"] = LabelArray.new(labels) + + if len(class_ids): + is_splat = len(class_ids) == 1 + comps[is_splat]["rerun.class_id"] = ClassIdArray.from_numpy(class_ids) + + if ext: + _add_extension_components(comps[0], comps[1], ext, None) + + if comps[1]: + comps[1]["rerun.instance_key"] = InstanceArray.splat() + bindings.log_arrow_msg(entity_path, components=comps[1], timeless=timeless, recording=recording) + + # Always the primary component last so range-based queries will include the other data. See(#1215) + bindings.log_arrow_msg(entity_path, components=comps[0], timeless=timeless, recording=recording) From 4e70ceead8a3452771ea4a0b34f63a4c7aa71681 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 26 Jul 2023 01:28:12 +0200 Subject: [PATCH 2/8] Handle empty lists for obb half_width / position / rotations --- rerun_py/rerun_sdk/rerun/components/box.py | 2 +- rerun_py/rerun_sdk/rerun/log/bounding_box.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/rerun_py/rerun_sdk/rerun/components/box.py b/rerun_py/rerun_sdk/rerun/components/box.py index 252d391e70e9..31fc87f39ec6 100644 --- a/rerun_py/rerun_sdk/rerun/components/box.py +++ b/rerun_py/rerun_sdk/rerun/components/box.py @@ -15,7 +15,7 @@ class Box3DArray(pa.ExtensionArray): # type: ignore[misc] def from_numpy(array: npt.NDArray[np.float32]) -> Box3DArray: """Build a `Box3DArray` from an Nx3 numpy array.""" - assert array.shape[1] == 3 + assert len(array) == 0 or array.shape[1] == 3 storage = pa.FixedSizeListArray.from_arrays(array.flatten(), type=Box3DType.storage_type) # TODO(john) enable extension type wrapper # return cast(Box3DArray, pa.ExtensionArray.from_storage(Box3DType(), storage)) diff --git a/rerun_py/rerun_sdk/rerun/log/bounding_box.py b/rerun_py/rerun_sdk/rerun/log/bounding_box.py index 41acdfd52e9a..7fea76e0e334 100644 --- a/rerun_py/rerun_sdk/rerun/log/bounding_box.py +++ b/rerun_py/rerun_sdk/rerun/log/bounding_box.py @@ -214,7 +214,7 @@ def log_obbs( if half_sizes is not None: half_sizes = np.require(half_sizes, dtype="float32") - if half_sizes.shape[1] == 3: + if len(half_sizes) == 0 or half_sizes.shape[1] == 3: comps[0]["rerun.box3d"] = Box3DArray.from_numpy(half_sizes) else: raise TypeError("half_size should be Nx3") @@ -222,7 +222,7 @@ def log_obbs( if positions is not None: positions = np.require(positions, dtype="float32") - if positions.shape[1] == 3: + if len(positions) == 0 or positions.shape[1] == 3: comps[0]["rerun.vec3d"] = Vec3DArray.from_numpy(positions) else: raise TypeError("position should be 1x3") @@ -230,7 +230,7 @@ def log_obbs( if rotations_q is not None: rotations_q = np.require(rotations_q, dtype="float32") - if rotations_q.shape[1] == 4: + if len(rotations_q) == 0 or rotations_q.shape[1] == 4: comps[0]["rerun.quaternion"] = QuaternionArray.from_numpy(rotations_q) else: raise TypeError("rotation should be 1x4") From f0946f5478865581d464a36727ad43be4c400730 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 26 Jul 2023 01:33:23 +0200 Subject: [PATCH 3/8] Handle empty arrays --- rerun_py/rerun_sdk/rerun/components/quaternion.py | 2 +- rerun_py/rerun_sdk/rerun/components/vec.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/rerun_py/rerun_sdk/rerun/components/quaternion.py b/rerun_py/rerun_sdk/rerun/components/quaternion.py index 443b63e9abcf..e6c0b5cb705e 100644 --- a/rerun_py/rerun_sdk/rerun/components/quaternion.py +++ b/rerun_py/rerun_sdk/rerun/components/quaternion.py @@ -29,7 +29,7 @@ def __array__(self) -> npt.NDArray[np.float32]: class QuaternionArray(pa.ExtensionArray): # type: ignore[misc] def from_numpy(array: npt.NDArray[np.float32]) -> QuaternionArray: """Build a `QuaternionArray` from an Nx4 numpy array.""" - assert array.shape[1] == 4 + assert len(array) == 0 or array.shape[1] == 4 storage = pa.FixedSizeListArray.from_arrays(array.flatten(), type=QuaternionType.storage_type) # TODO(john) enable extension type wrapper # return cast(QuaternionArray, pa.ExtensionArray.from_storage(QuaternionType(), storage)) diff --git a/rerun_py/rerun_sdk/rerun/components/vec.py b/rerun_py/rerun_sdk/rerun/components/vec.py index 5c222159cc5c..eeca3fb9a338 100644 --- a/rerun_py/rerun_sdk/rerun/components/vec.py +++ b/rerun_py/rerun_sdk/rerun/components/vec.py @@ -17,7 +17,7 @@ class Vec2DArray(pa.ExtensionArray): # type: ignore[misc] def from_numpy(array: npt.NDArray[np.float32]) -> Vec2DArray: """Build a `Vec2DArray` from an Nx2 numpy array.""" - assert array.shape[1] == 2 + assert len(array) == 0 or array.shape[1] == 2 storage = pa.FixedSizeListArray.from_arrays(array.flatten(), type=Vec2DType.storage_type) # TODO(john) enable extension type wrapper # return cast(Vec2DArray, pa.ExtensionArray.from_storage(Vec2DType(), storage)) @@ -32,7 +32,7 @@ def from_numpy(array: npt.NDArray[np.float32]) -> Vec2DArray: class Vec3DArray(pa.ExtensionArray): # type: ignore[misc] def from_numpy(array: npt.NDArray[np.float32]) -> Vec3DArray: """Build a `Vec3DArray` from an Nx3 numpy array.""" - assert array.shape[1] == 3 + assert len(array) == 0 or array.shape[1] == 3 storage = pa.FixedSizeListArray.from_arrays(array.flatten(), type=Vec3DType.storage_type) # TODO(john) enable extension type wrapper # return cast(Vec3DArray, pa.ExtensionArray.from_storage(Vec3DType(), storage)) From 660ae870f5c7c6c57a7601cab0db905724f74408 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 26 Jul 2023 01:46:35 +0200 Subject: [PATCH 4/8] add log_obbs to __all__ --- rerun_py/rerun_sdk/rerun/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/rerun_py/rerun_sdk/rerun/__init__.py b/rerun_py/rerun_sdk/rerun/__init__.py index 8dcde3a38020..a87572852d70 100644 --- a/rerun_py/rerun_sdk/rerun/__init__.py +++ b/rerun_py/rerun_sdk/rerun/__init__.py @@ -50,6 +50,7 @@ "log_mesh_file", "log_meshes", "log_obb", + "log_obbs", "log_path", "log_pinhole", "log_point", From f1543930d413744c55df297f8d234b5d1fe1d7ca Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 26 Jul 2023 17:24:52 +0200 Subject: [PATCH 5/8] Box3D batch example --- docs/code-examples/box3d_batch.py | 21 +++++++++++++++++++++ docs/content/reference/data_types/box3d.md | 12 ++++++++++++ 2 files changed, 33 insertions(+) create mode 100644 docs/code-examples/box3d_batch.py diff --git a/docs/code-examples/box3d_batch.py b/docs/code-examples/box3d_batch.py new file mode 100644 index 000000000000..7688b2510838 --- /dev/null +++ b/docs/code-examples/box3d_batch.py @@ -0,0 +1,21 @@ +"""Log a batch of oriented bounding boxes.""" +import rerun as rr +from scipy.spatial.transform import Rotation + +rr.init("box3d", spawn=True) + +rr.log_annotation_context( + "/", + [ + rr.ClassDescription(info=rr.AnnotationInfo(1, "red", (255, 0, 0))), + rr.ClassDescription(info=rr.AnnotationInfo(2, "green", (0, 255, 0))), + ], +) +rr.log_obbs( + "batch", + half_sizes=[[2.0, 2.0, 1.0], [1.0, 1.0, 0.5]], + rotations_q=[Rotation.from_euler("xyz", [0, 0, 0]).as_quat(), Rotation.from_euler("xyz", [0, 0, 45]).as_quat()], + positions=[[2, 0, 0], [-2, 0, 0]], + stroke_widths=0.05, + class_ids=[2, 1], +) diff --git a/docs/content/reference/data_types/box3d.md b/docs/content/reference/data_types/box3d.md index c0254eee9978..7517f89321bf 100644 --- a/docs/content/reference/data_types/box3d.md +++ b/docs/content/reference/data_types/box3d.md @@ -30,3 +30,15 @@ code-example: box3d_simple + +## Batch Example + +code-example: box3d_batch + + + + + + + + From 0d30a8cf852f7af0aa4528ce8d01e0055541b413 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 26 Jul 2023 17:26:50 +0200 Subject: [PATCH 6/8] Fix stroke_width --- docs/content/reference/data_types/box3d.md | 10 +++++----- rerun_py/rerun_sdk/rerun/log/bounding_box.py | 4 ++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/content/reference/data_types/box3d.md b/docs/content/reference/data_types/box3d.md index 7517f89321bf..6e0bdade5db0 100644 --- a/docs/content/reference/data_types/box3d.md +++ b/docs/content/reference/data_types/box3d.md @@ -36,9 +36,9 @@ code-example: box3d_simple code-example: box3d_batch - - - - - + + + + + diff --git a/rerun_py/rerun_sdk/rerun/log/bounding_box.py b/rerun_py/rerun_sdk/rerun/log/bounding_box.py index 7fea76e0e334..72d5413d9ce8 100644 --- a/rerun_py/rerun_sdk/rerun/log/bounding_box.py +++ b/rerun_py/rerun_sdk/rerun/log/bounding_box.py @@ -203,8 +203,8 @@ def log_obbs( recording = RecordingStream.to_native(recording) colors = _normalize_colors(colors) - radii = _normalize_radii(stroke_widths) - radii / 2 + stroke_widths = _normalize_radii(stroke_widths) + radii = stroke_widths / 2 labels = _normalize_labels(labels) class_ids = _normalize_ids(class_ids) From a272eee6004f9f6569ce21b16a1a4400c674c73f Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 26 Jul 2023 17:32:57 +0200 Subject: [PATCH 7/8] Split lines --- docs/code-examples/box3d_batch.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/code-examples/box3d_batch.py b/docs/code-examples/box3d_batch.py index 7688b2510838..02e3071078e4 100644 --- a/docs/code-examples/box3d_batch.py +++ b/docs/code-examples/box3d_batch.py @@ -14,7 +14,10 @@ rr.log_obbs( "batch", half_sizes=[[2.0, 2.0, 1.0], [1.0, 1.0, 0.5]], - rotations_q=[Rotation.from_euler("xyz", [0, 0, 0]).as_quat(), Rotation.from_euler("xyz", [0, 0, 45]).as_quat()], + rotations_q=[ + Rotation.from_euler("xyz", [0, 0, 0]).as_quat(), + Rotation.from_euler("xyz", [0, 0, 45]).as_quat(), + ], positions=[[2, 0, 0], [-2, 0, 0]], stroke_widths=0.05, class_ids=[2, 1], From e7eb45dea4b5db81406f57af13bcb59d8890e2bb Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 26 Jul 2023 18:09:29 +0200 Subject: [PATCH 8/8] Update the source hash --- crates/re_types/source_hash.txt | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/crates/re_types/source_hash.txt b/crates/re_types/source_hash.txt index 0b1e6fd2e451..9fa1862b0818 100644 --- a/crates/re_types/source_hash.txt +++ b/crates/re_types/source_hash.txt @@ -1,8 +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. -<<<<<<< HEAD -89d2d068457a4278ed2299d9226df927fd795ff4096c5f078ef701eba7f4dae2 -======= -9340c2d668f1bb96d364568f660e6a16b2144f2580f7bd5dc8bfae88a728517a ->>>>>>> main +b1231093ce5784a399aa1c02f0b2ad6563b044d1420b4553cbd85830b8365817