fix existing tests

bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_utils.py

@@ -1,15 +1,12 @@
 """This file contains utility functions helpful for using the NuScenes dataset."""
 
 import numpy as np
-from av2.geometry.geometry import compute_interior_points_mask
-from av2.utils.typing import NDArrayBool, NDArrayFloat
 from nuscenes.nuscenes import NuScenes
 from nuscenes.utils.data_classes import Box
 from pyquaternion import Quaternion
 
 from bucketed_scene_flow_eval.datastructures.se3 import SE3
 
-
 class InstanceBox(Box):
     def __init__(
         self,
@@ -27,8 +24,8 @@ def __init__(
         self.instance_token = instance_token
 
     def compute_interior_points(
-        self, points_xyz_m: NDArrayFloat, wlh_factor: float = 1.0
-    ) -> tuple[NDArrayFloat, NDArrayBool]:
+        self, points_xyz_m: np.ndarray[float], wlh_factor: float = 1.0
+    ) -> tuple[np.ndarray[float], np.ndarray[bool]]:
         """Given a query point cloud, filter to points interior to the cuboid, and provide mask.
 
         Note: comparison is to cuboid vertices in the destination reference frame.
@@ -157,3 +154,51 @@ def get_boxes_with_instance_token(self, sample_data_token: str) -> list[Instance
 
                 boxes.append(box)
         return boxes
+
+def compute_interior_points_mask(points_xyz: np.ndarray[float], cuboid_vertices: np.ndarray[float]) -> np.ndarray[bool]:
+    """Compute the interior points mask for the cuboid. Copied from av2.geometry.geometry.compute_interior_points_mask.
+
+    Reference: https://math.stackexchange.com/questions/1472049/check-if-a-point-is-inside-a-rectangular-shaped-area-3d
+
+            5------4
+            |\\    |\\
+            | \\   | \\
+            6--\\--7  \\
+            \\  \\  \\ \\
+        l    \\  1-------0    h
+         e    \\ ||   \\ ||   e
+          n    \\||    \\||   i
+           g    \\2------3    g
+            t      width.     h
+             h.               t.
+
+    Args:
+        points_xyz: (N,3) Array representing a point cloud in Cartesian coordinates (x,y,z).
+        cuboid_vertices: (8,3) Array representing 3D cuboid vertices, ordered as shown above.
+
+    Returns:
+        (N,) An array of boolean flags indicating whether the points are interior to the cuboid.
+    """
+    # Get three corners of the cuboid vertices.
+    vertices: np.ndarray[float] = np.stack((cuboid_vertices[6], cuboid_vertices[3], cuboid_vertices[1]))  # (3,3)
+
+    # Choose reference vertex.
+    # vertices and choice of ref_vertex are coupled.
+    ref_vertex = cuboid_vertices[2]  # (3,)
+
+    # Compute orthogonal edges of the cuboid.
+    uvw = ref_vertex - vertices  # (3,3)
+
+    # Compute signed values which are proportional to the distance from the vector.
+    sim_uvw_points = points_xyz @ uvw.transpose()  # (N,3)
+    sim_uvw_ref = uvw @ ref_vertex  # (3,)
+
+    # Only care about the diagonal.
+    sim_uvw_vertices: np.ndarray[float] = np.diag(uvw @ vertices.transpose())  # type: ignore # (3,)
+
+    # Check 6 conditions (2 for each of the 3 orthogonal directions).
+    # Refer to the linked reference for additional information.
+    constraint_a = np.logical_and(sim_uvw_ref <= sim_uvw_points, sim_uvw_points <= sim_uvw_vertices)
+    constraint_b = np.logical_and(sim_uvw_ref >= sim_uvw_points, sim_uvw_points >= sim_uvw_vertices)
+    is_interior: np.ndarray[bool] = np.logical_or(constraint_a, constraint_b).all(axis=1)
+    return is_interior

tests/datasets/nuscenes/nuscenes_tests.py

@@ -1,23 +1,20 @@
 from pathlib import Path
-from typing import Optional
 
-import numpy as np
 import pytest
-import tqdm
 
-from bucketed_scene_flow_eval.datasets.nuscenes import NuScenesLoader
+from bucketed_scene_flow_eval.datasets.nuscenes import NuScenesRawSequenceLoader
 
 
 @pytest.fixture
-def nuscenes_loader() -> NuScenesLoader:
-    return NuScenesLoader(
+def nuscenes_loader() -> NuScenesRawSequenceLoader:
+    return NuScenesRawSequenceLoader(
         sequence_dir=Path("/tmp/nuscenes"),
         version="v1.0-mini",
         verbose=False,
     )
 
 
-def test_nuscenes_loader_basic_load_and_len_check(nuscenes_loader: NuScenesLoader):
+def test_nuscenes_loader_basic_load_and_len_check(nuscenes_loader: NuScenesRawSequenceLoader):
     assert len(nuscenes_loader) > 0, f"no sequences found in {nuscenes_loader}"
     expected_lens = [236, 239, 236, 236, 233, 223, 239, 231, 231, 228]
     assert len(nuscenes_loader) == len(
@@ -35,11 +32,3 @@ def test_nuscenes_loader_basic_load_and_len_check(nuscenes_loader: NuScenesLoade
     assert num_loop_iterations == len(
         expected_lens
     ), f"expected {len(expected_lens)} loop iterations, got {num_loop_iterations}"
-
-
-def test_nuscenes_first_sequence_investigation(nuscenes_loader: NuScenesLoader):
-    sequence_id = nuscenes_loader.get_sequence_ids()[0]
-    nusc_seq = nuscenes_loader.load_sequence(sequence_id)
-    assert len(nusc_seq) > 0, f"no frames found in {sequence_id}"
-
-    first_frame = nusc_seq.load(0, 0)