Minimal disk dataparser (#1128)

* wip

* merge upstream

* clkeanup

* cleanup

* cleanup

* cleanup

* more cleanup

* cleanup

* cleanup docs

* lint

* style fix

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(),

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,

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

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()