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 <machenmusik@users.noreply.github.com>

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

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
 
 

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)
 

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)

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)