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

Merge branch 'main' into generic-model-interface

Browse Source
This commit is contained in:
Jonáš Kulhánek 2023-02-27 01:52:00 +01:00 committed by GitHub
commit 8d24a3bcab
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 377 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -35,6 +35,9 @@ from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
from nerfstudio.cameras.cameras import CameraType
from nerfstudio.cameras.rays import RayBundle
from nerfstudio.configs.base_config import InstantiateConfig
from nerfstudio.data.dataparsers.arkitscenes_dataparser import (
ARKitScenesDataParserConfig,
)
from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs
from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig
@ -48,6 +51,7 @@ from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserCo
from nerfstudio.data.dataparsers.phototourism_dataparser import (
PhototourismDataParserConfig,
)
from nerfstudio.data.dataparsers.scannet_dataparser import ScanNetDataParserConfig
from nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig
from nerfstudio.data.dataparsers.sitcoms3d_dataparser import Sitcoms3DDataParserConfig
from nerfstudio.data.datasets.base_dataset import InputDataset
@ -73,12 +77,14 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of
{
"nerfstudio-data": NerfstudioDataParserConfig(),
"minimal-parser": MinimalDataParserConfig(),
"arkit-data": ARKitScenesDataParserConfig(),
"blender-data": BlenderDataParserConfig(),
"instant-ngp-data": InstantNGPDataParserConfig(),
"nuscenes-data": NuScenesDataParserConfig(),
"dnerf-data": DNeRFDataParserConfig(),
"phototourism-data": PhototourismDataParserConfig(),
"dycheck-data": DycheckDataParserConfig(),
"scannet-data": ScanNetDataParserConfig(),
"sdfstudio-data": SDFStudioDataParserConfig(),
"sitcoms3d-data": Sitcoms3DDataParserConfig(),
},

211
nerfstudio/data/dataparsers/arkitscenes_dataparser.py Normal file
View File

@ -0,0 +1,211 @@
# 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 ARKitScenes dataset"""
import math
from dataclasses import dataclass, field
from pathlib import Path
from typing import Type
import cv2
import numpy as np
import torch
from nerfstudio.cameras.cameras import Cameras, CameraType
from nerfstudio.data.dataparsers.base_dataparser import (
DataParser,
DataParserConfig,
DataparserOutputs,
)
from nerfstudio.data.scene_box import SceneBox
# Taken from https://github.com/apple/ARKitScenes/blob/main/threedod/benchmark_scripts/utils/tenFpsDataLoader.py
def traj_string_to_matrix(traj_string: str):
"""convert traj_string into translation and rotation matrices
Args:
traj_string: A space-delimited file where each line represents a camera position at a particular timestamp.
The file has seven columns:
* Column 1: timestamp
* Columns 2-4: rotation (axis-angle representation in radians)
* Columns 5-7: translation (usually in meters)
Returns:
ts: translation matrix
Rt: rotation matrix
"""
tokens = traj_string.split()
assert len(tokens) == 7
ts = tokens[0]
# Rotation in angle axis
angle_axis = [float(tokens[1]), float(tokens[2]), float(tokens[3])]
r_w_to_p, _ = cv2.Rodrigues(np.asarray(angle_axis))
# Translation
t_w_to_p = np.asarray([float(tokens[4]), float(tokens[5]), float(tokens[6])])
extrinsics = np.eye(4, 4)
extrinsics[:3, :3] = r_w_to_p
extrinsics[:3, -1] = t_w_to_p
Rt = np.linalg.inv(extrinsics)
return (ts, Rt)
@dataclass
class ARKitScenesDataParserConfig(DataParserConfig):
"""ARKitScenes dataset config.
ARKitScenes dataset (http://github.com/apple/ARKitScenes) is a large-scale 3D dataset of indoor scenes.
This dataparser uses 3D deteciton subset of the ARKitScenes dataset.
"""
_target: Type = field(default_factory=lambda: ARKitScenes)
"""target class to instantiate"""
data: Path = Path("data/ARKitScenes/3dod/Validation/41069021")
"""Path to ARKitScenes folder with densely extracted scenes."""
scene_scale: float = 1.0
"""How much to scale the region of interest by."""
center_poses: bool = True
"""Whether to center the poses."""
scale_poses: bool = True
"""Whether to automatically scale the poses to fit in +/- 1 bounding box."""
train_split_fraction: float = 0.9
"""The fraction of images to use for training. The remaining images are for eval."""
depth_unit_scale_factor: float = 1e-3
"""Scales the depth values to meters. Default value is 0.001 for a millimeter to meter conversion."""
@dataclass
class ARKitScenes(DataParser):
"""ARKitScenes DatasetParser"""
config: ARKitScenesDataParserConfig
def _generate_dataparser_outputs(self, split="train"):
video_id = self.config.data.name
image_dir = self.config.data / f"{video_id}_frames" / "lowres_wide"
depth_dir = self.config.data / f"{video_id}_frames" / "lowres_depth"
intrinsics_dir = self.config.data / f"{video_id}_frames" / "lowres_wide_intrinsics"
pose_file = self.config.data / f"{video_id}_frames" / "lowres_wide.traj"
frame_ids = [x.name for x in sorted(depth_dir.iterdir())]
frame_ids = [x.split(".png")[0].split("_")[1] for x in frame_ids]
frame_ids.sort()
poses_from_traj = {}
with open(pose_file, "r", encoding="utf-8") as f:
traj = f.readlines()
for line in traj:
poses_from_traj[f"{round(float(line.split(' ')[0]), 3):.3f}"] = np.array(
traj_string_to_matrix(line)[1].tolist()
)
image_filenames, depth_filenames, intrinsics, poses = [], [], [], []
w, h, _, _, _, _ = np.loadtxt(list(sorted(intrinsics_dir.iterdir()))[0]) # Get image size from first intrinsic
for frame_id in frame_ids:
intrinsic = self._get_intrinsic(intrinsics_dir, frame_id, video_id)
frame_pose = self._get_pose(frame_id, poses_from_traj)
intrinsics.append(intrinsic)
image_filenames.append(image_dir / f"{video_id}_{frame_id}.png")
depth_filenames.append(depth_dir / f"{video_id}_{frame_id}.png")
poses.append(frame_pose)
# filter image_filenames and poses based on train/eval split percentage
num_images = len(image_filenames)
num_train_images = math.ceil(num_images * self.config.train_split_fraction)
num_eval_images = num_images - num_train_images
i_all = np.arange(num_images)
i_train = np.linspace(
0, num_images - 1, num_train_images, dtype=int
) # equally spaced training images 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_images
if split == "train":
indices = i_train
elif split in ["val", "test"]:
indices = i_eval
else:
raise ValueError(f"Unknown dataparser split {split}")
poses = torch.from_numpy(np.stack(poses).astype(np.float32))
intrinsics = torch.from_numpy(np.stack(intrinsics).astype(np.float32))
if self.config.scale_poses:
poses[:, :3, 3] -= poses[:, :3, 3].mean(dim=0)
if self.config.scale_poses:
poses[:, :3, 3] /= poses[:, :3, 3].abs().max()
# 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]
depth_filenames = [depth_filenames[i] for i in indices] if len(depth_filenames) > 0 else []
intrinsics = intrinsics[indices.tolist()]
poses = poses[indices.tolist()]
# in x,y,z order
# assumes that the scene is centered at the origin
aabb_scale = self.config.scene_scale
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=int(h),
width=int(w),
camera_to_worlds=poses[:, :3, :4],
camera_type=CameraType.PERSPECTIVE,
)
dataparser_outputs = DataparserOutputs(
image_filenames=image_filenames,
cameras=cameras,
scene_box=scene_box,
metadata={
"depth_filenames": depth_filenames if len(depth_filenames) > 0 else None,
"depth_unit_scale_factor": self.config.depth_unit_scale_factor,
},
)
return dataparser_outputs
@staticmethod
def _get_intrinsic(intrinsics_dir: Path, frame_id: str, video_id: str):
intrinsic_fn = intrinsics_dir / f"{video_id}_{frame_id}.pincam"
if not intrinsic_fn.exists():
intrinsic_fn = intrinsics_dir / f"{video_id}_{float(frame_id) - 0.001:.3f}.pincam"
if not intrinsic_fn.exists():
intrinsic_fn = intrinsics_dir / f"{video_id}_{float(frame_id) + 0.001:.3f}.pincam"
_, _, fx, fy, hw, hh = np.loadtxt(intrinsic_fn)
intrinsic = np.asarray([[fx, 0, hw], [0, fy, hh], [0, 0, 1]])
return intrinsic
@staticmethod
def _get_pose(frame_id: str, poses_from_traj: dict):
if str(frame_id) in poses_from_traj:
frame_pose = np.array(poses_from_traj[str(frame_id)])
else:
for my_key in poses_from_traj:
if abs(float(frame_id) - float(my_key)) < 0.005:
frame_pose = np.array(poses_from_traj[str(my_key)])
frame_pose[0:3, 1:3] *= -1
return frame_pose

160
nerfstudio/data/dataparsers/scannet_dataparser.py Normal file
View File

@ -0,0 +1,160 @@
# 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 ScanNet dataset"""
import math
from dataclasses import dataclass, field
from pathlib import Path
from typing import Type
import cv2
import numpy as np
import torch
from nerfstudio.cameras.cameras import Cameras, CameraType
from nerfstudio.data.dataparsers.base_dataparser import (
DataParser,
DataParserConfig,
DataparserOutputs,
)
from nerfstudio.data.scene_box import SceneBox
@dataclass
class ScanNetDataParserConfig(DataParserConfig):
"""ScanNet dataset config.
ScanNet dataset (https://www.scan-net.org/) is a large-scale 3D dataset of indoor scenes.
This dataparser assumes that the dense stream was extracted from .sens files.
Expected structure of scene directory:
- color/
- depth/
- intrinsic/
- pose/
"""
_target: Type = field(default_factory=lambda: ScanNet)
"""target class to instantiate"""
data: Path = Path("data/scannet/scene0423_02")
"""Path to ScanNet folder with densely extracted scenes."""
scene_scale: float = 1.0
"""How much to scale the region of interest by."""
center_poses: bool = True
"""Whether to center the poses."""
scale_poses: bool = True
"""Whether to automatically scale the poses to fit in +/- 1 bounding box."""
train_split_fraction: float = 0.9
"""The fraction of images to use for training. The remaining images are for eval."""
depth_unit_scale_factor: float = 1e-3
"""Scales the depth values to meters. Default value is 0.001 for a millimeter to meter conversion."""
@dataclass
class ScanNet(DataParser):
"""ScanNet DatasetParser"""
config: ScanNetDataParserConfig
def _generate_dataparser_outputs(self, split="train"):
image_dir = self.config.data / "color"
depth_dir = self.config.data / "depth"
pose_dir = self.config.data / "pose"
img_dir_sorted = list(sorted(image_dir.iterdir(), key=lambda x: int(x.name.split(".")[0])))
depth_dir_sorted = list(sorted(depth_dir.iterdir(), key=lambda x: int(x.name.split(".")[0])))
pose_dir_sorted = list(sorted(pose_dir.iterdir(), key=lambda x: int(x.name.split(".")[0])))
first_img = cv2.imread(img_dir_sorted[0].as_posix())
h, w, _ = first_img.shape
image_filenames = []
depth_filenames = []
intrinsics = []
poses = []
K = np.loadtxt(self.config.data / "intrinsic" / "intrinsic_color.txt")
for img, depth, pose in zip(img_dir_sorted, depth_dir_sorted, pose_dir_sorted):
pose = np.loadtxt(pose)
pose[0:3, 1:3] *= -1
# We cannot accept files directly, as some of the poses are invalid
if np.isinf(pose).any():
continue
poses.append(pose)
intrinsics.append(K)
image_filenames.append(img)
depth_filenames.append(depth)
# filter image_filenames and poses based on train/eval split percentage
num_images = len(image_filenames)
num_train_images = math.ceil(num_images * self.config.train_split_fraction)
num_eval_images = num_images - num_train_images
i_all = np.arange(num_images)
i_train = np.linspace(
0, num_images - 1, num_train_images, dtype=int
) # equally spaced training images 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_images
if split == "train":
indices = i_train
elif split in ["val", "test"]:
indices = i_eval
else:
raise ValueError(f"Unknown dataparser split {split}")
poses = torch.from_numpy(np.stack(poses).astype(np.float32))
intrinsics = torch.from_numpy(np.stack(intrinsics).astype(np.float32))
if self.config.center_poses:
poses[:, :3, 3] -= poses[:, :3, 3].mean(dim=0)
if self.config.scale_poses:
poses[:, :3, 3] /= poses[:, :3, 3].abs().max()
# 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]
depth_filenames = [depth_filenames[i] for i in indices] if len(depth_filenames) > 0 else []
intrinsics = intrinsics[indices.tolist()]
poses = poses[indices.tolist()]
# in x,y,z order
# assumes that the scene is centered at the origin
aabb_scale = self.config.scene_scale
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=h,
width=w,
camera_to_worlds=poses[:, :3, :4],
camera_type=CameraType.PERSPECTIVE,
)
dataparser_outputs = DataparserOutputs(
image_filenames=image_filenames,
cameras=cameras,
scene_box=scene_box,
metadata={
"depth_filenames": depth_filenames if len(depth_filenames) > 0 else None,
"depth_unit_scale_factor": self.config.depth_unit_scale_factor,
},
)
return dataparser_outputs