adding masked metrics

nerfstudio/utils/lpips_utils.py

@@ -0,0 +1,56 @@
+# Copyright 2022 The Nerfstudio Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This file contains utils for LPIPS in order to support masking.
+We effectively copy of the relevant code from torchmetrics.image.lpip
+and then return full images which we can later mask. Some functions are copied over
+simply to override functions.
+"""
+
+from torch import Tensor
+from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+
+
+def _valid_img(img: Tensor, normalize: bool) -> bool:
+    """check that input is a valid image to the network."""
+    value_check = img.max() <= 1.0 and img.min() >= 0.0 if normalize else img.min() >= -1
+    return img.ndim == 4 and img.shape[1] == 3 and value_check
+
+
+class LearnedPerceptualImagePatchSimilarityWithMasking(LearnedPerceptualImagePatchSimilarity):
+    """LearnedPerceptualImagePatchSimilarity module that will allow for masking capabilities."""
+
+    def update(self, img1: Tensor, img2: Tensor) -> None:  # pylint: disable=arguments-differ
+        """Update internal states with lpips score with masking."""
+
+        # hardcode this to True for now to avoid touching a lot of the torchmetrics code
+        self.net.spatial = True
+
+        if not (_valid_img(img1, self.normalize) and _valid_img(img2, self.normalize)):
+            raise ValueError(
+                "Expected both input arguments to be normalized tensors with shape [N, 3, H, W]."
+                f" Got input with shape {img1.shape} and {img2.shape} and values in range"
+                f" {[img1.min(), img1.max()]} and {[img2.min(), img2.max()]} when all values are"
+                f" expected to be in the {[0,1] if self.normalize else [-1,1]} range."
+            )
+        loss = self.net(img1, img2, normalize=self.normalize)
+        # now loss is the shape [batch size, H, W]
+        # we set loss to self.sum_scores to use the existing API from torchvision
+        self.sum_scores = loss  # pylint: disable=attribute-defined-outside-init
+
+    def compute(self) -> Tensor:
+        """Compute final perceptual similarity metric."""
+        # note that we don't use self.reduction anymore
+        return self.sum_scores

nerfstudio/utils/metrics.py

@@ -1,8 +1,156 @@
+# Copyright 2022 The Nerfstudio Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """
 This file contains NeRF metrics with masking capabilities.
 """
 
-from torchmetrics.functional import structural_similarity_index_measure
-from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from abc import abstractmethod
+from typing import Optional
 
-# TODO: masked PSNR, SSIM, LPIPS
+import torch
+from torch import Tensor, nn
+from torchmetrics.functional import (
+    peak_signal_noise_ratio,
+    structural_similarity_index_measure,
+)
+from torchtyping import TensorType
+
+from nerfstudio.utils.lpips_utils import (
+    LearnedPerceptualImagePatchSimilarityWithMasking,
+)
+
+
+class ImageMetricModule(nn.Module):
+    """Computes image metrics with masking capabilities.
+    We assume that the pred and target inputs are in the range [0, 1].
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.populate_modules()
+
+    def populate_modules(self):
+        """Populates the modules that will be used to compute the metric."""
+
+    @abstractmethod
+    def forward(
+        self,
+        preds: TensorType["bs", 3, "H", "W"],
+        target: TensorType["bs", 3, "H", "W"],
+        mask: Optional[TensorType["bs", 1, "H", "W"]] = None,
+    ) -> Tensor:
+        """Computes the metric.
+
+        Args:
+            preds: Predictions.
+            target: Ground truth.
+            mask: Mask to use to only compute the metrics where the mask is True.
+
+        Returns:
+            TensorType["bs"]: Metric value.
+        """
+
+
+class PSNRModule(ImageMetricModule):
+    """Computes PSNR with masking capabilities."""
+
+    def forward(
+        self,
+        preds: TensorType["bs", 3, "H", "W"],
+        target: TensorType["bs", 3, "H", "W"],
+        mask: Optional[TensorType["bs", 1, "H", "W"]] = None,
+    ) -> Tensor:
+
+        bs, h, w = preds.shape[0], preds.shape[2], preds.shape[3]
+        hw = h * w
+
+        psnr_image = peak_signal_noise_ratio(preds=preds, target=target, reduction="none", data_range=1.0, dim=1)
+        assert psnr_image.shape == (bs, h, w)
+
+        # the non-masked version
+        if mask is None:
+            psnr = psnr_image.view(bs, hw).mean(1)
+            return psnr
+
+        # the masked version
+        psnr_reshaped = psnr_image.view(bs, hw)
+        mask_reshaped = mask.view(bs, hw)
+        psnr = (psnr_reshaped * mask_reshaped).sum(-1) / hw
+        return psnr
+
+
+class SSIMModule(ImageMetricModule):
+    """Computes PSNR with masking capabilities."""
+
+    def forward(
+        self,
+        preds: TensorType["bs", 3, "H", "W"],
+        target: TensorType["bs", 3, "H", "W"],
+        mask: Optional[TensorType["bs", 1, "H", "W"]] = None,
+    ) -> Tensor:
+
+        bs, h, w = preds.shape[0], preds.shape[2], preds.shape[3]
+        hw = h * w
+
+        _, ssim_image = structural_similarity_index_measure(
+            preds=preds, target=target, reduction="none", data_range=1.0, return_full_image=True
+        )
+        ssim_image = ssim_image.mean(1)  # average over the channels
+        assert ssim_image.shape == (bs, h, w)
+
+        # the non-masked version
+        if mask is None:
+            ssim = ssim_image.view(bs, hw).mean(1)
+            return ssim
+
+        # the masked version
+        ssim_reshaped = ssim_image.view(bs, hw)
+        mask_reshaped = mask.view(bs, hw)
+        ssim = (ssim_reshaped * mask_reshaped).sum(-1) / hw
+        return ssim
+
+
+class LPIPSModule(ImageMetricModule):
+    """Computes LPIPS with masking capabilities."""
+
+    def populate_modules(self):
+        # by setting normalize=True, we assume that the pred and target inputs are in the range [0, 1]
+        self.lpips_with_masking = LearnedPerceptualImagePatchSimilarityWithMasking(normalize=True)
+
+    def forward(
+        self,
+        preds: TensorType["bs", 3, "H", "W"],
+        target: TensorType["bs", 3, "H", "W"],
+        mask: Optional[TensorType["bs", 1, "H", "W"]] = None,
+    ) -> Tensor:
+
+        bs, h, w = preds.shape[0], preds.shape[2], preds.shape[3]
+        hw = h * w
+
+        with torch.no_grad():
+            lpips_image = self.lpips_with_masking(preds, target)
+        lpips_image = lpips_image.mean(1)  # average over the channels
+        assert lpips_image.shape == (bs, h, w)
+
+        # the non-masked version
+        if mask is None:
+            lpips = lpips_image.view(bs, hw).mean(1)
+            return lpips
+
+        # the masked version
+        lpips_reshaped = lpips_image.view(bs, hw)
+        mask_reshaped = mask.view(bs, hw)
+        lpips = (lpips_reshaped * mask_reshaped).sum(-1) / hw
+        return lpips

tests/utils/test_metrics.py

@@ -0,0 +1,68 @@
+"""
+Test metrics
+"""
+import torch
+
+from nerfstudio.utils.metrics import PSNRModule, SSIMModule, LPIPSModule
+
+
+def test_psnr():
+    """Test PSNR metric"""
+    bs, h, w = 10, 100, 200
+    pred = torch.rand((bs, 3, h, w))
+    target = torch.rand((bs, 3, h, w))
+    mask = torch.rand((bs, 1, h, w)) > 0.5
+
+    psnr_module = PSNRModule()
+
+    # test non-masked version
+    psnr = psnr_module(preds=pred, target=target)
+    assert psnr.shape == (bs,)
+    assert torch.min(psnr) >= 0
+
+    # test masked version
+    psnr = psnr_module(preds=pred, target=target, mask=mask)
+    assert psnr.shape == (bs,)
+    assert torch.min(psnr) >= 0
+
+
+def test_ssim():
+    """Test SSIM metrics"""
+    bs, h, w = 10, 100, 200
+    pred = torch.rand((bs, 3, h, w))
+    target = torch.rand((bs, 3, h, w))
+    mask = torch.rand((bs, 1, h, w)) > 0.5
+
+    ssim_module = SSIMModule()
+
+    # test non-masked version
+    ssim = ssim_module(preds=pred, target=target)
+    assert ssim.shape == (bs,)
+
+    # test masked version
+    ssim = ssim_module(preds=pred, target=target, mask=mask)
+    assert ssim.shape == (bs,)
+
+
+def test_lpips():
+    """Test LPIPS metrics"""
+    bs, h, w = 10, 100, 200
+    pred = torch.rand((bs, 3, h, w))
+    target = torch.rand((bs, 3, h, w))
+    mask = torch.rand((bs, 1, h, w)) > 0.5
+
+    lpips_module = LPIPSModule()
+
+    # test non-masked version
+    lpips = lpips_module(preds=pred, target=target)
+    assert lpips.shape == (bs,)
+
+    # test masked version
+    lpips = lpips_module(preds=pred, target=target, mask=mask)
+    assert lpips.shape == (bs,)
+
+
+if __name__ == "__main__":
+    test_psnr()
+    test_ssim()
+    test_lpips()