mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-14 03:44:02 +08:00
Ethan WeberandMatt Tancik
9a39298e2f
* Move parser configs to files * fixed circular import errors * update to configs * merge fixes Co-authored-by: Matt Tancik <tancik@berkeley.edu>
107 lines
3.7 KiB
Python
107 lines
3.7 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.
|
|
|
|
"""Data parser for blender dataset"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Type
|
|
|
|
import imageio
|
|
import numpy as np
|
|
import torch
|
|
|
|
from nerfactory.cameras.cameras import Cameras, CameraType
|
|
from nerfactory.datamanagers.dataparsers.base import DataParser, DataParserConfig
|
|
from nerfactory.datamanagers.structs import DatasetInputs, SceneBounds
|
|
from nerfactory.utils.colors import get_color
|
|
from nerfactory.utils.io import get_absolute_path, load_from_json
|
|
|
|
|
|
@dataclass
|
|
class BlenderDataParserConfig(DataParserConfig):
|
|
"""Blender dataset parser config"""
|
|
|
|
_target: Type = field(default_factory=lambda: Blender)
|
|
"""target class to instantiate"""
|
|
data_directory: Path = Path("data/blender/lego")
|
|
"""directory specifying location of data"""
|
|
scale_factor: float = 1.0
|
|
"""How much to scale the camera origins by."""
|
|
alpha_color: str = "white"
|
|
"""alpha color of background"""
|
|
|
|
|
|
@dataclass
|
|
class Blender(DataParser):
|
|
"""Blender Dataset
|
|
Some of this code comes from https://github.com/yenchenlin/nerf-pytorch/blob/master/load_blender.py#L37.
|
|
"""
|
|
|
|
config: BlenderDataParserConfig
|
|
|
|
def __init__(self, config: BlenderDataParserConfig):
|
|
super().__init__(config=config)
|
|
self.data_directory: Path = config.data_directory
|
|
self.scale_factor: float = config.scale_factor
|
|
self.alpha_color = config.alpha_color
|
|
|
|
def _generate_dataset_inputs(self, split="train"):
|
|
if self.alpha_color is not None:
|
|
alpha_color_tensor = get_color(self.alpha_color)
|
|
else:
|
|
alpha_color_tensor = None
|
|
|
|
abs_dir = get_absolute_path(self.data_directory)
|
|
meta = load_from_json(abs_dir / f"transforms_{split}.json")
|
|
image_filenames = []
|
|
poses = []
|
|
for frame in meta["frames"]:
|
|
fname = abs_dir / Path(frame["file_path"].replace("./", "") + ".png")
|
|
image_filenames.append(fname)
|
|
poses.append(np.array(frame["transform_matrix"]))
|
|
poses = np.array(poses).astype(np.float32)
|
|
|
|
img_0 = imageio.imread(image_filenames[0])
|
|
image_height, image_width = img_0.shape[:2]
|
|
camera_angle_x = float(meta["camera_angle_x"])
|
|
focal_length = 0.5 * image_width / np.tan(0.5 * camera_angle_x)
|
|
|
|
cx = image_width / 2.0
|
|
cy = image_height / 2.0
|
|
camera_to_world = torch.from_numpy(poses[:, :3]) # camera to world transform
|
|
|
|
# in x,y,z order
|
|
camera_to_world[..., 3] *= self.scale_factor
|
|
scene_bounds = SceneBounds(aabb=torch.tensor([[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]], dtype=torch.float32))
|
|
|
|
cameras = Cameras(
|
|
camera_to_worlds=camera_to_world,
|
|
fx=focal_length,
|
|
fy=focal_length,
|
|
cx=cx,
|
|
cy=cy,
|
|
camera_type=CameraType.PERSPECTIVE,
|
|
)
|
|
|
|
dataset_inputs = DatasetInputs(
|
|
image_filenames=image_filenames,
|
|
cameras=cameras,
|
|
alpha_color=alpha_color_tensor,
|
|
scene_bounds=scene_bounds,
|
|
)
|
|
|
|
return dataset_inputs
|