v2.0.25: Updated vis infra for orbbec data

bucketed_scene_flow_eval/datasets/orbbec_astra/dataset.py

@@ -6,6 +6,7 @@
 
 from bucketed_scene_flow_eval.datastructures import (
     SE3,
+    ColoredSupervisedPointCloudFrame,
     EgoLidarFlow,
     PointCloud,
     PoseInfo,
@@ -16,7 +17,7 @@
 )
 from bucketed_scene_flow_eval.eval import EmptyEvaluator, Evaluator
 from bucketed_scene_flow_eval.interfaces import AbstractDataset, LoaderType
-from bucketed_scene_flow_eval.utils import load_by_extension
+from bucketed_scene_flow_eval.utils import load_feather, load_pickle
 
 
 class OrbbecAstra(AbstractDataset):
@@ -25,7 +26,7 @@ def __init__(
         root_dir: Path,
         flow_dir: Path | None,
         subsequence_length: int = 2,
-        extension_name: str = ".pcd",
+        extension_name: str = ".pkl",
     ) -> None:
         root_dir = Path(root_dir)
         self.data_dir = root_dir
@@ -41,12 +42,22 @@ def __init__(
             len(self.pointclouds) >= subsequence_length
         ), f"Need at least {subsequence_length} frames, found {len(self.pointclouds)} in {root_dir}/*{extension_name}"
 
+        # Magic numbers to scale the pointclouds to be in the same range as the argoverse data we are training on
+        # These numbers are derived from looking at the pointclouds themselves
+
+        self.scale = 35
+        self.center_translation = -np.array([1.13756592, 0.21126675, -1.04425789])
+        self.bg_delete_x_max = (1.7 + self.center_translation[0]) * self.scale
+        self.bg_delete_z_min = (-1.2 + self.center_translation[2]) * self.scale
+
     def __len__(self):
         return len(self.pointclouds) - self.subsequence_length + 1
 
-    def _load_file(self, pcd_file: Path) -> PointCloud:
-        data = load_by_extension(pcd_file, verbose=False)
-        return PointCloud(data)
+    def _load_file(self, data_file: Path) -> tuple[PointCloud, np.ndarray]:
+        data = load_pickle(data_file, verbose=False)
+        points = data[:, :3]
+        colors = data[:, 3:]
+        return PointCloud(points), colors
 
     def evaluator(self) -> Evaluator:
         return EmptyEvaluator()
@@ -55,6 +66,7 @@ def loader_type(self):
         return LoaderType.NON_CAUSAL
 
     def _load_pose_info(self) -> PoseInfo:
+        # Convert from standard sensor coordinate system to right hand coordinate system we use.
         sensor_to_right_hand = SE3(
             # fmt: off
             rotation_matrix=np.array([[0, 0, 1],
@@ -64,6 +76,7 @@ def _load_pose_info(self) -> PoseInfo:
             translation=np.array([0.0, 0.0, 0.0]),
         )
 
+        # The sensor is rotated down 30 degrees, so we need to rotate it back up so the table is level.
         theta_degrees = 30
         theta_radians = np.radians(theta_degrees)
         rotation_matrix_y = np.array(
@@ -79,22 +92,46 @@ def _load_pose_info(self) -> PoseInfo:
         )
 
         return PoseInfo(
-            sensor_to_ego=right_hand_to_ego.compose(sensor_to_right_hand),
+            sensor_to_ego=right_hand_to_ego.compose(sensor_to_right_hand)
+            .translate(self.center_translation)
+            .scale(self.scale),
             ego_to_global=SE3.identity(),
         )
 
+    def _load_flow(self, idx: int, pc_frame: SupervisedPointCloudFrame) -> EgoLidarFlow:
+        if self.flow_dir is None:
+            return EgoLidarFlow(full_flow=np.zeros_like(pc_frame.full_pc), mask=pc_frame.mask)
+        flow_file = self.flow_dir / f"{idx:010d}.feather"
+        if not flow_file.exists():
+            return EgoLidarFlow(full_flow=np.zeros_like(pc_frame.full_pc), mask=pc_frame.mask)
+        flow_feather = load_feather(flow_file, verbose=False)
+        flow_x = flow_feather["flow_tx_m"].to_numpy()
+        flow_y = flow_feather["flow_ty_m"].to_numpy()
+        flow_z = flow_feather["flow_tz_m"].to_numpy()
+        flow = np.stack([flow_x, flow_y, flow_z], axis=-1)
+        flow_mask = flow_feather["is_valid"].to_numpy()
+        assert len(flow) == len(
+            pc_frame.full_pc
+        ), f"Expected {len(pc_frame.full_pc)} points, found {len(flow)}"
+        assert np.all(
+            flow_mask == pc_frame.mask
+        ), f"Founds {np.sum(flow_mask)} masked points, expected {np.sum(pc_frame.mask)}"
+        return EgoLidarFlow(full_flow=flow, mask=pc_frame.mask)
+
     def _get_sequence_frame(self, idx: int) -> TimeSyncedSceneFlowFrame:
-        pc = self.pointclouds[idx]
+        pc, color = self.pointclouds[idx]
 
         semantics = np.zeros(pc.shape[0], dtype=SemanticClassId)
         pose_info = self._load_pose_info()
 
         ego_pc = pc.transform(pose_info.sensor_to_ego)
 
-        is_valid_mask = (ego_pc.points[:, 0] < 1.7) & (ego_pc.points[:, 2] > -1.2)
-        pc_frame = SupervisedPointCloudFrame(pc, pose_info, is_valid_mask, semantics)
+        is_valid_mask = (ego_pc.points[:, 0] < self.bg_delete_x_max) & (
+            ego_pc.points[:, 2] > self.bg_delete_z_min
+        )
+        pc_frame = ColoredSupervisedPointCloudFrame(pc, pose_info, is_valid_mask, semantics, color)
         rgb_frame_lookup = RGBFrameLookup.empty()
-        gt_flow = EgoLidarFlow(full_flow=np.zeros_like(pc.points), mask=is_valid_mask)
+        gt_flow = self._load_flow(idx, pc_frame)
 
         return TimeSyncedSceneFlowFrame(
             pc=pc_frame,

bucketed_scene_flow_eval/datastructures/__init__.py

@@ -1,6 +1,7 @@
 from .camera_projection import CameraModel, CameraProjection
 from .dataclasses import (
     BoundingBox,
+    ColoredSupervisedPointCloudFrame,
     EgoLidarDistance,
     EgoLidarFlow,
     MaskArray,
@@ -18,6 +19,7 @@
     TimeSyncedSceneFlowFrame,
     VectorArray,
 )
+from .line_mesh import LineMesh
 from .o3d_visualizer import O3DVisualizer
 from .pointcloud import PointCloud, from_fixed_array, to_fixed_array
 from .rgb_image import RGBImage, RGBImageCrop
@@ -37,6 +39,9 @@
     "SemanticClassId",
     "SemanticClassIdArray",
     "SupervisedPointCloudFrame",
+    "ColoredSupervisedPointCloudFrame",
+    "TimeSyncedSceneFlowBoxFrame",
+    "TimeSyncedSceneFlowFrame",
     "TimeSyncedAVLidarData",
     "TimeSyncedBaseAuxilaryData",
     "TimeSyncedRawFrame",

bucketed_scene_flow_eval/datastructures/dataclasses.py

@@ -234,6 +234,26 @@ def mask_points(self, mask: MaskArray) -> "SupervisedPointCloudFrame":
         )
 
 
+@dataclass
+class ColoredSupervisedPointCloudFrame(SupervisedPointCloudFrame):
+    colors: np.ndarray
+
+    def __post_init__(self):
+        assert isinstance(
+            self.colors, np.ndarray
+        ), f"colors must be an ndarray, got {type(self.colors)}"
+        assert self.colors.ndim == 2, f"colors must be a 2D array, got {self.colors.ndim}"
+        assert self.colors.shape[1] == 3, f"colors must have 3 columns, got {self.colors.shape[1]}"
+        assert len(self.colors) == len(
+            self.full_pc
+        ), f"colors must have the same length as full_pc, got {len(self.colors)} and {len(self.full_pc)}"
+        assert self.colors.dtype == np.float32, f"colors must be float32, got {self.colors.dtype}"
+        # Ensure all colors are between 0 and 1
+        assert np.all(
+            (self.colors >= 0) & (self.colors <= 1)
+        ), f"colors must be between 0 and 1, got min {np.min(self.colors)} and max {np.max(self.colors)}"
+
+
 @dataclass
 class RGBFrame:
     rgb: RGBImage

bucketed_scene_flow_eval/datastructures/line_mesh.py

@@ -0,0 +1,92 @@
+import numpy as np
+import open3d as o3d
+
+
+def align_vector_to_another(a=np.array([0, 0, 1]), b=np.array([1, 0, 0])):
+    """
+    Aligns vector a to vector b with axis angle rotation
+    """
+    if np.array_equal(a, b):
+        return None, None
+    axis_ = np.cross(a, b)
+    axis_ = axis_ / np.linalg.norm(axis_)
+    angle = np.arccos(np.dot(a, b))
+
+    return axis_, angle
+
+
+def normalized(a, axis=-1, order=2):
+    """Normalizes a numpy array of points"""
+    l2 = np.atleast_1d(np.linalg.norm(a, order, axis))
+    l2[l2 == 0] = 1
+    return a / np.expand_dims(l2, axis), l2
+
+
+class LineMesh(object):
+    def __init__(self, points, lines=None, colors=[0, 1, 0], radius=0.15):
+        """Creates a line represented as sequence of cylinder triangular meshes
+
+        Arguments:
+            points {ndarray} -- Numpy array of ponts Nx3.
+
+        Keyword Arguments:
+            lines {list[list] or None} -- List of point index pairs denoting line segments. If None, implicit lines from ordered pairwise points. (default: {None})
+            colors {list} -- list of colors, or single color of the line (default: {[0, 1, 0]})
+            radius {float} -- radius of cylinder (default: {0.15})
+        """
+        self.points = np.array(points)
+        self.lines = (
+            np.array(lines) if lines is not None else self.lines_from_ordered_points(self.points)
+        )
+        self.colors = np.array(colors)
+        self.radius = radius
+        self.cylinder_segments = []
+
+        self._create_line_mesh()
+
+    @staticmethod
+    def lines_from_ordered_points(points):
+        lines = [[i, i + 1] for i in range(0, points.shape[0] - 1, 1)]
+        return np.array(lines)
+
+    def _create_line_mesh(self):
+        first_points = self.points[self.lines[:, 0], :]
+        second_points = self.points[self.lines[:, 1], :]
+        line_segments = second_points - first_points
+        line_segments_unit, line_lengths = normalized(line_segments)
+
+        z_axis = np.array([0, 0, 1])
+        # Create triangular mesh cylinder segments of line
+        for i in range(line_segments_unit.shape[0]):
+            line_segment = line_segments_unit[i, :]
+            line_length = line_lengths[i]
+            # get axis angle rotation to allign cylinder with line segment
+            axis, angle = align_vector_to_another(z_axis, line_segment)
+            # Get translation vector
+            translation = first_points[i, :] + line_segment * line_length * 0.5
+            # create cylinder and apply transformations
+            cylinder_segment = o3d.geometry.TriangleMesh.create_cylinder(self.radius, line_length)
+            cylinder_segment = cylinder_segment.translate(translation, relative=False)
+            if axis is not None:
+                axis_a = axis * angle
+                cylinder_segment = cylinder_segment.rotate(
+                    R=o3d.geometry.get_rotation_matrix_from_axis_angle(axis_a),
+                    center=cylinder_segment.get_center(),
+                )
+                # cylinder_segment = cylinder_segment.rotate(
+                #   axis_a, center=True, type=o3d.geometry.RotationType.AxisAngle)
+            # color cylinder
+            color = self.colors if self.colors.ndim == 1 else self.colors[i, :]
+            cylinder_segment.paint_uniform_color(color)
+
+            self.cylinder_segments.append(cylinder_segment)
+
+    def add_line(self, vis):
+        """Adds this line to the visualizer"""
+        for cylinder in self.cylinder_segments:
+            vis.add_geometry(cylinder)
+
+    def remove_line(self, vis):
+        """Removes this line from the visualizer"""
+        for cylinder in self.cylinder_segments:
+            vis.remove_geometry(cylinder)

bucketed_scene_flow_eval/datastructures/o3d_visualizer.py

@@ -6,15 +6,19 @@
 import open3d as o3d
 
 from .dataclasses import EgoLidarFlow, PointCloudFrame, RGBFrame, VectorArray
+from .line_mesh import LineMesh
 from .pointcloud import PointCloud
 from .se3 import SE3
 
 ColorType = Union[np.ndarray, tuple[float, float, float], list[tuple[float, float, float]]]
 
 
 class O3DVisualizer:
-    def __init__(self, point_size: float = 0.1, add_world_frame: bool = True):
+    def __init__(
+        self, point_size: float = 0.1, line_width: float = 1.0, add_world_frame: bool = True
+    ):
         self.point_size = point_size
+        self.line_width = line_width
         self.geometry_list = []
 
         if add_world_frame:
@@ -192,16 +196,14 @@ def add_trajectory(
         for i in range(len(trajectory) - 1):
             self.add_sphere(trajectory[i], radius, color)
 
-        # Add line set between trajectory points
-        line_set = o3d.geometry.LineSet()
-        line_set.points = o3d.utility.Vector3dVector(trajectory)
-        line_set.lines = o3d.utility.Vector2iVector(
+        points = o3d.utility.Vector3dVector(trajectory)
+        lines = o3d.utility.Vector2iVector(
             np.array([[i, i + 1] for i in range(len(trajectory) - 1)])
         )
-        line_set.colors = o3d.utility.Vector3dVector(
-            np.tile(np.array(color), (len(trajectory) - 1, 1))
-        )
-        self.add_geometry(line_set)
+        colors = o3d.utility.Vector3dVector(np.tile(np.array(color), (len(trajectory) - 1, 1)))
+
+        line_mesh = LineMesh(points=points, lines=lines, colors=colors, radius=self.line_width / 20)
+        self.add_geometry(line_mesh.cylinder_segments)
 
     def render(self, vis, reset_view: bool = True):
         for geometry in self.geometry_list:
@@ -210,7 +212,9 @@ def render(self, vis, reset_view: bool = True):
     def run(self, vis=o3d.visualization.Visualizer()):
         print("Running visualizer on geometry list of length", len(self.geometry_list))
         vis.create_window(window_name="Benchmark Visualizer")
-        vis.get_render_option().point_size = self.point_size
+
+        ro = vis.get_render_option()
+        ro.point_size = self.point_size
 
         self.render(vis)
 

bucketed_scene_flow_eval/datastructures/se3.py

@@ -33,9 +33,11 @@ def identity() -> "SE3":
 
     @staticmethod
     def from_rot_x_y_z_translation_x_y_z(rx, ry, rz, tx, ty, tz) -> "SE3":
-        rotation_matrix = Quaternion(axis=[1, 0, 0], angle=rx).rotation_matrix @ Quaternion(
-            axis=[0, 1, 0], angle=ry
-        ).rotation_matrix @ Quaternion(axis=[0, 0, 1], angle=rz).rotation_matrix
+        rotation_matrix = (
+            Quaternion(axis=[1, 0, 0], angle=rx).rotation_matrix
+            @ Quaternion(axis=[0, 1, 0], angle=ry).rotation_matrix
+            @ Quaternion(axis=[0, 0, 1], angle=rz).rotation_matrix
+        )
         translation = np.array([tx, ty, tz])
         return SE3(rotation_matrix, translation)
 
@@ -64,6 +66,13 @@ def translate(self, translation: np.ndarray) -> "SE3":
             translation=self.translation + translation,
         )
 
+    def scale(self, scale: float) -> "SE3":
+        """Return a new SE3 instance with the given scale applied."""
+        return SE3(
+            rotation_matrix=self.rotation_matrix * scale,
+            translation=self.translation * scale,
+        )
+
     def transform_points(self, point_cloud: np.ndarray) -> np.ndarray:
         """Apply the SE(3) transformation to this point cloud.
         Args:

pyproject.toml

@@ -15,7 +15,7 @@ exclude = [
 
 [project]
 name = "bucketed_scene_flow_eval"
-version = "2.0.24"
+version = "2.0.25"
 authors = [
   { name="Kyle Vedder", email="[email protected]" },
 ]