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)