wehub-resource-sync
112 lines
3.2 KiB
Python
112 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
import itertools
|
|
from typing import TYPE_CHECKING, cast
|
|
|
|
import numpy as np
|
|
import rerun as rr
|
|
from rerun.components import Position3DBatch, Vector3DBatch
|
|
|
|
from .common_arrays import (
|
|
class_ids_arrays,
|
|
class_ids_expected,
|
|
colors_arrays,
|
|
colors_expected,
|
|
labels_arrays,
|
|
labels_expected,
|
|
radii_arrays,
|
|
radii_expected,
|
|
vec3ds_arrays,
|
|
vec3ds_expected,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from rerun.datatypes import (
|
|
ClassIdArrayLike,
|
|
Float32ArrayLike,
|
|
Float64ArrayLike,
|
|
Rgba32ArrayLike,
|
|
Utf8ArrayLike,
|
|
Vec3DArrayLike,
|
|
)
|
|
|
|
|
|
def test_arrows3d() -> None:
|
|
vectors_arrays = vec3ds_arrays
|
|
origins_arrays = vec3ds_arrays
|
|
|
|
all_arrays = itertools.zip_longest(
|
|
vectors_arrays,
|
|
origins_arrays,
|
|
radii_arrays,
|
|
colors_arrays,
|
|
labels_arrays,
|
|
class_ids_arrays,
|
|
)
|
|
|
|
for vectors, origins, radii, colors, labels, class_ids in all_arrays:
|
|
vectors = vectors if vectors is not None else vectors_arrays[-1]
|
|
origins = origins if origins is not None else origins_arrays[-1]
|
|
|
|
# make Pyright happy as it's apparently not able to track typing info trough zip_longest
|
|
vectors = cast("Vec3DArrayLike", vectors)
|
|
origins = cast("Vec3DArrayLike | None", origins)
|
|
radii = cast("Float32ArrayLike | None", radii)
|
|
colors = cast("Rgba32ArrayLike | None", colors)
|
|
labels = cast("Utf8ArrayLike | None", labels)
|
|
class_ids = cast("ClassIdArrayLike | None", class_ids)
|
|
|
|
print(
|
|
f"E: rr.Arrows3D(\n"
|
|
f" vectors={vectors}\n"
|
|
f" origins={origins!r}\n"
|
|
f" radii={radii!r}\n"
|
|
f" colors={colors!r}\n"
|
|
f" labels={labels!r}\n"
|
|
f" class_ids={class_ids!r}\n"
|
|
f")",
|
|
)
|
|
arch = rr.Arrows3D(
|
|
vectors=vectors,
|
|
origins=origins,
|
|
radii=radii,
|
|
colors=colors,
|
|
labels=labels,
|
|
class_ids=class_ids,
|
|
)
|
|
print(f"A: {arch}\n")
|
|
|
|
assert arch.vectors == vec3ds_expected(vectors, Vector3DBatch)
|
|
assert arch.origins == vec3ds_expected(origins, Position3DBatch)
|
|
assert arch.radii == radii_expected(radii)
|
|
assert arch.colors == colors_expected(colors)
|
|
assert arch.labels == labels_expected(labels)
|
|
assert arch.class_ids == class_ids_expected(class_ids)
|
|
|
|
|
|
CASES: list[tuple[Float64ArrayLike, Float64ArrayLike]] = [
|
|
(
|
|
[],
|
|
[],
|
|
),
|
|
(
|
|
[[1.0, 1.0, 1.0]],
|
|
[[[1.0, 1.0, 1.0]]],
|
|
),
|
|
([[[1.0, 1.0, 1.0]] for _ in range(3)], [[[1.0, 1.0, 1.0]] for _ in range(3)]),
|
|
([[2.1, 3.2, 4.3] for _ in range(333)], [[[2.1, 3.2, 4.3]] for _ in range(333)]),
|
|
(np.ones((30, 3)), np.ones((30, 1, 3)).tolist()),
|
|
(np.ones((3, 1, 3)), np.ones((3, 1, 3)).tolist()),
|
|
(np.ones((1, 3)), np.ones((1, 1, 3)).tolist()),
|
|
]
|
|
|
|
|
|
def test_arrows3d_columnar() -> None:
|
|
for input, expected in CASES:
|
|
data = [*rr.Arrows3D.columns(vectors=input, origins=np.zeros_like(input))]
|
|
assert np.allclose(np.asarray(data[0].as_arrow_array().to_pylist()), np.asarray(expected))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_arrows3d()
|