tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

NuScenes dataparser (#939)

Browse Source 
* nuscenes dataparser

* cleanup

* style

* cleanup

* cleanup

* cleanup

Co-authored-by: zmurez <zmurez@eng.ucsd.edu>
This commit is contained in:
Zak Murez 2022-11-11 18:40:36 -08:00 committed by GitHub
parent fdbeb2ae31
commit 6ea7dbf934
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 346 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
nerfstudio/data/datamanagers/base_datamanager.py
View File

@ -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(),
},

212
nerfstudio/data/dataparsers/nuscenes_dataparser.py Normal file
View File

@ -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

1
pyproject.toml
View File

@ -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"
]

131
scripts/datasets/process_nuscenes_masks.py Normal file
View File

@ -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()