diff --git a/nerfstudio/utils/lpips_utils.py b/nerfstudio/utils/lpips_utils.py new file mode 100644 index 000000000..dd4d9b50d --- /dev/null +++ b/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 diff --git a/nerfstudio/utils/metrics.py b/nerfstudio/utils/metrics.py index f9a4b4f66..3fab13f0e 100644 --- a/nerfstudio/utils/metrics.py +++ b/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 \ No newline at end of file +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 diff --git a/tests/utils/test_metrics.py b/tests/utils/test_metrics.py new file mode 100644 index 000000000..599069924 --- /dev/null +++ b/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()