From 6e6cedd438348d03f390ba0d9b8afddaea4163c1 Mon Sep 17 00:00:00 2001 From: trinh31201 <141480383+trinh31201@users.noreply.github.com> Date: Wed, 6 Mar 2024 06:00:39 +0900 Subject: [PATCH 1/2] Handle alpha images properly for instant-ngp and tensorf models (#2979) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Đinh Trinh --- nerfstudio/models/instant_ngp.py | 2 +- nerfstudio/models/tensorf.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index 1b0304d5b..2503e7e22 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/nerfstudio/models/instant_ngp.py @@ -223,7 +223,7 @@ class NGPModel(Model): 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 index d19b0e2fb..6c149b4b6 100644 --- a/nerfstudio/models/tensorf.py +++ b/nerfstudio/models/tensorf.py @@ -308,7 +308,7 @@ class TensoRFModel(Model): 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"], From ae6c46cfeaebbe28f9cd48fa986755e27b5f0ae2 Mon Sep 17 00:00:00 2001 From: Nicolas Zunhammer Date: Fri, 8 Mar 2024 18:18:30 +0100 Subject: [PATCH 2/2] Fix: Masks having wrong dimension after undistortion (causes crash in splatfacto) (#2987) * Add third dim to masks after undistortion to fix error in multiplying color images with masks. * Revert auto format of imports to fix ruff error --------- Co-authored-by: Nicolas Zunhammer --- nerfstudio/data/datamanagers/full_images_datamanager.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py index 6e041b33d..47925acce 100644 --- a/nerfstudio/data/datamanagers/full_images_datamanager.py +++ b/nerfstudio/data/datamanagers/full_images_datamanager.py @@ -387,6 +387,8 @@ def _undistort_image( mask = cv2.undistort(mask, K, distortion_params, None, newK) # type: ignore mask = mask[y : y + h, x : x + w] mask = torch.from_numpy(mask).bool() + if len(mask.shape) == 2: + mask = mask[:, :, None] K = newK elif camera.camera_type.item() == CameraType.FISHEYE.value: @@ -406,6 +408,8 @@ def _undistort_image( mask = mask.astype(np.uint8) * 255 mask = cv2.fisheye.undistortImage(mask, K, distortion_params, None, newK) mask = torch.from_numpy(mask).bool() + if len(mask.shape) == 2: + mask = mask[:, :, None] K = newK elif camera.camera_type.item() == CameraType.FISHEYE624.value: fisheye624_params = torch.cat( @@ -498,6 +502,8 @@ def _undistort_image( ) / 255.0 ).bool()[..., None] + if len(mask.shape) == 2: + mask = mask[:, :, None] assert mask.shape == (undist_h, undist_w, 1) K = undist_K.numpy() else: