diff --git a/bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_utils.py b/bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_utils.py index faf7550..7955fd0 100644 --- a/bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_utils.py +++ b/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 diff --git a/tests/datasets/nuscenes/nuscenes_tests.py b/tests/datasets/nuscenes/nuscenes_tests.py index e76aa2b..bb6988c 100644 --- a/tests/datasets/nuscenes/nuscenes_tests.py +++ b/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)