From 6ea7dbf934dd53f8403ceca4d0c5e2929550cdf5 Mon Sep 17 00:00:00 2001 From: Zak Murez <33580232+zmurez@users.noreply.github.com> Date: Fri, 11 Nov 2022 18:40:36 -0800 Subject: [PATCH] NuScenes dataparser (#939) * nuscenes dataparser * cleanup * style * cleanup * cleanup * cleanup Co-authored-by: zmurez --- .../data/datamanagers/base_datamanager.py | 2 + .../data/dataparsers/nuscenes_dataparser.py | 212 ++++++++++++++++++ pyproject.toml | 1 + scripts/datasets/process_nuscenes_masks.py | 131 +++++++++++ 4 files changed, 346 insertions(+) create mode 100644 nerfstudio/data/dataparsers/nuscenes_dataparser.py create mode 100644 scripts/datasets/process_nuscenes_masks.py diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 1beb29963..5ea8c8329 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -40,6 +40,7 @@ from nerfstudio.data.dataparsers.instant_ngp_dataparser import ( InstantNGPDataParserConfig, ) from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig +from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserConfig from nerfstudio.data.dataparsers.record3d_dataparser import Record3DDataParserConfig from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.pixel_samplers import PixelSampler @@ -61,6 +62,7 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of "blender-data": BlenderDataParserConfig(), "friends-data": FriendsDataParserConfig(), "instant-ngp-data": InstantNGPDataParserConfig(), + "nuscenes-data": NuScenesDataParserConfig(), "record3d-data": Record3DDataParserConfig(), "dnerf-data": DNeRFDataParserConfig(), }, diff --git a/nerfstudio/data/dataparsers/nuscenes_dataparser.py b/nerfstudio/data/dataparsers/nuscenes_dataparser.py new file mode 100644 index 000000000..88e215611 --- /dev/null +++ b/nerfstudio/data/dataparsers/nuscenes_dataparser.py @@ -0,0 +1,212 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Data parser for NuScenes dataset""" +import math +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional, Tuple, Type + +import numpy as np +import torch +from nuscenes.nuscenes import NuScenes as NuScenesDatabase +from typing_extensions import Literal + +from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.data.dataparsers.base_dataparser import ( + DataParser, + DataParserConfig, + DataparserOutputs, +) +from nerfstudio.data.scene_box import SceneBox +from nerfstudio.utils.colmap_utils import qvec2rotmat + + +def rotation_translation_to_pose(r_vec, t_vec): + """Convert quaternion rotation and translation vectors to 4x4 matrix""" + + pose = np.eye(4) + pose[:3, :3] = qvec2rotmat(r_vec) + pose[:3, 3] = t_vec + return pose + + +@dataclass +class NuScenesDataParserConfig(DataParserConfig): + """NuScenes dataset config. + NuScenes (https://www.nuscenes.org/nuscenes) is an autonomous driving dataset containing 1000 20s clips. + Each clip was recorded with a suite of sensors including 6 souround cameras. + It also includes 3D cuboid annotations around obects. + We optionally use these cuboids to mask dynamic objects by specifying the mask_dir flag. + To create these masks use scripts/datasets/process_nuscenes_masks.py. + """ + + _target: Type = field(default_factory=lambda: NuScenes) + """target class to instantiate""" + data: Path = Path("scene-0103") # TODO: rename to scene but keep checkpoint saving name? + """Name of the scene.""" + data_dir: Path = Path("/mnt/local/NuScenes") + """Path to NuScenes dataset.""" + version: Literal["v1.0-mini", "v1.0-trainval"] = "v1.0-mini" + """Dataset version.""" + cameras: Tuple[Literal["FRONT", "FRONT_LEFT", "FRONT_RIGHT", "BACK", "BACK_LEFT", "BACK_RIGHT"], ...] = ("FRONT",) + """Which cameras to use.""" + mask_dir: Optional[Path] = None + """Path to masks of dynamic objects.""" + + train_split_percentage: float = 0.9 + """The percent of images to use for training. The remaining images are for eval.""" + + verbose: bool = False + """Load dataset with verbose messaging""" + + +@dataclass +class NuScenes(DataParser): + """NuScenes DatasetParser""" + + config: NuScenesDataParserConfig + + def _generate_dataparser_outputs(self, split="train"): + # pylint: disable=too-many-statements + + nusc = NuScenesDatabase(version=self.config.version, dataroot=self.config.data_dir, verbose=self.config.verbose) + cameras = ["CAM_" + camera for camera in self.config.cameras] + + assert ( + len(cameras) == 1 + ), "waiting on multiple camera support" # TODO: remove once multiple cameras are supported + + # get samples for scene + samples = [ + samp for samp in nusc.sample if nusc.get("scene", samp["scene_token"])["name"] == str(self.config.data) + ] + + # sort by timestamp (only to make chronological viz easier) + samples.sort(key=lambda x: (x["scene_token"], x["timestamp"])) + + transform1 = np.array( + [ + [0, -1, 0, 0], + [0, 0, -1, 0], + [1, 0, 0, 0], + [0, 0, 0, 1], + ] + ) + transform2 = np.array( + [ + [0, 0, 1, 0], + [0, 1, 0, 0], + [-1, 0, 0, 0], + [0, 0, 0, 1], + ] + ) + + # get image filenames and camera data + image_filenames = [] + mask_filenames = [] + mask_dir = self.config.mask_dir if self.config.mask_dir is not None else Path("") + intrinsics = [] + poses = [] + for sample in samples: + for camera in cameras: + camera_data = nusc.get("sample_data", sample["data"][camera]) + calibrated_sensor_data = nusc.get("calibrated_sensor", camera_data["calibrated_sensor_token"]) + ego_pose_data = nusc.get("ego_pose", camera_data["ego_pose_token"]) + + ego_pose = rotation_translation_to_pose(ego_pose_data["rotation"], ego_pose_data["translation"]) + cam_pose = rotation_translation_to_pose( + calibrated_sensor_data["rotation"], calibrated_sensor_data["translation"] + ) + pose = ego_pose @ cam_pose + + # rotate to opencv frame + pose = transform1 @ pose + + # convert from opencv camera to nerfstudio camera + pose[0:3, 1:3] *= -1 + pose = pose[np.array([1, 0, 2, 3]), :] + pose[2, :] *= -1 + + # rotate to z-up in viewer + pose = transform2 @ pose + + image_filenames.append(self.config.data_dir / camera_data["filename"]) + mask_filenames.append( + mask_dir / "masks" / camera / os.path.split(camera_data["filename"])[1].replace("jpg", "png") + ) + intrinsics.append(calibrated_sensor_data["camera_intrinsic"]) + poses.append(pose) + poses = torch.from_numpy(np.stack(poses).astype(np.float32)) + intrinsics = torch.from_numpy(np.array(intrinsics).astype(np.float32)) + + # center poses + poses[:, :3, 3] -= poses[:, :3, 3].mean(dim=0) + + # scale poses + poses[:, :3, 3] /= poses[:, :3, 3].abs().max() + + # filter image_filenames and poses based on train/eval split percentage + num_snapshots = len(samples) + num_train_snapshots = math.ceil(num_snapshots * self.config.train_split_percentage) + num_eval_snapshots = num_snapshots - num_train_snapshots + i_all = np.arange(num_snapshots) + i_train = np.linspace( + 0, num_snapshots - 1, num_train_snapshots, dtype=int + ) # equally spaced training snapshots starting and ending at 0 and num_images-1 + i_eval = np.setdiff1d(i_all, i_train) # eval images are the remaining images + assert len(i_eval) == num_eval_snapshots + i_train = (i_train[None, :] * len(cameras) + np.arange(len(cameras))[:, None]).ravel() + i_eval = (i_eval[None, :] * len(cameras) + np.arange(len(cameras))[:, None]).ravel() + if split == "train": + indices = i_train + elif split in ["val", "test"]: + indices = i_eval + else: + raise ValueError(f"Unknown dataparser split {split}") + + # Choose image_filenames and poses based on split, but after auto orient and scaling the poses. + image_filenames = [image_filenames[i] for i in indices] + mask_filenames = [mask_filenames[i] for i in indices] + intrinsics = intrinsics[indices] + poses = poses[indices] + + # in x,y,z order + # assumes that the scene is centered at the origin + aabb_scale = 1.0 + scene_box = SceneBox( + aabb=torch.tensor( + [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32 + ) + ) + + cameras = Cameras( + fx=intrinsics[:, 0, 0], + fy=intrinsics[:, 1, 1], + cx=intrinsics[:, 0, 2], + cy=intrinsics[:, 1, 2], + height=900, + width=1600, + camera_to_worlds=poses[:, :3, :4], + camera_type=CameraType.PERSPECTIVE, + ) + + dataparser_outputs = DataparserOutputs( + image_filenames=image_filenames, + cameras=cameras, + scene_box=scene_box, + mask_filenames=mask_filenames if self.config.mask_dir is not None else None, + ) + return dataparser_outputs diff --git a/pyproject.toml b/pyproject.toml index 146897261..59a74727e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ dependencies = [ "torchtyping>=0.1.4", "torchvision>=0.13.0", "u-msgpack-python>=2.4.1", + "nuscenes-devkit>=1.1.1", "wandb>=0.13.3" ] diff --git a/scripts/datasets/process_nuscenes_masks.py b/scripts/datasets/process_nuscenes_masks.py new file mode 100644 index 000000000..cb9328f95 --- /dev/null +++ b/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()