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