From 8e8b98c905227dd832a894f753913718d12b8640 Mon Sep 17 00:00:00 2001 From: Maemaemaeko Date: Thu, 27 Feb 2025 00:11:41 +0900 Subject: [PATCH] fix calculation of PSNR(if only_in_mask) --- nr3d_lib/graphics/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/nr3d_lib/graphics/utils.py b/nr3d_lib/graphics/utils.py index 6f3174e..4ffdfb7 100644 --- a/nr3d_lib/graphics/utils.py +++ b/nr3d_lib/graphics/utils.py @@ -94,7 +94,9 @@ def PSNR(x: torch.Tensor, y: torch.Tensor, mask: torch.BoolTensor=None, only_in_ if (mask is not None): mask = mask.view(*x.shape[:-1]) if only_in_mask: - mse = ((x - y)**2)[mask].sum() / mask.sum().clip(1e-5) + # NOTE: The mask should be expanded to match the shape of x (or y); otherwise, the PSNR may be lower than expected. + mask_expanded = mask.view(*x.shape[:-1]).unsqueeze(-1).expand_as(x) + mse = ((x - y)**2)[mask_expanded].sum() / mask_expanded.sum().clip(1e-5) else: # NOTE: x should not be masked; # and the convergence of mask will also affect performance here (which is expected)