NuScenes dataparser (#939)

* nuscenes dataparser

* cleanup

* style

* cleanup

* cleanup

* cleanup

Co-authored-by: zmurez <zmurez@eng.ucsd.edu>

scripts/datasets/process_nuscenes_masks.py

@@ -0,0 +1,131 @@
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Tuple
+
+import cv2
+import numpy as np
+import tyro
+from nuscenes.nuscenes import NuScenes as NuScenesDatabase
+from nuscenes.utils.data_classes import Box
+from nuscenes.utils.geometry_utils import BoxVisibility, view_points
+from tqdm import tqdm
+from typing_extensions import Literal
+
+
+@dataclass
+class ProcessNuScenesMasks:
+    """Use cuboid detections to render masks for dynamic objects."""
+
+    data_dir: Path
+    """Path to NuScenes dataset."""
+    output_dir: Path
+    """Path to the output directory."""
+    version: Literal["v1.0-mini", "v1.0-trainval"] = "v1.0-mini"
+    """Which version of the dataset to process."""
+    velocity_thresh: float = 0.75
+    """Minimum speed for object to be considered dynamic."""
+    cameras: Tuple[Literal["FRONT", "FRONT_LEFT", "FRONT_RIGHT", "BACK", "BACK_LEFT", "BACK_RIGHT"], ...] = (
+        "FRONT",
+        "FRONT_LEFT",
+        "FRONT_RIGHT",
+        "BACK",
+        "BACK_LEFT",
+        "BACK_RIGHT",
+    )
+    """Which cameras to use."""
+
+    verbose: bool = False
+    """If True, print extra logging and visualize images on screen."""
+
+    def main(self) -> None:
+        """Generate NuScenes dynamic object masks."""
+
+        nusc = NuScenesDatabase(version=self.version, dataroot=self.data_dir, verbose=self.verbose)
+        cameras = ["CAM_" + camera for camera in self.cameras]
+
+        for camera in cameras:
+            (self.output_dir / "masks" / camera).mkdir(parents=True, exist_ok=True)
+
+        # get samples for scene
+        samples = [samp for samp in nusc.sample]
+
+        # sort by timestamp (only to make chronological viz easier)
+        samples.sort(key=lambda x: (x["scene_token"], x["timestamp"]))
+
+        # get which instances are moving in any frame (these are what we mask)
+        instances = nusc.instance
+        for instance in instances:
+            is_dynamic = False
+            ann_token = instance["first_annotation_token"]
+            while ann_token:
+                velocity = nusc.box_velocity(ann_token)
+                if not np.linalg.norm(velocity) < self.velocity_thresh:
+                    is_dynamic = True
+                    break
+                ann_token = nusc.get("sample_annotation", ann_token)["next"]
+            instance["is_dynamic"] = is_dynamic
+        instances_is_dynamic = {instance["token"]: instance["is_dynamic"] for instance in instances}
+
+        for sample in tqdm(samples):
+            viz = []
+            for camera in cameras:
+                camera_data = nusc.get("sample_data", sample["data"][camera])
+                calibrated_sensor = nusc.get("calibrated_sensor", camera_data["calibrated_sensor_token"])
+                intrinsics = np.array(calibrated_sensor["camera_intrinsic"])
+
+                _, boxes, _ = nusc.get_sample_data(sample["data"][camera], box_vis_level=BoxVisibility.ANY)
+                # TODO: BoxVisibility.ANY misses boxes that are partially behind the camera leading to missed masks
+                # Instead use BoxVisibility.NONE and make sure to rasterize box faces correctly
+
+                mask = np.ones((900, 1600), dtype=np.uint8)
+                for box in boxes:
+
+                    # Dont mask out static objects (static in all frames)
+                    instance_token = nusc.get("sample_annotation", box.token)["instance_token"]
+                    if not instances_is_dynamic[instance_token]:
+                        continue
+
+                    # project box to image plane and rasterize each face
+                    corners_3d = box.corners()
+                    corners = view_points(corners_3d, intrinsics, normalize=True)[:2, :]
+                    corners = np.round(corners).astype(int).T
+                    cv2.fillPoly(mask, [corners[[0, 1, 2, 3]]], 0)  # front
+                    cv2.fillPoly(mask, [corners[[4, 5, 6, 7]]], 0)  # back
+                    cv2.fillPoly(mask, [corners[[0, 1, 5, 4]]], 0)  # top
+                    cv2.fillPoly(mask, [corners[[2, 3, 7, 6]]], 0)  # bottom
+                    cv2.fillPoly(mask, [corners[[0, 3, 7, 4]]], 0)  # left
+                    cv2.fillPoly(mask, [corners[[1, 2, 6, 5]]], 0)  # right
+
+                maskname = os.path.split(camera_data["filename"])[1].replace("jpg", "png")
+                cv2.imwrite(str(self.output_dir / "masks" / camera / maskname), mask * 255)
+
+                if self.verbose:
+                    img = cv2.imread(str(self.data_dir / camera_data["filename"]))
+                    mask = ~mask.astype(bool)
+                    img[mask, :] -= np.minimum(img[mask, :], 100)
+                    viz.append(img)
+
+            if self.verbose:
+                if len(viz) == 6:
+                    viz = np.vstack((np.hstack(viz[:3]), np.hstack(viz[3:])))
+                    viz = cv2.resize(viz, (int(1600 * 3 / 3), int(900 * 2 / 3)))
+                elif len(viz) == 3:
+                    viz = np.hstack(viz[:3])
+                    viz = cv2.resize(viz, (int(1600 * 3 / 3), int(900 / 3)))
+                elif len(viz) == 1:
+                    viz = viz[0]
+                else:
+                    raise ValueError("Only support 1 or 3 or 6 cameras for viz")
+                cv2.imshow("", viz)
+                cv2.waitKey(1)
+
+
+def entrypoint():
+    """Entrypoint for use with pyproject scripts."""
+    tyro.extras.set_accent_color("bright_yellow")
+    tyro.cli(ProcessNuScenesMasks).main()
+
+
+if __name__ == "__main__":
+    entrypoint()