Refactor Data Pipeline (#896)

* add none check

* black formatting

* cleanup

* cleanup

* fix tests

* remove new line

* init

* toy example working

* semantic nerf working

* simplify pixel samplers

* cleanup + fix tests

* comments

* fix tests

* fix tests

docs/developer_guides/pipelines/fields.md

@@ -69,8 +69,6 @@ class FieldHeadNames(Enum):
     TRANSIENT_RGB = "transient_rgb"
     TRANSIENT_DENSITY = "transient_density"
     SEMANTICS = "semantics"
-    SEMANTICS_STUFF = "semantics_stuff"
-    SEMANTICS_THING = "semantics_thing"
 ```
 
 ```{button-link} https://github.com/nerfstudio-project/nerfstudio/blob/master/nerfstudio/field_components/field_heads.py

docs/reference/api/data.rst

@@ -13,7 +13,7 @@ Data Parsers
 Datasets
 ----------------
 
-.. automodule:: nerfstudio.data.utils.datasets
+.. automodule:: nerfstudio.data.datasets
    :members:
    :show-inheritance:
 

nerfstudio/configs/method_configs.py

@@ -29,7 +29,8 @@ from nerfstudio.configs.base_config import (
     TrainerConfig,
     ViewerConfig,
 )
-from nerfstudio.data.datamanagers import VanillaDataManagerConfig
+from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManagerConfig
+from nerfstudio.data.datamanagers.semantic_datamanager import SemanticDataManagerConfig
 from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
 from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig
 from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
@@ -127,7 +128,7 @@ method_configs["semantic-nerfw"] = Config(
     method_name="semantic-nerfw",
     trainer=TrainerConfig(steps_per_eval_batch=500, steps_per_save=2000, mixed_precision=True),
     pipeline=VanillaPipelineConfig(
-        datamanager=VanillaDataManagerConfig(
+        datamanager=SemanticDataManagerConfig(
             dataparser=FriendsDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=8192
         ),
         model=SemanticNerfWModelConfig(eval_num_rays_per_chunk=1 << 16),

nerfstudio/data/datamanagers/base_datamanager.py

@@ -41,13 +41,13 @@ from nerfstudio.data.dataparsers.instant_ngp_dataparser import (
 )
 from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
 from nerfstudio.data.dataparsers.record3d_dataparser import Record3DDataParserConfig
+from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.data.pixel_samplers import PixelSampler
 from nerfstudio.data.utils.dataloaders import (
     CacheDataloader,
     FixedIndicesEvalDataloader,
     RandIndicesEvalDataloader,
 )
-from nerfstudio.data.utils.datasets import InputDataset
 from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes
 from nerfstudio.model_components.ray_generators import RayGenerator
 from nerfstudio.utils.misc import IterableWrapper
@@ -296,13 +296,22 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         self.world_size = world_size
         self.local_rank = local_rank
         self.sampler = None
+        self.test_mode = test_mode
 
-        self.train_dataset = InputDataset(config.dataparser.setup().get_dataparser_outputs(split="train"))
-        self.eval_dataset = InputDataset(
-            config.dataparser.setup().get_dataparser_outputs(split="val" if not test_mode else "test")
-        )
+        self.train_dataset = self.create_train_dataset()
+        self.eval_dataset = self.create_eval_dataset()
         super().__init__()
 
+    def create_train_dataset(self) -> InputDataset:
+        """Sets up the data loaders for training"""
+        return InputDataset(self.config.dataparser.setup().get_dataparser_outputs(split="train"))
+
+    def create_eval_dataset(self) -> InputDataset:
+        """Sets up the data loaders for evaluation"""
+        return InputDataset(
+            self.config.dataparser.setup().get_dataparser_outputs(split="val" if not self.test_mode else "test")
+        )
+
     def setup_train(self):
         """Sets up the data loaders for training"""
         assert self.train_dataset is not None

nerfstudio/data/datamanagers/semantic_datamanager.py

@@ -0,0 +1,45 @@
+# 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.
+
+from dataclasses import dataclass, field
+from typing import Type
+
+from nerfstudio.data.datamanagers.base_datamanager import (
+    VanillaDataManager,
+    VanillaDataManagerConfig,
+)
+from nerfstudio.data.datasets.semantic_dataset import SemanticDataset
+
+
+@dataclass
+class SemanticDataManagerConfig(VanillaDataManagerConfig):
+    """A semantic datamanager - required to use with .setup()"""
+
+    _target: Type = field(default_factory=lambda: SemanticDataManager)
+
+
+class SemanticDataManager(VanillaDataManager):  # pylint: disable=abstract-method
+    """Data manager implementation for data that also requires processing semantic data.
+
+    Args:
+        config: the DataManagerConfig used to instantiate class
+    """
+
+    def create_train_dataset(self) -> SemanticDataset:
+        return SemanticDataset(self.config.dataparser.setup().get_dataparser_outputs(split="train"))
+
+    def create_eval_dataset(self) -> SemanticDataset:
+        return SemanticDataset(
+            self.config.dataparser.setup().get_dataparser_outputs(split="val" if not self.test_mode else "test")
+        )

nerfstudio/data/dataparsers/base_dataparser.py

@@ -34,18 +34,14 @@ from nerfstudio.data.scene_box import SceneBox
 class Semantics:
     """Dataclass for semantic labels."""
 
-    stuff_filenames: List[Path]
-    """filenames to load "stuff"/background data"""
-    stuff_classes: List[str]
-    """class labels for "stuff" data"""
-    stuff_colors: torch.Tensor
-    """color mapping for "stuff" classes"""
-    thing_filenames: List[Path]
-    """filenames to load "thing"/foreground data"""
-    thing_classes: List[str]
-    """class labels for "thing" data"""
-    thing_colors: torch.Tensor
-    """color mapping for "thing" classes"""
+    filenames: List[Path]
+    """filenames to load semantic data"""
+    classes: List[str]
+    """class labels for data"""
+    colors: torch.Tensor
+    """color mapping for classes"""
+    mask_classes: List[str] = field(default_factory=lambda: [])
+    """classes to mask out from training for all modalities"""
 
 
 @dataclass
@@ -61,16 +57,11 @@ class DataparserOutputs:
     """Color of dataset background."""
     scene_box: SceneBox = SceneBox()
     """Scene box of dataset. Used to bound the scene or provide the scene scale depending on model."""
-    semantics: Optional[Semantics] = None
-    """Semantics information."""
-    times: Optional[TensorType[1]] = None
-    """Time in range [0,1] for when each image was taken."""
-    additional_inputs: Dict[str, Any] = to_immutable_dict({})
-    """Dictionary of additional dataset information (e.g. semantics/point clouds/masks).
-    {input_name:
-    ... {"func": function to process additional dataparser outputs,
-    ... "kwargs": dictionary of data to pass into "func"}
-    }
+    mask_filenames: Optional[List[Path]] = None
+    """Filenames for any masks that are required"""
+    metadata: Dict[str, Any] = to_immutable_dict({})
+    """Dictionary of any metadata that be required for the given experiment.
+    Will be processed by the InputDataset to create any additional tensors that may be required.
     """
 
     def as_dict(self) -> dict:

nerfstudio/data/dataparsers/friends_dataparser.py

@@ -19,9 +19,7 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Type
 
-import numpy as np
 import torch
-from PIL import Image
 from rich.console import Console
 
 from nerfstudio.cameras.cameras import Cameras, CameraType
@@ -37,31 +35,6 @@ from nerfstudio.utils.io import load_from_json
 CONSOLE = Console()
 
 
-def get_semantics_and_masks(image_idx: int, semantics: Semantics):
-    """function to process additional semantics and mask information
-
-    Args:
-        image_idx: specific image index to work with
-        semantics: semantics data
-    """
-    # handle mask
-    person_index = semantics.thing_classes.index("person")
-    thing_image_filename = semantics.thing_filenames[image_idx]
-    pil_image = Image.open(thing_image_filename)
-    thing_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None]
-    mask = (thing_semantics != person_index).to(torch.float32)  # 1 where valid
-    # handle semantics
-    # stuff
-    stuff_image_filename = semantics.stuff_filenames[image_idx]
-    pil_image = Image.open(stuff_image_filename)
-    stuff_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None]
-    # thing
-    thing_image_filename = semantics.thing_filenames[image_idx]
-    pil_image = Image.open(thing_image_filename)
-    thing_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None]
-    return {"mask": mask, "semantics_stuff": stuff_semantics, "semantics_thing": thing_semantics}
-
-
 @dataclass
 class FriendsDataParserConfig(DataParserConfig):
     """Friends dataset parser config"""
@@ -142,9 +115,8 @@ class Friends(DataParser):
         camera_to_worlds[..., 3] *= scale  # cameras
 
         # --- semantics ---
-        semantics = None
         if self.config.include_semantics:
-            thing_filenames = [
+            filenames = [
                 Path(
                     str(image_filename)
                     .replace(f"/{images_folder}/", f"/{segmentations_folder}/thing/")
@@ -152,27 +124,10 @@ class Friends(DataParser):
                 )
                 for image_filename in image_filenames
             ]
-            stuff_filenames = [
-                Path(
-                    str(image_filename)
-                    .replace(f"/{images_folder}/", f"/{segmentations_folder}/stuff/")
-                    .replace(".jpg", ".png")
-                )
-                for image_filename in image_filenames
-            ]
             panoptic_classes = load_from_json(self.config.data / "panoptic_classes.json")
-            stuff_classes = panoptic_classes["stuff"]
-            stuff_colors = torch.tensor(panoptic_classes["stuff_colors"], dtype=torch.float32) / 255.0
-            thing_classes = panoptic_classes["thing"]
-            thing_colors = torch.tensor(panoptic_classes["thing_colors"], dtype=torch.float32) / 255.0
-            semantics = Semantics(
-                stuff_classes=stuff_classes,
-                stuff_colors=stuff_colors,
-                stuff_filenames=stuff_filenames,
-                thing_classes=thing_classes,
-                thing_colors=thing_colors,
-                thing_filenames=thing_filenames,
-            )
+            classes = panoptic_classes["thing"]
+            colors = torch.tensor(panoptic_classes["thing_colors"], dtype=torch.float32) / 255.0
+            semantics = Semantics(filenames=filenames, classes=classes, colors=colors, mask_classes=["person"])
 
         assert torch.all(cx[0] == cx), "Not all cameras have the same cx. Our Cameras class does not support this."
         assert torch.all(cy[0] == cy), "Not all cameras have the same cy. Our Cameras class does not support this."
@@ -191,7 +146,6 @@ class Friends(DataParser):
             image_filenames=image_filenames,
             cameras=cameras,
             scene_box=scene_box,
-            additional_inputs={"semantics": {"func": get_semantics_and_masks, "kwargs": {"semantics": semantics}}},
-            semantics=semantics,
+            metadata={"semantics": semantics} if self.config.include_semantics else {},
         )
         return dataparser_outputs

nerfstudio/data/datasets/base_dataset.py

@@ -27,6 +27,7 @@ from torch.utils.data import Dataset
 from torchtyping import TensorType
 
 from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs
+from nerfstudio.data.utils.data_utils import get_image_mask_tensor_from_path
 
 
 class InputDataset(Dataset):
@@ -39,6 +40,7 @@ class InputDataset(Dataset):
     def __init__(self, dataparser_outputs: DataparserOutputs):
         super().__init__()
         self.dataparser_outputs = dataparser_outputs
+        self.has_masks = self.dataparser_outputs.mask_filenames is not None
 
     def __len__(self):
         return len(self.dataparser_outputs.image_filenames)
@@ -71,7 +73,7 @@ class InputDataset(Dataset):
             image = image[:, :, :3]
         return image
 
-    def get_data(self, image_idx) -> Dict:
+    def get_data(self, image_idx: int) -> Dict:
         """Returns the ImageDataset data as a dictionary.
 
         Args:
@@ -80,13 +82,23 @@ class InputDataset(Dataset):
         image = self.get_image(image_idx)
         data = {"image_idx": image_idx}
         data["image"] = image
-        for _, data_func_dict in self.dataparser_outputs.additional_inputs.items():
-            assert "func" in data_func_dict, "Missing function to process data: specify `func` in `additional_inputs`"
-            func = data_func_dict["func"]
-            assert "kwargs" in data_func_dict, "No data to process: specify `kwargs` in `additional_inputs`"
-            data.update(func(image_idx, **data_func_dict["kwargs"]))
+        if self.has_masks:
+            mask_filepath = self.dataparser_outputs.mask_filenames[image_idx]
+            data["mask"] = get_image_mask_tensor_from_path(filepath=mask_filepath)
+        metadata = self.get_metadata(data)
+        data.update(metadata)
         return data
 
-    def __getitem__(self, image_idx):
+    # pylint: disable=no-self-use
+    def get_metadata(self, data: Dict) -> Dict:
+        """Method that can be used to process any additional metadata that may be part of the model inputs.
+
+        Args:
+            image_idx: The image index in the dataset.
+        """
+        del data
+        return {}
+
+    def __getitem__(self, image_idx: int) -> Dict:
         data = self.get_data(image_idx)
         return data

nerfstudio/data/datasets/semantic_dataset.py

@@ -0,0 +1,49 @@
+# 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.
+
+from typing import Dict
+
+import torch
+
+from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs, Semantics
+from nerfstudio.data.datasets.base_dataset import InputDataset
+from nerfstudio.data.utils.data_utils import get_semantics_and_mask_tensors_from_path
+
+
+class SemanticDataset(InputDataset):
+    """Dataset that returns images and semantics and masks.
+
+    Args:
+        dataparser_outputs: description of where and how to read input images.
+    """
+
+    def __init__(self, dataparser_outputs: DataparserOutputs):
+        super().__init__(dataparser_outputs)
+        assert "semantics" in dataparser_outputs.metadata.keys() and isinstance(
+            dataparser_outputs.metadata["semantics"], Semantics
+        )
+        self.semantics = dataparser_outputs.metadata["semantics"]
+        self.mask_indices = torch.tensor(
+            [self.semantics.classes.index(mask_class) for mask_class in self.semantics.mask_classes]
+        ).view(1, 1, -1)
+
+    def get_metadata(self, data: Dict) -> Dict:
+        # handle mask
+        filepath = self.semantics.filenames[data["image_idx"]]
+        semantic_label, mask = get_semantics_and_mask_tensors_from_path(
+            filepath=filepath, mask_indices=self.mask_indices
+        )
+        if "mask" in data.keys():
+            mask = mask & data["mask"]
+        return {"mask": mask, "semantics": semantic_label}

nerfstudio/data/pixel_samplers.py

@@ -48,30 +48,13 @@ def collate_image_dataset_batch(batch: Dict, num_rays_per_batch: int, keep_full_
         ).long()
 
     c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1))
-    image = batch["image"][c, y, x]
-    mask, semantics_stuff, semantics_thing = None, None, None
-    if "mask" in batch:
-        mask = batch["mask"][c, y, x]
-    if "semantics_stuff" in batch:
-        semantics_stuff = batch["semantics_stuff"][c, y, x]
-    if "semantics_thing" in batch:
-        semantics_thing = batch["semantics_thing"][c, y, x]
-    assert image.shape == (num_rays_per_batch, 3), image.shape
+    collated_batch = {key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and value is not None}
+
+    assert collated_batch["image"].shape == (num_rays_per_batch, 3), collated_batch["image"].shape
 
     # Needed to correct the random indices to their actual camera idx locations.
-    local_indices = indices.clone()
     indices[:, 0] = batch["image_idx"][c]
-    collated_batch = {
-        "local_indices": local_indices,  # local to the batch returned
-        "indices": indices,  # with the abs camera indices
-        "image": image,
-    }
-    if mask is not None:
-        collated_batch["mask"] = mask
-    if semantics_stuff is not None:
-        collated_batch["semantics_stuff"] = semantics_stuff
-    if semantics_thing is not None:
-        collated_batch["semantics_thing"] = semantics_thing
+    collated_batch["indices"] = indices  # with the abs camera indices
 
     if keep_full_image:
         collated_batch["full_image"] = batch["image"]

nerfstudio/data/utils/data_utils.py

@@ -0,0 +1,45 @@
+# 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.
+
+"""Utility functions to allow easy re-use of common operations across dataloaders"""
+from pathlib import Path
+from typing import List, Tuple, Union
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+def get_image_mask_tensor_from_path(filepath: Path) -> torch.Tensor:
+    """
+    Utility function to read a mask image from the given path and return a boolean tensor
+    """
+    pil_mask = Image.open(filepath)
+    mask_tensor = torch.from_numpy(np.array(pil_mask)).unsqueeze(-1).bool()
+    return mask_tensor
+
+
+def get_semantics_and_mask_tensors_from_path(
+    filepath: Path, mask_indices: Union[List, torch.Tensor]
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Utility function to read segmentation from the given filepath
+    If no mask is required - use mask_indices = []
+    """
+    if isinstance(mask_indices, List):
+        mask_indices = torch.tensor(mask_indices, dtype="int64").view(1, 1, -1)
+    pil_image = Image.open(filepath)
+    semantics = torch.from_numpy(np.array(pil_image, dtype="int64"))[..., None]
+    mask = torch.sum(semantics == mask_indices, dim=-1, keepdim=True) == 0
+    return semantics, mask

nerfstudio/data/utils/dataloaders.py

@@ -27,7 +27,7 @@ from torch.utils.data.dataloader import DataLoader
 
 from nerfstudio.cameras.cameras import Cameras
 from nerfstudio.cameras.rays import RayBundle
-from nerfstudio.data.utils.datasets import InputDataset
+from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.utils.misc import get_dict_to_torch
 
 CONSOLE = Console(width=120)

nerfstudio/field_components/field_heads.py

@@ -181,33 +181,3 @@ class SemanticFieldHead(FieldHead):
 
     def __init__(self, num_classes: int, in_dim: Optional[int] = None) -> None:
         super().__init__(in_dim=in_dim, out_dim=num_classes, field_head_name=FieldHeadNames.SEMANTICS, activation=None)
-
-
-class SemanticStuffFieldHead(FieldHead):
-    """Semantic stuff output
-
-    Args:
-        num_classes: Number of semantic classes
-        in_dim: input dimension. If not defined in constructor, it must be set later.
-        activation: output head activation
-    """
-
-    def __init__(self, num_classes: int, in_dim: Optional[int] = None) -> None:
-        super().__init__(
-            in_dim=in_dim, out_dim=num_classes, field_head_name=FieldHeadNames.SEMANTICS_STUFF, activation=None
-        )
-
-
-class SemanticThingFieldHead(FieldHead):
-    """Semantic thing output
-
-    Args:
-        num_classes: Number of semantic classes
-        in_dim: input dimension. If not defined in constructor, it must be set later.
-        activation: output head activation
-    """
-
-    def __init__(self, num_classes: int, in_dim: Optional[int] = None) -> None:
-        super().__init__(
-            in_dim=in_dim, out_dim=num_classes, field_head_name=FieldHeadNames.SEMANTICS_THING, activation=None
-        )

nerfstudio/fields/nerfacto_field.py

@@ -35,8 +35,7 @@ from nerfstudio.field_components.field_heads import (
     FieldHead,
     FieldHeadNames,
     RGBFieldHead,
-    SemanticStuffFieldHead,
-    SemanticThingFieldHead,
+    SemanticFieldHead,
     TransientDensityFieldHead,
     TransientRGBFieldHead,
     UncertaintyFieldHead,
@@ -95,8 +94,7 @@ class TCNNNerfactoField(Field):
         transient_embedding_dim: int = 16,
         use_transient_embedding: bool = False,
         use_semantics: bool = False,
-        num_semantics_thing_classes: int = 100,
-        num_semantics_stuff_classes: int = 100,
+        num_semantic_classes: int = 100,
         use_average_appearance_embedding: bool = False,
         spatial_distortion: Optional[SpatialDistortion] = None,
     ) -> None:
@@ -180,11 +178,8 @@ class TCNNNerfactoField(Field):
                     "n_hidden_layers": 1,
                 },
             )
-            self.field_head_semantics_stuff = SemanticStuffFieldHead(
-                in_dim=self.mlp_semantics.n_output_dims, num_classes=num_semantics_stuff_classes
-            )
-            self.field_head_semantics_thing = SemanticThingFieldHead(
-                in_dim=self.mlp_semantics.n_output_dims, num_classes=num_semantics_thing_classes
+            self.field_head_semantics = SemanticFieldHead(
+                in_dim=self.mlp_semantics.n_output_dims, num_classes=num_semantic_classes
             )
 
         self.mlp_head = tcnn.Network(
@@ -266,8 +261,7 @@ class TCNNNerfactoField(Field):
             )
             # print(semantics_input)
             x = self.mlp_semantics(semantics_input).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions)
-            outputs[FieldHeadNames.SEMANTICS_STUFF] = self.field_head_semantics_stuff(x)
-            outputs[FieldHeadNames.SEMANTICS_THING] = self.field_head_semantics_thing(x)
+            outputs[FieldHeadNames.SEMANTICS] = self.field_head_semantics(x)
 
         h = torch.cat(
             [

nerfstudio/models/semantic_nerfw.py

@@ -72,8 +72,9 @@ class SemanticNerfWModel(Model):
 
     config: SemanticNerfWModelConfig
 
-    def __init__(self, config: SemanticNerfWModelConfig, semantics: Semantics, **kwargs) -> None:
-        self.semantics = semantics
+    def __init__(self, config: SemanticNerfWModelConfig, metadata: Dict, **kwargs) -> None:
+        assert "semantics" in metadata.keys() and isinstance(metadata["semantics"], Semantics)
+        self.semantics = metadata["semantics"]
         super().__init__(config=config, **kwargs)
 
     def populate_modules(self):
@@ -93,8 +94,7 @@ class SemanticNerfWModel(Model):
             use_average_appearance_embedding=self.config.use_average_appearance_embedding,
             use_transient_embedding=self.config.use_transient_embedding,
             use_semantics=True,
-            num_semantics_stuff_classes=len(self.semantics.stuff_classes),
-            num_semantics_thing_classes=len(self.semantics.thing_classes),
+            num_semantic_classes=len(self.semantics.classes),
         )
 
         # Build the proposal network(s)
@@ -204,20 +204,13 @@ class SemanticNerfWModel(Model):
             outputs["density_transient"] = field_outputs[FieldHeadNames.TRANSIENT_DENSITY]
 
         # semantics
-        outputs["semantics_thing"] = self.renderer_semantics(
-            field_outputs[FieldHeadNames.SEMANTICS_THING], weights=weights_static.detach()
-        )
-        outputs["semantics_stuff"] = self.renderer_semantics(
-            field_outputs[FieldHeadNames.SEMANTICS_STUFF], weights=weights_static.detach()
+        outputs["semantics"] = self.renderer_semantics(
+            field_outputs[FieldHeadNames.SEMANTICS], weights=weights_static.detach()
         )
 
         # semantics colormaps
-        semantic_labels = torch.argmax(torch.nn.functional.softmax(outputs["semantics_stuff"], dim=-1), dim=-1)
-        outputs["semantics_colormap_stuff"] = self.semantics.stuff_colors[semantic_labels]
-        semantic_labels = torch.argmax(torch.nn.functional.softmax(outputs["semantics_thing"], dim=-1), dim=-1)
-        outputs["semantics_colormap_thing"] = self.semantics.thing_colors[semantic_labels]
-        # TODO: adding these together is incorrect and produces weird results
-        outputs["semantics_colormap"] = outputs["semantics_colormap_stuff"] + outputs["semantics_colormap_thing"]
+        semantic_labels = torch.argmax(torch.nn.functional.softmax(outputs["semantics"], dim=-1), dim=-1)
+        outputs["semantics_colormap"] = self.semantics.colors[semantic_labels]
 
         return outputs
 
@@ -247,12 +240,7 @@ class SemanticNerfWModel(Model):
             loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"])
 
         # semantic loss
-        loss_dict["semantics_stuff_loss"] = self.cross_entropy_loss(
-            outputs["semantics_stuff"], batch["semantics_stuff"][..., 0].long()
-        )
-        loss_dict["semantics_thing_loss"] = self.cross_entropy_loss(
-            outputs["semantics_thing"], batch["semantics_thing"][..., 0].long()
-        )
+        loss_dict["semantics_loss"] = self.cross_entropy_loss(outputs["semantics"], batch["semantics"][..., 0].long())
         return loss_dict
 
     def get_image_metrics_and_images(
@@ -295,14 +283,8 @@ class SemanticNerfWModel(Model):
             images_dict[key] = prop_depth_i
 
         # semantics
-        semantic_labels = torch.argmax(torch.nn.functional.softmax(outputs["semantics_stuff"], dim=-1), dim=-1)
-        images_dict["semantics_colormap_stuff"] = self.semantics.stuff_colors[semantic_labels]
-        semantic_labels = torch.argmax(torch.nn.functional.softmax(outputs["semantics_thing"], dim=-1), dim=-1)
-        images_dict["semantics_colormap_thing"] = self.semantics.thing_colors[semantic_labels]
-        # TODO: adding these together is incorrect and produces weird results
-        images_dict["semantics_colormap"] = (
-            images_dict["semantics_colormap_stuff"] + images_dict["semantics_colormap_thing"]
-        )
+        semantic_labels = torch.argmax(torch.nn.functional.softmax(outputs["semantics"], dim=-1), dim=-1)
+        images_dict["semantics_colormap"] = self.semantics.colors[semantic_labels]
 
         # valid mask
         images_dict["mask"] = batch["mask"].repeat(1, 1, 3)

nerfstudio/pipelines/base_pipeline.py

@@ -37,7 +37,7 @@ from torch.nn import Parameter
 from torch.nn.parallel import DistributedDataParallel as DDP
 
 from nerfstudio.configs import base_config as cfg
-from nerfstudio.data.datamanagers import (
+from nerfstudio.data.datamanagers.base_datamanager import (
     DataManager,
     VanillaDataManager,
     VanillaDataManagerConfig,
@@ -220,8 +220,7 @@ class VanillaPipeline(Pipeline):
         self._model = config.model.setup(
             scene_box=self.datamanager.train_dataset.dataparser_outputs.scene_box,
             num_train_data=len(self.datamanager.train_dataset),
-            semantics=self.datamanager.train_dataset.dataparser_outputs.semantics,
-            device=device,
+            metadata=self.datamanager.train_dataset.dataparser_outputs.metadata,
         )
         self.model.to(device)
 

nerfstudio/pipelines/dynamic_batch.py

@@ -21,7 +21,7 @@ from typing import Type
 
 import torch
 
-from nerfstudio.data.datamanagers import VanillaDataManager
+from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager
 from nerfstudio.pipelines.base_pipeline import VanillaPipeline, VanillaPipelineConfig
 
 

nerfstudio/viewer/server/viewer_utils.py

@@ -35,7 +35,7 @@ from rich.console import Console
 from nerfstudio.cameras.cameras import Cameras
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.configs import base_config as cfg
-from nerfstudio.data.utils.datasets import InputDataset
+from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.models.base_model import Model
 from nerfstudio.utils import colormaps, profiler, writer
 from nerfstudio.utils.decorators import check_main_thread, decorate_all
@@ -446,12 +446,12 @@ class ViewerState:
             self.camera_moving = True
         return camera_object
 
-    def _apply_colormap(self, outputs: Dict[str, Any], stuff_colors: torch.Tensor = None, eps=1e-6):
+    def _apply_colormap(self, outputs: Dict[str, Any], colors: torch.Tensor = None, eps=1e-6):
         """Determines which colormap to use based on set colormap type
 
         Args:
             outputs: the output tensors for which to apply colormaps on
-            stuff_colors: is only set if colormap is for semantics. Defaults to None.
+            colors: is only set if colormap is for semantics. Defaults to None.
             eps: epsilon to handle floating point comparisons
         """
         if self.output_list:
@@ -486,8 +486,8 @@ class ViewerState:
         ):
             logits = outputs[reformatted_output]
             labels = torch.argmax(torch.nn.functional.softmax(logits, dim=-1), dim=-1)  # type: ignore
-            assert stuff_colors is not None
-            return stuff_colors[labels]
+            assert colors is not None
+            return colors[labels]
 
         # rendering boolean outputs
         if self.prev_colormap_type == ColormapTypes.BOOLEAN or (
@@ -532,12 +532,12 @@ class ViewerState:
         for video_track in self.video_tracks:
             video_track.put_frame(image)
 
-    def _send_output_to_viewer(self, outputs: Dict[str, Any], stuff_colors: torch.Tensor = None, eps=1e-6):
+    def _send_output_to_viewer(self, outputs: Dict[str, Any], colors: torch.Tensor = None, eps=1e-6):
         """Chooses the correct output and sends it to the viewer
 
         Args:
             outputs: the dictionary of outputs to choose from, from the graph
-            stuff_colors: is only set if colormap is for semantics. Defaults to None.
+            colors: is only set if colormap is for semantics. Defaults to None.
             eps: epsilon to handle floating point comparisons
         """
         if self.output_list is None:
@@ -564,7 +564,7 @@ class ViewerState:
             self.output_type_changed = False
             self.vis["renderingState/colormap_choice"].write(self.prev_colormap_type)
             self.vis["renderingState/colormap_options"].write(colormap_options)
-        selected_output = (self._apply_colormap(outputs, stuff_colors) * 255).type(torch.uint8)
+        selected_output = (self._apply_colormap(outputs, colors) * 255).type(torch.uint8)
         image = selected_output.cpu().numpy()
         self.set_image(image)
 
@@ -764,8 +764,8 @@ class ViewerState:
         graph.train()
         outputs = render_thread.vis_outputs
         if outputs is not None:
-            stuff_colors = graph.stuff_colors if hasattr(graph, "stuff_colors") else None
-            self._send_output_to_viewer(outputs, stuff_colors=stuff_colors)
+            colors = graph.colors if hasattr(graph, "colors") else None
+            self._send_output_to_viewer(outputs, colors=colors)
             self._update_viewer_stats(
                 vis_t.duration, num_rays=len(camera_ray_bundle), image_height=image_height, image_width=image_width
             )

scripts/viewer/view_dataset.py

@@ -15,7 +15,7 @@ import tyro
 
 from nerfstudio.configs import method_configs as cfg
 from nerfstudio.data.datamanagers import AnnotatedDataParserUnion
-from nerfstudio.data.utils.datasets import InputDataset
+from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.viewer.server import viewer_utils
 
 DEFAULT_TIMEOUT = timedelta(minutes=30)