diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 140d018ba..49b759a9d 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -42,6 +42,7 @@ from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConf from nerfstudio.data.dataparsers.instant_ngp_dataparser import ( InstantNGPDataParserConfig, ) +from nerfstudio.data.dataparsers.minimal_dataparser import MinimalDataParserConfig from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserConfig from nerfstudio.data.dataparsers.phototourism_dataparser import ( @@ -65,6 +66,7 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of tyro.extras.subcommand_type_from_defaults( { "nerfstudio-data": NerfstudioDataParserConfig(), + "minimal-parser": MinimalDataParserConfig(), "blender-data": BlenderDataParserConfig(), "friends-data": FriendsDataParserConfig(), "instant-ngp-data": InstantNGPDataParserConfig(), diff --git a/nerfstudio/data/datamanagers/semantic_datamanager.py b/nerfstudio/data/datamanagers/semantic_datamanager.py index 188c7f6c3..322b1d3bd 100644 --- a/nerfstudio/data/datamanagers/semantic_datamanager.py +++ b/nerfstudio/data/datamanagers/semantic_datamanager.py @@ -41,6 +41,7 @@ class SemanticDataManager(VanillaDataManager): # pylint: disable=abstract-metho """ def create_train_dataset(self) -> SemanticDataset: + self.train_dataparser_outputs = self.dataparser.get_dataparser_outputs(split="train") return SemanticDataset( dataparser_outputs=self.train_dataparser_outputs, scale_factor=self.config.camera_res_scale_factor, diff --git a/nerfstudio/data/dataparsers/minimal_dataparser.py b/nerfstudio/data/dataparsers/minimal_dataparser.py new file mode 100644 index 000000000..b6ffd8518 --- /dev/null +++ b/nerfstudio/data/dataparsers/minimal_dataparser.py @@ -0,0 +1,97 @@ +# 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 pre-prepared datasets for all cameras, with no additional processing needed +Optional fields - semantics, mask_filenames, cameras.distortion_params, cameras.times +""" + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Type + +import numpy as np +import torch + +from nerfstudio.cameras.cameras import Cameras +from nerfstudio.data.dataparsers.base_dataparser import ( + DataParser, + DataParserConfig, + DataparserOutputs, + Semantics, +) +from nerfstudio.data.scene_box import SceneBox + + +@dataclass +class MinimalDataParserConfig(DataParserConfig): + """Minimal dataset config""" + + _target: Type = field(default_factory=lambda: MinimalDataParser) + """target class to instantiate""" + data: Path = Path("/home/nikhil/nerfstudio-main/tests/data/lego_test/minimal_parser") + + +@dataclass +class MinimalDataParser(DataParser): + """Minimal DatasetParser""" + + config: MinimalDataParserConfig + + def _generate_dataparser_outputs(self, split="train"): + filepath = self.config.data / f"{split}.npz" + data = np.load(filepath, allow_pickle=True) + + image_filenames = data["image_filenames"].tolist() + mask_filenames = data["mask_filenames"].tolist() if "mask_filenames" in data.keys() else None + + metadata = None + if "semantics" in data.keys(): + semantics = data["semantics"].item() + metadata = { + "semantics": Semantics( + filenames=semantics["filenames"].tolist(), + classes=semantics["classes"].tolist(), + colors=semantics["colors"].tolist(), + mask_classes=semantics["mask_classes"].tolist(), + ) + } + + scene_box_aabb = torch.from_numpy(data["scene_box"]) + scene_box = SceneBox(aabb=scene_box_aabb) + + camera_np = data["cameras"].item() + distortion_params = None + if "distortion_params" in camera_np.keys(): + distortion_params = torch.from_numpy(camera_np["distortion_params"]) + cameras = Cameras( + fx=torch.from_numpy(camera_np["fx"]), + fy=torch.from_numpy(camera_np["fy"]), + cx=torch.from_numpy(camera_np["cx"]), + cy=torch.from_numpy(camera_np["cy"]), + distortion_params=distortion_params, + height=torch.from_numpy(camera_np["height"]), + width=torch.from_numpy(camera_np["width"]), + camera_to_worlds=torch.from_numpy(camera_np["camera_to_worlds"])[:, :3, :4], + camera_type=torch.from_numpy(camera_np["camera_type"]), + times=torch.from_numpy(camera_np["times"]) if "times" in camera_np.keys() else None, + ) + + dataparser_outputs = DataparserOutputs( + image_filenames=image_filenames, + cameras=cameras, + scene_box=scene_box, + mask_filenames=mask_filenames, + metadata=metadata, + ) + return dataparser_outputs diff --git a/tests/data/minimal_parser/train.npz b/tests/data/minimal_parser/train.npz new file mode 100644 index 000000000..9d8acbb15 Binary files /dev/null and b/tests/data/minimal_parser/train.npz differ diff --git a/tests/data/minimal_parser/val.npz b/tests/data/minimal_parser/val.npz new file mode 100644 index 000000000..9d8acbb15 Binary files /dev/null and b/tests/data/minimal_parser/val.npz differ diff --git a/tests/test_train.py b/tests/test_train.py index f7e860275..8b564369b 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -12,6 +12,7 @@ import pytest from nerfstudio.configs.experiment_config import ExperimentConfig from nerfstudio.configs.method_configs import method_configs from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig +from nerfstudio.data.dataparsers.minimal_dataparser import MinimalDataParserConfig from scripts.train import train_loop BLACKLIST = ["base", "semantic-nerfw", "instant-ngp", "instant-ngp-bounded", "nerfacto", "phototourism"] @@ -60,5 +61,14 @@ def test_train(): train_loop(local_rank=0, world_size=0, config=config) +def test_simple_io(): + """test to check minimal data IO works correctly""" + config = method_configs["vanilla-nerf"] + config.pipeline.datamanager.dataparser = MinimalDataParserConfig(data=Path("tests/data/minimal_parser")) + config = set_reduced_config(config) + train_loop(local_rank=0, world_size=0, config=config) + + if __name__ == "__main__": test_train() + test_simple_io()