mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-13 03:14:00 +08:00
Ethan Weber
457788d689
* renaming dataloader to data. adding some dataloader classes * restructuring configs * fixed instant ngp tcnn direction rescaling * working through test cases * black and lint and tests
98 lines
3.6 KiB
Python
98 lines
3.6 KiB
Python
"""
|
|
Code for loading the dataset inputs and for caching the data for fast loading.
|
|
"""
|
|
|
|
import logging
|
|
from typing import Optional, Union
|
|
|
|
from omegaconf import ListConfig
|
|
|
|
from pyrad.data.format.blender import load_blender_data
|
|
from pyrad.data.format.friends import load_friends_data
|
|
from pyrad.data.format.instant_ngp import load_instant_ngp_data
|
|
from pyrad.data.structs import DatasetInputs
|
|
from pyrad.utils.colors import get_color
|
|
from pyrad.utils.io import get_absolute_path, load_from_pkl, make_dir, write_to_pkl
|
|
from pyrad.utils.misc import get_hash_str_from_dict
|
|
|
|
|
|
def get_cache_filename_from_kwargs(kwargs: dict, split: str):
|
|
"""Creates a cache filename from the dataset inputs arguments."""
|
|
dataset_config_hash = get_hash_str_from_dict(kwargs)
|
|
dataset_config_hash_filename = make_dir(
|
|
get_absolute_path(f"cache/dataset_inputs/{dataset_config_hash}-{split}.pkl")
|
|
)
|
|
return dataset_config_hash_filename
|
|
|
|
|
|
def save_dataset_inputs_kwargs_to_cache(kwargs: dict, split: str):
|
|
"""Saves the dataset inputs to cache."""
|
|
dataset_inputs = get_dataset_inputs(**kwargs, split=split)
|
|
dataset_config_hash_filename = get_cache_filename_from_kwargs(kwargs, split)
|
|
write_to_pkl(dataset_config_hash_filename, dataset_inputs)
|
|
|
|
|
|
def get_dataset_inputs_from_cache(kwargs: dict, split: str):
|
|
"""Loads the dataset inputs from cache."""
|
|
dataset_config_hash_filename = get_cache_filename_from_kwargs(kwargs, split)
|
|
return load_from_pkl(dataset_config_hash_filename)
|
|
|
|
|
|
def get_dataset_inputs(
|
|
data_directory: str,
|
|
dataset_format: str,
|
|
split: str,
|
|
downscale_factor: int = 1,
|
|
alpha_color: Optional[Union[str, list, ListConfig]] = None,
|
|
) -> DatasetInputs:
|
|
"""Returns the dataset inputs, which will be used with an ImageDataset and RayGenerator.
|
|
|
|
Args:
|
|
data_directory (str): Location of data
|
|
dataset_format (str): Name of dataset type
|
|
downscale_factor (int, optional): How much to downscale images. Defaults to 1.0.
|
|
alpha_color (str, list, optional): Sets transparent regions to specified color, otherwise black.
|
|
load_from_cache (bool)
|
|
Returns:
|
|
DatasetInputs: The inputs needed for generating rays.
|
|
"""
|
|
|
|
if alpha_color is not None:
|
|
alpha_color = get_color(alpha_color)
|
|
|
|
if dataset_format == "blender":
|
|
dataset_inputs = load_blender_data(
|
|
get_absolute_path(data_directory),
|
|
downscale_factor=downscale_factor,
|
|
split=split,
|
|
alpha_color=alpha_color,
|
|
)
|
|
elif dataset_format == "friends":
|
|
# TODO(ethan): change this hack, and do proper splits for the friends dataset
|
|
# currently we assume that there is only the training set of images
|
|
dataset_inputs = load_friends_data(
|
|
get_absolute_path(data_directory),
|
|
downscale_factor=downscale_factor,
|
|
split="train",
|
|
include_point_cloud=False,
|
|
)
|
|
elif dataset_format == "instant_ngp":
|
|
dataset_inputs = load_instant_ngp_data(
|
|
get_absolute_path(data_directory),
|
|
downscale_factor=downscale_factor,
|
|
split="train",
|
|
)
|
|
else:
|
|
raise NotImplementedError(f"{dataset_format} is not a valid dataset type")
|
|
|
|
return dataset_inputs
|
|
|
|
|
|
def get_dataset_inputs_from_dataset_config(*, split: str, use_cache: bool = False, **kwargs):
|
|
"""This may optionally use the cache to return the dataset inputs."""
|
|
if use_cache:
|
|
logging.info("Loading from cache! Be careful with using this when making changes!")
|
|
print(kwargs)
|
|
return get_dataset_inputs_from_cache(kwargs, split)
|
|
return get_dataset_inputs(**kwargs, split=split)
|