89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import itertools
|
|
from typing import TYPE_CHECKING, cast
|
|
|
|
import numpy as np
|
|
import rerun as rr
|
|
import rerun.blueprint as rrb
|
|
|
|
from .common_arrays import none_empty_or_value
|
|
|
|
if TYPE_CHECKING:
|
|
from rerun.blueprint.datatypes import TensorDimensionIndexSliderArrayLike
|
|
|
|
|
|
def test_tensor_slice_selection() -> None:
|
|
widths = [
|
|
None,
|
|
2,
|
|
rr.datatypes.TensorDimensionSelection(dimension=2, invert=False),
|
|
rr.components.TensorWidthDimension(dimension=2, invert=False),
|
|
]
|
|
heights = [
|
|
None,
|
|
3,
|
|
rr.datatypes.TensorDimensionSelection(dimension=3, invert=False),
|
|
rr.components.TensorHeightDimension(dimension=3, invert=False),
|
|
]
|
|
indices_arrays = [
|
|
[
|
|
rr.components.TensorDimensionIndexSelection(dimension=1, index=3),
|
|
rr.components.TensorDimensionIndexSelection(dimension=2, index=2),
|
|
rr.components.TensorDimensionIndexSelection(dimension=3, index=1),
|
|
],
|
|
None,
|
|
]
|
|
slider_arrays = [
|
|
None,
|
|
[1, 2, 3],
|
|
[
|
|
rrb.components.TensorDimensionIndexSlider(1),
|
|
rrb.components.TensorDimensionIndexSlider(2),
|
|
rrb.components.TensorDimensionIndexSlider(3),
|
|
],
|
|
np.array([1, 2, 3]),
|
|
]
|
|
|
|
all_arrays = itertools.zip_longest(
|
|
widths,
|
|
heights,
|
|
indices_arrays,
|
|
slider_arrays,
|
|
)
|
|
|
|
for width, height, indices, slider in all_arrays:
|
|
width = cast("rr.datatypes.TensorDimensionSelectionLike | None", width)
|
|
height = cast("rr.datatypes.TensorDimensionSelectionLike | None", height)
|
|
indices = cast("rr.datatypes.TensorDimensionIndexSelectionArrayLike | None", indices)
|
|
slider = cast("TensorDimensionIndexSliderArrayLike | None", slider)
|
|
|
|
print(
|
|
f"rr.TensorSliceSelection(\n"
|
|
f" width={width!r}\n"
|
|
f" height={height!r}\n"
|
|
f" indices={indices!r}\n"
|
|
f" slider={slider!r}\n"
|
|
f")",
|
|
)
|
|
arch = rrb.TensorSliceSelection(
|
|
width=width,
|
|
height=height,
|
|
indices=indices,
|
|
slider=slider,
|
|
)
|
|
print(f"{arch}\n")
|
|
|
|
assert arch.width == rr.components.TensorWidthDimensionBatch._converter(
|
|
none_empty_or_value(width, rr.components.TensorWidthDimension(dimension=2, invert=False)),
|
|
)
|
|
assert arch.height == rr.components.TensorHeightDimensionBatch._converter(
|
|
none_empty_or_value(height, rr.components.TensorHeightDimension(dimension=3, invert=False)),
|
|
)
|
|
assert arch.indices == rr.components.TensorDimensionIndexSelectionBatch._converter(
|
|
none_empty_or_value(indices, indices_arrays[0]),
|
|
)
|
|
assert arch.slider == rrb.components.TensorDimensionIndexSliderBatch._converter(
|
|
none_empty_or_value(slider, [1, 2, 3]),
|
|
)
|