Handle alpha images properly for instant-ngp and tensorf models (nerf…

…studio-project#2979) Co-authored-by: Đinh Trinh <[email protected]>
Michael-Spleenlab · Apr 26, 2024 · a155f40 · a155f40
1 parent 040346c
commit a155f40
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py
@@ -223,7 +223,7 @@ def get_metrics_dict(self, outputs, batch):
         return metrics_dict
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None):
-        image = batch["image"][..., :3].to(self.device)
+        image = batch["image"].to(self.device)
         pred_rgb, image = self.renderer_rgb.blend_background_for_loss_computation(
             pred_image=outputs["rgb"],
             pred_accumulation=outputs["accumulation"],

diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py
@@ -308,7 +308,7 @@ def get_outputs(self, ray_bundle: RayBundle):
     def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:
         # Scaling metrics by coefficients to create the losses.
         device = outputs["rgb"].device
-        image = batch["image"][..., :3].to(device)
+        image = batch["image"].to(device)
         pred_image, image = self.renderer_rgb.blend_background_for_loss_computation(
             pred_image=outputs["rgb"],
             pred_accumulation=outputs["accumulation"],