From a91b92f89d401115b42e6ad295a8b850e02ee86f Mon Sep 17 00:00:00 2001 From: Nikhil Mohan <47728081+nikmo33@users.noreply.github.com> Date: Thu, 23 Feb 2023 03:21:04 +0000 Subject: [PATCH] semantic_nerfw_configurability (#1467) * fix semseg vis * semantic configs * change defaults * Fix type * format * fix imports * suggested fix for issue 1465 * fix docstring placement; avoid log spam and runtime errors `nerfstudio\models\semantic_nerfw.py:81: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor). self.colormap = torch.tensor(self.semantics.colors).to(self.device)` * hide semantics output choice, since it kills the viewer * revert some changes * caps * move both image and mask to GPU to avoid runtime error as suggested in https://github.com/nerfstudio-project/nerfstudio/issues/1465#issuecomment-1441117562 --------- Co-authored-by: machenmusik --- .../data/dataparsers/minimal_dataparser.py | 2 +- nerfstudio/data/utils/dataloaders.py | 4 ++-- nerfstudio/fields/nerfacto_field.py | 13 ++++++------- nerfstudio/models/semantic_nerfw.py | 17 +++++++++++++---- nerfstudio/viewer/server/viewer_utils.py | 3 +++ 5 files changed, 25 insertions(+), 14 deletions(-) diff --git a/nerfstudio/data/dataparsers/minimal_dataparser.py b/nerfstudio/data/dataparsers/minimal_dataparser.py index b6ffd8518..2fa2218a4 100644 --- a/nerfstudio/data/dataparsers/minimal_dataparser.py +++ b/nerfstudio/data/dataparsers/minimal_dataparser.py @@ -62,7 +62,7 @@ class MinimalDataParser(DataParser): "semantics": Semantics( filenames=semantics["filenames"].tolist(), classes=semantics["classes"].tolist(), - colors=semantics["colors"].tolist(), + colors=torch.from_numpy(semantics["colors"]), mask_classes=semantics["mask_classes"].tolist(), ) } diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py index 23d712d7e..c4e155102 100644 --- a/nerfstudio/data/utils/dataloaders.py +++ b/nerfstudio/data/utils/dataloaders.py @@ -116,7 +116,7 @@ class CacheDataloader(DataLoader): """Returns a collated batch.""" batch_list = self._get_batch_list() collated_batch = self.collate_fn(batch_list) - collated_batch = get_dict_to_torch(collated_batch, device=self.device, exclude=["image"]) + collated_batch = get_dict_to_torch(collated_batch, device=self.device) return collated_batch def __iter__(self): @@ -183,7 +183,7 @@ class EvalDataloader(DataLoader): """ ray_bundle = self.cameras.generate_rays(camera_indices=image_idx, keep_shape=True) batch = self.input_dataset[image_idx] - batch = get_dict_to_torch(batch, device=self.device, exclude=["image"]) + batch = get_dict_to_torch(batch, device=self.device) return ray_bundle, batch diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index 1c3865fee..746a538a1 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/nerfstudio/fields/nerfacto_field.py @@ -109,6 +109,7 @@ class TCNNNerfactoField(Field): use_transient_embedding: bool = False, use_semantics: bool = False, num_semantic_classes: int = 100, + pass_semantic_gradients: bool = False, use_pred_normals: bool = False, use_average_appearance_embedding: bool = False, spatial_distortion: Optional[SpatialDistortion] = None, @@ -126,6 +127,7 @@ class TCNNNerfactoField(Field): self.use_transient_embedding = use_transient_embedding self.use_semantics = use_semantics self.use_pred_normals = use_pred_normals + self.pass_semantic_gradients = pass_semantic_gradients base_res = 16 features_per_level = 2 @@ -291,13 +293,10 @@ class TCNNNerfactoField(Field): # semantics if self.use_semantics: - density_embedding_copy = density_embedding.clone().detach() - semantics_input = torch.cat( - [ - density_embedding_copy.view(-1, self.geo_feat_dim), - ], - dim=-1, - ) + semantics_input = density_embedding.view(-1, self.geo_feat_dim) + if not self.pass_semantic_gradients: + semantics_input = semantics_input.detach() + x = self.mlp_semantics(semantics_input).view(*outputs_shape, -1).to(directions) outputs[FieldHeadNames.SEMANTICS] = self.field_head_semantics(x) diff --git a/nerfstudio/models/semantic_nerfw.py b/nerfstudio/models/semantic_nerfw.py index 22957d268..4324e1eeb 100644 --- a/nerfstudio/models/semantic_nerfw.py +++ b/nerfstudio/models/semantic_nerfw.py @@ -61,6 +61,8 @@ class SemanticNerfWModelConfig(NerfactoModelConfig): _target: Type = field(default_factory=lambda: SemanticNerfWModel) use_transient_embedding: bool = False """Whether to use transient embedding.""" + semantic_loss_weight: float = 1.0 + pass_semantic_gradients: bool = False class SemanticNerfWModel(Model): @@ -76,6 +78,7 @@ class SemanticNerfWModel(Model): assert "semantics" in metadata.keys() and isinstance(metadata["semantics"], Semantics) self.semantics = metadata["semantics"] super().__init__(config=config, **kwargs) + self.colormap = self.semantics.colors.clone().detach().to(self.device) def populate_modules(self): """Set the fields and modules.""" @@ -98,6 +101,7 @@ class SemanticNerfWModel(Model): use_transient_embedding=self.config.use_transient_embedding, use_semantics=True, num_semantic_classes=len(self.semantics.classes), + pass_semantic_gradients=self.config.pass_semantic_gradients, ) # Build the proposal network(s) @@ -207,13 +211,16 @@ class SemanticNerfWModel(Model): outputs["density_transient"] = field_outputs[FieldHeadNames.TRANSIENT_DENSITY] # semantics + semantic_weights = weights_static + if not self.config.pass_semantic_gradients: + semantic_weights = semantic_weights.detach() outputs["semantics"] = self.renderer_semantics( - field_outputs[FieldHeadNames.SEMANTICS], weights=weights_static.detach() + field_outputs[FieldHeadNames.SEMANTICS], weights=semantic_weights ) # semantics colormaps semantic_labels = torch.argmax(torch.nn.functional.softmax(outputs["semantics"], dim=-1), dim=-1) - outputs["semantics_colormap"] = self.semantics.colors[semantic_labels] + outputs["semantics_colormap"] = self.colormap.to(self.device)[semantic_labels] return outputs @@ -243,7 +250,9 @@ class SemanticNerfWModel(Model): loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"]) # semantic loss - loss_dict["semantics_loss"] = self.cross_entropy_loss(outputs["semantics"], batch["semantics"][..., 0].long()) + loss_dict["semantics_loss"] = self.config.semantic_loss_weight * self.cross_entropy_loss( + outputs["semantics"], batch["semantics"][..., 0].long() + ) return loss_dict def get_image_metrics_and_images( @@ -287,7 +296,7 @@ class SemanticNerfWModel(Model): # semantics semantic_labels = torch.argmax(torch.nn.functional.softmax(outputs["semantics"], dim=-1), dim=-1) - images_dict["semantics_colormap"] = self.semantics.colors[semantic_labels] + images_dict["semantics_colormap"] = self.colormap.to(self.device)[semantic_labels] # valid mask images_dict["mask"] = batch["mask"].repeat(1, 1, 3) diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 2c1541467..30d1fdb54 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -737,6 +737,9 @@ class ViewerState: if OutputTypes.RGB_FINE in self.output_list: viewer_output_list.remove(OutputTypes.RGB_FINE) viewer_output_list.insert(0, OutputTypes.RGB) + # remove semantics, which crashes viewer; semantics_colormap is OK + if "semantics" in self.output_list: + viewer_output_list.remove("semantics") self.vis["renderingState/output_options"].write(viewer_output_list) reformatted_output = self._process_invalid_output(self.prev_output_type)