diff --git a/nerfstudio/scripts/metrics.py b/nerfstudio/scripts/metrics.py new file mode 100644 index 000000000..613605ae6 --- /dev/null +++ b/nerfstudio/scripts/metrics.py @@ -0,0 +1,299 @@ +# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. 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. +from __future__ import annotations + +import glob +import json +import os +import sys +from collections import defaultdict +from dataclasses import dataclass +from pathlib import Path +from typing import Union + +import mediapy as media +import numpy as np +import torch +import torch.nn.functional as F +import tyro +from typing_extensions import Annotated + +from nerfstudio.utils.metrics import LPIPSModule, PSNRModule, SSIMModule + + +@dataclass +class BaseMetrics: + """Base class for metrics.""" + + input_folder: Path = Path("input-folder") + """Folder containing the renders.""" + output_folder: Path = Path("output-folder") + """Folder to save the metrics.""" + device: str = "cuda:0" + """Device to use for metrics.""" + + def main(self) -> None: + """Main function.""" + + # image metrics + psnr_module = PSNRModule().to(self.device) + ssim_module = SSIMModule().to(self.device) + lpips_module = LPIPSModule().to(self.device) + + rgb_gt_filenames = sorted(glob.glob(str(self.input_folder / "rgb_gt" / "*"))) + metrics = defaultdict() + video = [] + + experiment_name = os.path.basename(self.input_folder) + print(f"Processing experiment: {experiment_name} ...") + if len(rgb_gt_filenames) == 0: + print("No rgb_gt images found, skipping experiment") + sys.exit(0) + + for idx, rgb_gt_filename in enumerate(rgb_gt_filenames): + rgb = media.read_image(rgb_gt_filename.replace("rgb_gt", "rgb")) + rgb_gt = media.read_image(rgb_gt_filename) + + # move images to torch and to the correct device + rgb = torch.from_numpy(rgb).float().to(self.device) / 255.0 + rgb_gt = torch.from_numpy(rgb_gt).float().to(self.device) / 255.0 # (H, W, 3) + + # compute the image metrics + # reshape the images to (1, C, H, W) + x = rgb.permute(2, 0, 1).unsqueeze(0) # (1, 3, H, W) + x_gt = rgb_gt.permute(2, 0, 1).unsqueeze(0) # (1, 3, H, W) + + psnr = float(psnr_module(x, x_gt)[0]) + ssim = float(ssim_module(x, x_gt)[0]) + lpips = float(lpips_module(x, x_gt)[0]) + + metrics["psnr_list"].append(psnr) + metrics["ssim_list"].append(ssim) + metrics["lpips_list"].append(lpips) + + # save the images + rgb_gt = (rgb_gt * 255.0).cpu().numpy().astype(np.uint8) + rgb = (rgb * 255.0).cpu().numpy().astype(np.uint8) + image = np.concatenate([rgb_gt, rgb], axis=1) + image_filename = self.input_folder / "composited" / f"{idx:04d}.png" + image_filename.parent.mkdir(parents=True, exist_ok=True) + media.write_image(image_filename, image) + video.append(image) + + # write out the video + video_filename = self.input_folder / f"{experiment_name}.mp4" + media.write_video(video_filename, video, fps=30) + + # convert metrics dict to a proper dictionary + metrics = dict(metrics) + metrics["psnr"] = np.mean(metrics["psnr_list"]) + metrics["ssim"] = np.mean(metrics["ssim_list"]) + metrics["lpips"] = np.mean(metrics["lpips_list"]) + + for metric_name in sorted(metrics.keys()): + if "_list" not in metric_name: + print(f"{metric_name}: {metrics[metric_name]}") + + # write to a json file + metrics_filename = self.output_folder / f"{experiment_name}.json" + with open(metrics_filename, "w") as f: + json.dump(metrics, f, indent=4) + + +@dataclass +class NerfbusterMetrics(BaseMetrics): + """Compute metrics for nerfbusters on the renders.""" + + pseudo_gt_experiment_name: Path = Path("pseudo_gt_experiment_name") + """Name of the experiment to use for pseudo ground truth.""" + max_depth: float = 2 + """Maximum depth to use for metrics.""" + min_views: int = 1 + """Minimum number of views to use for metrics.""" + + def main(self) -> None: + """Main function.""" + + print("Using visibility masks from experiment: ", self.pseudo_gt_experiment_name) + + # image metrics + psnr_module = PSNRModule().to(self.device) + ssim_module = SSIMModule().to(self.device) + lpips_module = LPIPSModule().to(self.device) + + rgb_gt_filenames = sorted(glob.glob(str(self.input_folder / "rgb_gt" / "*"))) + visibility_filenames = sorted( + glob.glob(str(self.input_folder / self.pseudo_gt_experiment_name / "visibility_median_count" / "*")) + ) + metrics = defaultdict() + video = [] # images to make a video + + experiment_name = os.path.basename(self.input_folder) + print(f"Processing experiment: {experiment_name} ...") + if len(rgb_gt_filenames) == 0: + print("No rgb_gt images found, skipping experiment") + sys.exit(0) + + for idx, rgb_gt_filename in enumerate(rgb_gt_filenames): + # read in the images + depth = media.read_image(rgb_gt_filename.replace("rgb_gt", "depth")) + depth_raw = np.load(rgb_gt_filename.replace("rgb_gt", "depth_raw").replace(".png", ".npy"))[..., 0] + normals = media.read_image(rgb_gt_filename.replace("rgb_gt", "normals")) + pseudo_gt_visibility = media.read_image(visibility_filenames[idx]) + psuedo_gt_depth_raw = np.load( + visibility_filenames[idx].replace("visibility_median_count", "depth_raw").replace(".png", ".npy") + )[..., 0] + psuedo_gt_normals = media.read_image( + visibility_filenames[idx].replace("visibility_median_count", "normals") + ) + rgb = media.read_image(rgb_gt_filename.replace("rgb_gt", "rgb")) + rgb_gt = media.read_image(rgb_gt_filename) + + # move images to torch and to the correct device + depth = torch.from_numpy(depth).float().to(self.device) / 255.0 # 'depth' is a colormap + depth_raw = torch.from_numpy(depth_raw).float().to(self.device) + normals = torch.from_numpy(normals).float().to(self.device) / 255.0 + normals_raw = normals * 2.0 - 1.0 + pseudo_gt_visibility = torch.from_numpy(pseudo_gt_visibility).long().to(self.device) + psuedo_gt_depth_raw = torch.from_numpy(psuedo_gt_depth_raw).float().to(self.device) + psuedo_gt_normals = torch.from_numpy(psuedo_gt_normals).float().to(self.device) / 255.0 + psuedo_gt_normals_raw = psuedo_gt_normals * 2.0 - 1.0 + rgb = torch.from_numpy(rgb).float().to(self.device) / 255.0 + rgb_gt = torch.from_numpy(rgb_gt).float().to(self.device) / 255.0 # (H, W, 3) + + # create masks + visibilty_mask = (pseudo_gt_visibility[..., 0] >= self.min_views).float() + depth_mask = (depth_raw < self.max_depth).float() + mask = visibilty_mask * depth_mask # (H, W) + mask = mask[..., None].repeat(1, 1, 3) + + # show masked versions + # shape is (H, W, 3) + rgb_gt_masked = rgb_gt * mask + rgb_masked = rgb * mask + depth_masked = depth * mask + normals_masked = normals * mask + + # compute the image metrics + # reshape the images to (1, C, H, W) + x = rgb_masked.permute(2, 0, 1).unsqueeze(0) # (1, 3, H, W) + x_gt = rgb_gt_masked.permute(2, 0, 1).unsqueeze(0) # (1, 3, H, W) + m = mask.permute(2, 0, 1).unsqueeze(0)[:, 0:1] # (1, 1, H, W) + psnr = float(psnr_module(x, x_gt, m)[0]) + ssim = float(ssim_module(x, x_gt, m)[0]) + lpips = float(lpips_module(x, x_gt, m)[0]) + + # depth + depth_mask = mask[..., 0] == 1 + depth_mse = float(F.mse_loss(depth_raw[depth_mask], psuedo_gt_depth_raw[depth_mask])) + + # disparity + disparity_raw = 1.0 / depth_raw + psuedo_gt_disparity_raw = 1.0 / psuedo_gt_depth_raw + disparity = float(torch.abs(disparity_raw[depth_mask] - psuedo_gt_disparity_raw[depth_mask]).mean()) + + # make sure the normals raw are normalized + normals_raw = normals_raw / torch.norm(normals_raw, dim=-1, keepdim=True) + psuedo_gt_normals_raw = psuedo_gt_normals_raw / torch.norm(psuedo_gt_normals_raw, dim=-1, keepdim=True) + eps = 1e-8 + costheta = ( + (normals_raw[mask == 1].view(-1, 3) * psuedo_gt_normals_raw[mask == 1].view(-1, 3)) + .sum(dim=-1) + .clamp(-1 + eps, 1 - eps) + ) + theta = torch.abs(torch.acos(costheta) * 180.0 / np.pi) + normals_mse = float(theta.mean()) + normals_median = float(theta.median()) + + # the angle thresholds + metrics["normals_11.25"].append(float((theta < 11.25).float().mean())) + metrics["normals_22.5"].append(float((theta < 22.5).float().mean())) + metrics["normals_30"].append(float((theta < 30).float().mean())) + + # coverage + coverage = float(mask[..., 0].sum() / visibilty_mask.sum()) + + metrics["psnr_list"].append(psnr) + metrics["ssim_list"].append(ssim) + metrics["lpips_list"].append(lpips) + metrics["depth_list"].append(depth_mse) + metrics["disparity_list"].append(disparity) + metrics["normals_list"].append(normals_mse) + metrics["normals_median_list"].append(normals_median) + metrics["coverage_list"].append(coverage) + + # save the images + rgb_gt = (rgb_gt * 255.0).cpu().numpy().astype(np.uint8) + pseudo_gt_visibility = (pseudo_gt_visibility).cpu().numpy().astype(np.uint8) + rgb_gt_masked = (rgb_gt_masked * 255.0).cpu().numpy().astype(np.uint8) + rgb_masked = (rgb_masked * 255.0).cpu().numpy().astype(np.uint8) + depth_masked = (depth_masked * 255.0).cpu().numpy().astype(np.uint8) + normals_masked = (normals_masked * 255.0).cpu().numpy().astype(np.uint8) + image = np.concatenate( + [rgb_gt, pseudo_gt_visibility, rgb_gt_masked, rgb_masked, depth_masked, normals_masked], axis=1 + ) + image_filename = self.input_folder / "composited" / f"{idx:04d}.png" + image_filename.parent.mkdir(parents=True, exist_ok=True) + media.write_image(image_filename, image) + video.append(image) + + # write out the video + video_filename = self.input_folder / f"{experiment_name}.mp4" + media.write_video(video_filename, video, fps=30) + + # convert metrics dict to a proper dictionary + metrics = dict(metrics) + metrics["psnr"] = np.mean(metrics["psnr_list"]) + metrics["ssim"] = np.mean(metrics["ssim_list"]) + metrics["lpips"] = np.mean(metrics["lpips_list"]) + metrics["depth"] = np.mean(metrics["depth_list"]) + metrics["disparity"] = np.mean(metrics["disparity_list"]) + metrics["normals"] = np.mean(metrics["normals_list"]) + metrics["normals_median"] = np.mean(metrics["normals_median_list"]) + metrics["coverage"] = np.mean(metrics["coverage_list"]) + metrics["normals_11.25"] = np.mean(metrics["normals_11.25"]) + metrics["normals_22.5"] = np.mean(metrics["normals_22.5"]) + metrics["normals_30"] = np.mean(metrics["normals_30"]) + + for metric_name in sorted(metrics.keys()): + if "_list" not in metric_name: + print(f"{metric_name}: {metrics[metric_name]}") + + # write to a json file + metrics_filename = self.output_folder / f"{experiment_name}.json" + with open(metrics_filename, "w") as f: + json.dump(metrics, f, indent=4) + + +Commands = tyro.conf.FlagConversionOff[ + Union[ + Annotated[NerfbusterMetrics, tyro.conf.subcommand(name="nerfbusters-metrics")], + Annotated[BaseMetrics, tyro.conf.subcommand(name="base-metrics")], + ] +] + + +def entrypoint(): + """Entrypoint for use with pyproject scripts.""" + tyro.extras.set_accent_color("bright_yellow") + tyro.cli(Commands).main() + + +if __name__ == "__main__": + entrypoint() + + +def get_parser_fn(): + """Get the parser function for the sphinx docs.""" + return tyro.extras.get_parser(Commands) # noqa diff --git a/nerfstudio/utils/lpips_utils.py b/nerfstudio/utils/lpips_utils.py new file mode 100644 index 000000000..a8b118d4d --- /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): + """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 new file mode 100644 index 000000000..3482004e8 --- /dev/null +++ b/nerfstudio/utils/metrics.py @@ -0,0 +1,149 @@ +# 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 abc import abstractmethod +from typing import Optional + +import torch +from torch import nn +from torchmetrics.functional import 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, + ) -> TensorType["bs"]: + """Computes the metric. + Args: + preds: Predictions. + target: Ground truth. + mask: Mask to use to only compute the metrics where the mask is True. + Returns: + 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, + ) -> TensorType["bs"]: + bs, h, w = preds.shape[0], preds.shape[2], preds.shape[3] + hw = h * w + + preds_reshaped = preds.view(bs, 3, hw) + target_reshaped = target.view(bs, 3, hw) + num = (preds_reshaped - target_reshaped) ** 2 + # the non-masked version + if mask is None: + den = hw + else: + mask_reshaped = mask.view(bs, 1, hw) + num = num * mask_reshaped + den = mask_reshaped.sum(-1) + mse = num.sum(-1) / den + psnr = 10 * torch.log10(1.0 / mse) + psnr = psnr.mean(-1) + 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, + ) -> TensorType["bs"]: + 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) + den = mask_reshaped.sum(-1, keepdim=True) + ssim = (ssim_reshaped * mask_reshaped / den).sum(-1) + 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, + ) -> TensorType["bs"]: + 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) + den = mask_reshaped.sum(-1, keepdim=True) + lpips = (lpips_reshaped * mask_reshaped / den).sum(-1) + return lpips diff --git a/pyproject.toml b/pyproject.toml index e4de87116..c8bff54bd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -117,6 +117,7 @@ ns-train = "nerfstudio.scripts.train:entrypoint" ns-viewer = "nerfstudio.scripts.viewer.run_viewer:entrypoint" ns-eval = "nerfstudio.scripts.eval:entrypoint" ns-render = "nerfstudio.scripts.render:entrypoint" +ns-metrics = "nerfstudio.scripts.metrics:entrypoint" ns-export = "nerfstudio.scripts.exporter:entrypoint" ns-dev-test = "nerfstudio.scripts.github.run_actions:entrypoint" ns-dev-sync-viser-message-defs = "nerfstudio.scripts.viewer.sync_viser_message_defs:entrypoint"