Added support for PyTorch array to Boxes2D's array convenience argument

.github/workflows/contrib_rerun_py.yml

@@ -67,7 +67,7 @@ jobs:
         # TODO(jleibs): understand why deps can't be installed in the same step as the wheel
         shell: bash
         run: |
-          pip install attrs>=23.1.0 numpy>=1.23 pillow pyarrow==10.0.1 pytest==7.1.2 typing_extensions>=4.5
+          pip install attrs>=23.1.0 numpy>=1.23 pillow pyarrow==10.0.1 pytest==7.1.2 torch==2.1.0 typing_extensions>=4.5
 
       - name: Get version
         id: get-version

.github/workflows/reusable_build_and_test_wheels.yml

@@ -215,7 +215,7 @@ jobs:
         # TODO(jleibs): understand why deps can't be installed in the same step as the wheel
         shell: bash
         run: |
-          pip install attrs>=23.1.0 numpy>=1.23 pillow pyarrow==10.0.1 pytest==7.1.2 typing_extensions>=4.5
+          pip install attrs>=23.1.0 numpy>=1.23 pillow pyarrow==10.0.1 pytest==7.1.2 torch==2.1.0 typing_extensions>=4.5
 
       - name: Get version
         id: get-version

rerun_py/requirements-lint.txt

@@ -3,6 +3,7 @@ black==23.3.0
 blackdoc==0.3.8
 mypy==1.4.1
 numpy>=1.24 # For mypy plugin
+torch>=2.1.0
 pip-check-reqs==2.4.4 # Checks for missing deps in requirements.txt files
 pytest  # For mypy to work
 ruff==0.0.276

rerun_py/rerun_sdk/rerun/archetypes/boxes2d_ext.py

@@ -106,8 +106,10 @@ def __init__(
                 if centers is not None:
                     raise ValueError("Cannot specify both `array` and `centers` at the same time.")
 
+                if type(array) is not np.ndarray:
+                    array = np.array(array)
+
                 if np.any(array):
-                    array = np.asarray(array, dtype="float32")
                     if array.ndim == 1:
                         array = np.expand_dims(array, axis=0)
                 else:

rerun_py/tests/unit/common_arrays.py

@@ -3,6 +3,7 @@
 from typing import Any
 
 import numpy as np
+import torch
 from rerun.components import (
     ClassId,
     ClassIdBatch,
@@ -83,6 +84,8 @@ def none_empty_or_value(obj: Any, value: Any) -> Any:
     np.array([1, 2, 3, 4], dtype=np.float32),
     # Vec2DArrayLike: npt.NDArray[np.float32]
     np.array([1, 2, 3, 4], dtype=np.float32).reshape((2, 2, 1, 1, 1)),
+    # PyTorch array
+    torch.asarray([1, 2, 3, 4], dtype=torch.float32),
 ]
 
 
@@ -118,6 +121,8 @@ def vec2ds_expected(obj: Any, type_: Any | None = None) -> Any:
     np.array([1, 2, 3, 4, 5, 6], dtype=np.float32),
     # Vec3DArrayLike: npt.NDArray[np.float32]
     np.array([1, 2, 3, 4, 5, 6], dtype=np.float32).reshape((2, 3, 1, 1, 1)),
+    # PyTorch array
+    torch.asarray([1, 2, 3, 4, 5, 6], dtype=torch.float32),
 ]
 
 
@@ -153,6 +158,8 @@ def vec3ds_expected(obj: Any, type_: Any | None = None) -> Any:
     np.array([1, 2, 3, 4, 5, 6, 7, 8], dtype=np.float32),
     # Vec4DArrayLike: npt.NDArray[np.float32]
     np.array([1, 2, 3, 4, 5, 6, 7, 8], dtype=np.float32).reshape((2, 4, 1, 1, 1)),
+    # PyTorch array
+    torch.asarray([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.float32),
 ]
 
 

rerun_py/tests/unit/test_box2d.py

@@ -3,8 +3,11 @@
 import itertools
 from typing import Optional, cast
 
+import numpy as np
+import numpy.typing as npt
 import pytest
 import rerun as rr
+import torch
 from rerun.components import (
     DrawOrderLike,
     HalfSizes2DBatch,
@@ -137,6 +140,19 @@ def test_with_array_xcycw2h2() -> None:
     assert rr.Boxes2D(mins=[1, 1], sizes=[2, 4]) == rr.Boxes2D(array=[2, 3, 1, 2], array_format=rr.Box2DFormat.XCYCW2H2)
 
 
+@pytest.mark.parametrize(
+    "array",
+    [
+        [1, 2, 3, 4],
+        [1, 2, 3, 4],
+        np.array([1, 2, 3, 4], dtype=np.float32),
+        torch.asarray([1, 2, 3, 4], dtype=torch.float32),
+    ],
+)
+def test_with_array_types(array: npt.ArrayLike) -> None:
+    assert rr.Boxes2D(mins=[1, 2], sizes=[3, 4]) == rr.Boxes2D(array=array, array_format=rr.Box2DFormat.XYWH)
+
+
 def test_invalid_parameter_combinations() -> None:
     rr.set_strict_mode(True)