From fee8ab7d131610ac8192a973a230444a77550aa3 Mon Sep 17 00:00:00 2001 From: Jinhwan Choi Date: Wed, 21 Feb 2024 17:52:16 +0900 Subject: [PATCH] fix crash when both mask and downscale_factor are set (#2920) * fix crash when both mask and downscale_factor are set --- nerfstudio/models/splatfacto.py | 27 +++++++++++++++------------ 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/nerfstudio/models/splatfacto.py b/nerfstudio/models/splatfacto.py index e1f64ebcb..27319fb91 100644 --- a/nerfstudio/models/splatfacto.py +++ b/nerfstudio/models/splatfacto.py @@ -645,6 +645,17 @@ class SplatfactoModel(Model): else: return 1 + def _downscale_if_required(self, image): + d = self._get_downscale_factor() + if d > 1: + newsize = [image.shape[0] // d, image.shape[1] // d] + + # torchvision can be slow to import, so we do it lazily. + import torchvision.transforms.functional as TF + + return TF.resize(image.permute(2, 0, 1), newsize, antialias=None).permute(1, 2, 0) + return image + def get_outputs(self, camera: Cameras) -> Dict[str, Union[torch.Tensor, List]]: """Takes in a Ray Bundle and returns a dictionary of outputs. @@ -818,16 +829,7 @@ class SplatfactoModel(Model): """ if image.dtype == torch.uint8: image = image.float() / 255.0 - d = self._get_downscale_factor() - if d > 1: - newsize = [image.shape[0] // d, image.shape[1] // d] - - # torchvision can be slow to import, so we do it lazily. - import torchvision.transforms.functional as TF - - gt_img = TF.resize(image.permute(2, 0, 1), newsize, antialias=None).permute(1, 2, 0) - else: - gt_img = image + gt_img = self._downscale_if_required(image) return gt_img.to(self.device) def composite_with_background(self, image, background) -> torch.Tensor: @@ -873,8 +875,9 @@ class SplatfactoModel(Model): # This is a little bit sketchy for the SSIM loss. if "mask" in batch: # batch["mask"] : [H, W, 1] - assert batch["mask"].shape[:2] == gt_img.shape[:2] == pred_img.shape[:2] - mask = batch["mask"].to(self.device) + mask = self._downscale_if_required(batch["mask"]) + mask = mask.to(self.device) + assert mask.shape[:2] == gt_img.shape[:2] == pred_img.shape[:2] gt_img = gt_img * mask pred_img = pred_img * mask