From 1a24f3e58c544bc0211563e770d425426284256c Mon Sep 17 00:00:00 2001 From: Nikhil Mohan <47728081+nikmo33@users.noreply.github.com> Date: Sun, 6 Nov 2022 15:40:43 +0000 Subject: [PATCH] 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 | 2 - docs/reference/api/data.rst | 2 +- nerfstudio/configs/method_configs.py | 5 +- .../base_datamanager.py} | 19 +++++-- .../data/datamanagers/semantic_datamanager.py | 45 +++++++++++++++ .../data/dataparsers/base_dataparser.py | 35 +++++------- .../data/dataparsers/friends_dataparser.py | 56 ++----------------- .../datasets.py => datasets/base_dataset.py} | 26 ++++++--- nerfstudio/data/datasets/semantic_dataset.py | 49 ++++++++++++++++ nerfstudio/data/pixel_samplers.py | 25 ++------- nerfstudio/data/utils/data_utils.py | 45 +++++++++++++++ nerfstudio/data/utils/dataloaders.py | 2 +- nerfstudio/field_components/field_heads.py | 30 ---------- nerfstudio/fields/nerfacto_field.py | 16 ++---- nerfstudio/models/semantic_nerfw.py | 40 ++++--------- nerfstudio/pipelines/base_pipeline.py | 5 +- nerfstudio/pipelines/dynamic_batch.py | 2 +- nerfstudio/viewer/server/viewer_utils.py | 20 +++---- scripts/viewer/view_dataset.py | 2 +- 19 files changed, 229 insertions(+), 197 deletions(-) rename nerfstudio/data/{datamanagers.py => datamanagers/base_datamanager.py} (96%) create mode 100644 nerfstudio/data/datamanagers/semantic_datamanager.py rename nerfstudio/data/{utils/datasets.py => datasets/base_dataset.py} (77%) create mode 100644 nerfstudio/data/datasets/semantic_dataset.py create mode 100644 nerfstudio/data/utils/data_utils.py diff --git a/docs/developer_guides/pipelines/fields.md b/docs/developer_guides/pipelines/fields.md index 8c6c6005f..f4c9c6d2e 100644 --- a/docs/developer_guides/pipelines/fields.md +++ b/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 diff --git a/docs/reference/api/data.rst b/docs/reference/api/data.rst index a0ac5a2b2..de7b57abe 100644 --- a/docs/reference/api/data.rst +++ b/docs/reference/api/data.rst @@ -13,7 +13,7 @@ Data Parsers Datasets ---------------- -.. automodule:: nerfstudio.data.utils.datasets +.. automodule:: nerfstudio.data.datasets :members: :show-inheritance: diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index f99110569..377c782aa 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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), diff --git a/nerfstudio/data/datamanagers.py b/nerfstudio/data/datamanagers/base_datamanager.py similarity index 96% rename from nerfstudio/data/datamanagers.py rename to nerfstudio/data/datamanagers/base_datamanager.py index 2e908ad35..82cf4a0d5 100644 --- a/nerfstudio/data/datamanagers.py +++ b/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 diff --git a/nerfstudio/data/datamanagers/semantic_datamanager.py b/nerfstudio/data/datamanagers/semantic_datamanager.py new file mode 100644 index 000000000..a66e4f365 --- /dev/null +++ b/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") + ) diff --git a/nerfstudio/data/dataparsers/base_dataparser.py b/nerfstudio/data/dataparsers/base_dataparser.py index 5f19f126c..3b99f274e 100644 --- a/nerfstudio/data/dataparsers/base_dataparser.py +++ b/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: diff --git a/nerfstudio/data/dataparsers/friends_dataparser.py b/nerfstudio/data/dataparsers/friends_dataparser.py index 7a1afa69e..cd60e9951 100644 --- a/nerfstudio/data/dataparsers/friends_dataparser.py +++ b/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 diff --git a/nerfstudio/data/utils/datasets.py b/nerfstudio/data/datasets/base_dataset.py similarity index 77% rename from nerfstudio/data/utils/datasets.py rename to nerfstudio/data/datasets/base_dataset.py index b192a39c7..0c6c9fe15 100644 --- a/nerfstudio/data/utils/datasets.py +++ b/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 diff --git a/nerfstudio/data/datasets/semantic_dataset.py b/nerfstudio/data/datasets/semantic_dataset.py new file mode 100644 index 000000000..2fc551488 --- /dev/null +++ b/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} diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 8ba19ba63..6da0dff33 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/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"] diff --git a/nerfstudio/data/utils/data_utils.py b/nerfstudio/data/utils/data_utils.py new file mode 100644 index 000000000..953e7951b --- /dev/null +++ b/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 diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py index 825cb65a1..8f8873d65 100644 --- a/nerfstudio/data/utils/dataloaders.py +++ b/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) diff --git a/nerfstudio/field_components/field_heads.py b/nerfstudio/field_components/field_heads.py index c11a827d1..702cafd81 100644 --- a/nerfstudio/field_components/field_heads.py +++ b/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 - ) diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index 86e479e48..b70b6b4c3 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/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( [ diff --git a/nerfstudio/models/semantic_nerfw.py b/nerfstudio/models/semantic_nerfw.py index c30804439..3f9dcad86 100644 --- a/nerfstudio/models/semantic_nerfw.py +++ b/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) diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 7c64e268f..978ae41b0 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/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) diff --git a/nerfstudio/pipelines/dynamic_batch.py b/nerfstudio/pipelines/dynamic_batch.py index b2e18ef62..d7d9e7f97 100644 --- a/nerfstudio/pipelines/dynamic_batch.py +++ b/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 diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 4824262b7..e0ea6e36e 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/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 ) diff --git a/scripts/viewer/view_dataset.py b/scripts/viewer/view_dataset.py index 29317515f..68b54e42d 100644 --- a/scripts/viewer/view_dataset.py +++ b/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)