From ae6c46cfeaebbe28f9cd48fa986755e27b5f0ae2 Mon Sep 17 00:00:00 2001 From: Nicolas Zunhammer Date: Fri, 8 Mar 2024 18:18:30 +0100 Subject: [PATCH] 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: