mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-12 19:04:03 +08:00
119 lines
4.0 KiB
Python
119 lines
4.0 KiB
Python
# Copyright 2022 The Plenoptix 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.
|
|
|
|
"""
|
|
Code for loading mipnerf 360 dataset formats.
|
|
"""
|
|
|
|
import os
|
|
|
|
import imageio
|
|
import numpy as np
|
|
import torch
|
|
|
|
from nerfactory.data.structs import DatasetInputs
|
|
|
|
|
|
def load_mipnerf_360_data(
|
|
basedir: str,
|
|
downscale_factor: int = 1,
|
|
split: str = "train",
|
|
val_skip: int = 8,
|
|
auto_scale: bool = True,
|
|
) -> DatasetInputs:
|
|
"""Processes mipnerf 360 data.
|
|
Example data can be downloaded from https://jonbarron.info/mipnerf360/.
|
|
|
|
TODO: code currently assumes images were previously downscaled.
|
|
|
|
Args:
|
|
data_directory (str): Location of data
|
|
downscale_factor (int, optional): How much to downscale images. Must be a factor of 2. Defaults to 1.
|
|
split (str, optional): Which dataset split to generate.
|
|
val_skip (int, optional): 1/val_skip images to use for validation. Defaults to 8.
|
|
auto_scale (boolean, optional): Scale based on pose bounds. Defaults to True.
|
|
aabb_scale (float, optional): Scene scale, Defaults to 1.0.
|
|
|
|
Returns:
|
|
DatasetInputs
|
|
"""
|
|
|
|
image_dir = os.path.join(basedir, "images")
|
|
if downscale_factor > 1:
|
|
image_dir += f"_{downscale_factor}"
|
|
|
|
if not os.path.exists(image_dir):
|
|
raise ValueError(f"Image directory {image_dir} doesn't exist")
|
|
|
|
valid_formats = [".jpg", ".png"]
|
|
image_filenames = []
|
|
for f in os.listdir(image_dir):
|
|
ext = os.path.splitext(f)[1]
|
|
if ext.lower() not in valid_formats:
|
|
continue
|
|
image_filenames.append(os.path.join(image_dir, f))
|
|
image_filenames = sorted(image_filenames)
|
|
num_images = len(image_filenames)
|
|
|
|
poses_data = np.load(os.path.join(basedir, "poses_bounds.npy"))
|
|
poses = poses_data[:, :-2].reshape([-1, 3, 5]).astype(np.float32)
|
|
bounds = poses_data[:, -2:].transpose([1, 0])
|
|
|
|
if num_images != poses.shape[0]:
|
|
raise RuntimeError(f"Different number of images ({num_images}), and poses ({poses.shape[0]})")
|
|
|
|
idx_test = np.arange(num_images)[::val_skip]
|
|
idx_train = np.array([i for i in np.arange(num_images) if i not in idx_test])
|
|
idx = idx_train if split == "train" else idx_test
|
|
|
|
image_filenames = np.array(image_filenames)[idx]
|
|
poses = poses[idx]
|
|
|
|
img_0 = imageio.imread(image_filenames[0])
|
|
image_height, image_width = img_0.shape[:2]
|
|
|
|
poses[:, :2, 4] = np.array([image_height, image_width])
|
|
poses[:, 2, 4] = poses[:, 2, 4] * 1.0 / downscale_factor
|
|
|
|
# Reorder pose to match our convention
|
|
poses = np.concatenate([poses[:, :, 1:2], -poses[:, :, 0:1], poses[:, :, 2:]], axis=-1)
|
|
|
|
# Scale factor used in mipnerf
|
|
if auto_scale:
|
|
scale_factor = 1 / (np.min(bounds) * 0.75)
|
|
poses[:, :3, 3] *= scale_factor
|
|
bounds *= scale_factor
|
|
|
|
# Center poses
|
|
poses[:, :3, 3] = poses[:, :3, 3] - np.mean(poses[:, :3, :], axis=0)[:, 3]
|
|
|
|
focal_length = poses[0, -1, -1]
|
|
|
|
cx = image_width / 2.0
|
|
cy = image_height / 2.0
|
|
camera_to_world = torch.from_numpy(poses[:, :3, :4]) # camera to world transform
|
|
num_cameras = len(image_filenames)
|
|
num_intrinsics_params = 3
|
|
intrinsics = torch.ones((num_cameras, num_intrinsics_params), dtype=torch.float32)
|
|
intrinsics *= torch.tensor([cx, cy, focal_length])
|
|
|
|
dataset_inputs = DatasetInputs(
|
|
image_filenames=image_filenames,
|
|
downscale_factor=1,
|
|
intrinsics=intrinsics,
|
|
camera_to_world=camera_to_world,
|
|
)
|
|
|
|
return dataset_inputs
|