tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

Merge branch 'master' of github.com:ethanweber/mattport

Browse Source
This commit is contained in:
evonneng 2022-05-11 00:09:42 +00:00
commit 6a77ca89aa
4 changed files with 40 additions and 101 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.vscode/settings.json vendored
View File

@ -26,5 +26,4 @@
".vscode/*.py",
"external/**/*.py",
],
"python.analysis.typeCheckingMode": "basic",
}

2
environment/requirements.txt
View File

@ -10,9 +10,9 @@ mediapy
plotly==5.7.0
pylint==2.13.4
pytest
pytorch-msssim>=0.2.1
tensorboard==2.8.0
torch==1.11.0
torchmetrics[image]==0.8.2
torchtyping
torchvision==0.12.0
tqdm

75
mattport/nerf/graph/vanilla_nerf.py
View File

@ -8,8 +8,9 @@ from typing import Dict, List, Tuple
import torch
from torch import nn
from torch.nn import Parameter
from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
from mattport.nerf import metrics
from mattport.nerf.field_modules.encoding import NeRFEncoding
from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHeadNames, RGBFieldHead
from mattport.nerf.field_modules.mlp import MLP
@ -123,6 +124,11 @@ class NeRFGraph(Graph):
# losses
self.rgb_loss = MSELoss()
# metrics
self.psnr = PeakSignalNoiseRatio(data_range=1.0)
self.ssim = StructuralSimilarityIndexMeasure()
self.lpips = LearnedPerceptualImagePatchSimilarity()
def get_param_groups(self) -> Dict[str, List[Parameter]]:
"""Obtain the parameter groups for the optimizers
@ -164,30 +170,14 @@ class NeRFGraph(Graph):
fine_renderer_accumulation = self.renderer_accumulation(fine_weights) # RendererOutputs
fine_renderer_depth = self.renderer_depth(fine_weights, pdf_ray_samples.ts)
# TODO refactor this into "vis" section. Doesn't need to be run during training.
coarse_renderer_depth = visualization.apply_depth_colormap(
coarse_renderer_depth.depth,
accumulation=coarse_renderer_accumulation.accumulation,
near_plane=self.near_plane,
far_plane=self.far_plane,
)
fine_renderer_depth = visualization.apply_depth_colormap(
fine_renderer_depth.depth,
accumulation=fine_renderer_accumulation.accumulation,
near_plane=self.near_plane,
far_plane=self.far_plane,
)
coarse_renderer_accumulation = visualization.apply_colormap(coarse_renderer_accumulation.accumulation)
fine_renderer_accumulation = visualization.apply_colormap(fine_renderer_accumulation.accumulation)
# outputs:
outputs = {
"rgb_coarse": coarse_renderer_outputs.rgb,
"rgb_fine": fine_renderer_outputs.rgb,
"accumulation_coarse": coarse_renderer_accumulation,
"accumulation_fine": fine_renderer_accumulation,
"depth_coarse": coarse_renderer_depth,
"depth_fine": fine_renderer_depth,
"accumulation_coarse": coarse_renderer_accumulation.accumulation,
"accumulation_fine": fine_renderer_accumulation.accumulation,
"depth_coarse": coarse_renderer_depth.depth,
"depth_fine": fine_renderer_depth.depth,
}
return outputs
@ -203,43 +193,56 @@ class NeRFGraph(Graph):
def log_test_image_outputs(self, image_idx, step, image, outputs):
rgb_coarse = outputs["rgb_coarse"]
rgb_fine = outputs["rgb_fine"]
accumulation_coarse = outputs["accumulation_coarse"]
accumulation_fine = outputs["accumulation_fine"]
depth_coarse = outputs["depth_coarse"]
depth_fine = outputs["depth_fine"]
combined_image = torch.cat([image, rgb_coarse, rgb_fine], dim=1)
writer.write_event(
{"name": f"image_idx_{image_idx}-rgb_coarse_fine", "x": combined_image, "step": step, "group": "val_img"}
)
accumulation_coarse = visualization.apply_colormap(outputs["accumulation_coarse"])
accumulation_fine = visualization.apply_colormap(outputs["accumulation_fine"])
combined_image = torch.cat([accumulation_coarse, accumulation_fine], dim=1)
writer.write_event(
{"name": f"image_idx_{image_idx}", "x": combined_image, "step": step, "group": "val_accumulation"}
)
depth_coarse = visualization.apply_depth_colormap(
outputs["depth_coarse"],
accumulation=outputs["accumulation_coarse"],
near_plane=self.near_plane,
far_plane=self.far_plane,
)
depth_fine = visualization.apply_depth_colormap(
outputs["depth_fine"],
accumulation=outputs["accumulation_fine"],
near_plane=self.near_plane,
far_plane=self.far_plane,
)
combined_image = torch.cat([depth_coarse, depth_fine], dim=1)
writer.write_event({"name": f"image_idx_{image_idx}", "x": combined_image, "step": step, "group": "val_depth"})
coarse_psnr = metrics.get_psnr(image, rgb_coarse)
# Switch images from [H, W, C] to [1, C, H, W] for metrics computations
image = torch.moveaxis(image, -1, 0)[None, ...]
rgb_coarse = torch.moveaxis(rgb_coarse, -1, 0)[None, ...]
rgb_fine = torch.moveaxis(rgb_fine, -1, 0)[None, ...]
coarse_psnr = self.psnr(image, rgb_coarse)
writer.write_event(
{"name": f"image_idx_{image_idx}", "scalar": float(coarse_psnr), "step": step, "group": "val"}
{"name": f"val_{image_idx}-coarse", "scalar": float(coarse_psnr), "step": step, "group": "psnr"}
)
coarse_ssim = metrics.get_ssim(image, rgb_coarse)
writer.write_event(
{"name": f"image_idx_{image_idx}", "scalar": float(coarse_ssim), "step": step, "group": "ssim"}
)
fine_psnr = metrics.get_psnr(image, rgb_fine)
fine_psnr = self.psnr(image, rgb_fine)
stats_tracker.update_stats(
{"name": stats_tracker.Stats.CURR_TEST_PSNR, "value": float(fine_psnr), "step": step}
)
writer.write_event(
{"name": f"image_idx_{image_idx}-fine_psnr", "scalar": float(fine_psnr), "step": step, "group": "val"}
{"name": f"val_idx_{image_idx}-fine", "scalar": float(fine_psnr), "step": step, "group": "psnr"}
)
fine_ssim = metrics.get_ssim(image, rgb_fine)
fine_ssim = self.ssim(image, rgb_fine)
writer.write_event({"name": f"val_idx_{image_idx}", "scalar": float(fine_ssim), "step": step, "group": "ssim"})
fine_lpips = self.lpips(image, rgb_fine)
writer.write_event(
{"name": f"image_idx_{image_idx}-fine_ssim", "scalar": float(fine_ssim), "step": step, "group": "ssim"}
{"name": f"val_idx_{image_idx}", "scalar": float(fine_lpips), "step": step, "group": "lpips"}
)

63
mattport/nerf/metrics.py
View File

@ -1,63 +0,0 @@
"""
Collection of metrics.
"""
import torch
from pytorch_msssim import ssim, ms_ssim
from torchtyping import TensorType
def get_psnr(image_gt: TensorType, image_pred: TensorType) -> TensorType:
"""
Returns the PSNR (peak signal-to-noise ratio) for an image. TODO: add masking
Args:
image_gt (TensorType): Ground truth image
image_pred (TensorType): Predicted image
Returns:
float: PSNR
"""
mean = torch.mean((image_gt - image_pred) ** 2)
return -10 * torch.log10(mean)
def get_ssim(image_gt: TensorType, image_pred: TensorType) -> TensorType:
"""
Returns the SSIM (structural similarty) for an image.
Args:
image_gt (TensorType): Ground truth image
image_pred (TensorType): Predicted image
Returns:
float: SSIM
"""
assert image_gt.shape[-1] in [1, 3]
assert image_pred.shape[-1] in [1, 3]
assert image_gt.shape == image_pred.shape
if image_gt.dim() == 3:
image_gt = image_gt[None, ...]
if image_pred.dim() == 3:
image_pred = image_pred[None, ...]
return ssim(image_gt, image_pred, data_range=1.0)
def get_ms_ssim(image_gt: TensorType, image_pred: TensorType) -> TensorType:
"""
Returns the MS_SSIM (multi-scale structural similarty) for an image.
Args:
image_gt (TensorType): Ground truth image
image_pred (TensorType): Predicted image
Returns:
float: MS-SSIM
"""
assert image_gt.shape[-1] in [1, 3]
assert image_pred.shape[-1] in [1, 3]
assert image_gt.shape == image_pred.shape
if image_gt.dim() == 3:
image_gt = image_gt[None, ...]
if image_pred.dim() == 3:
image_pred = image_pred[None, ...]
return ms_ssim(image_gt, image_pred, data_range=1.0)