mirror of
https://github.com/graphdeco-inria/gaussian-splatting
synced 2024-11-22 16:28:32 +00:00
88 lines
3.4 KiB
Python
88 lines
3.4 KiB
Python
|
from pathlib import Path
|
||
|
|
import os
|
||
|
|
from PIL import Image
|
||
|
|
import torch
|
||
|
|
import torchvision.transforms.functional as tf
|
||
|
|
from utils.loss_utils import ssim
|
||
|
|
from lpipsPyTorch import lpips
|
||
|
|
import json
|
||
|
|
from tqdm import tqdm
|
||
|
|
from utils.image_utils import psnr
|
||
|
|
from argparse import ArgumentParser
|
||
|
|
|
||
|
|
def readImages(renders_dir, gt_dir):
|
||
|
|
renders = []
|
||
|
|
gts = []
|
||
|
|
image_names = []
|
||
|
|
for fname in os.listdir(renders_dir):
|
||
|
|
render = Image.open(renders_dir / fname)
|
||
|
|
gt = Image.open(gt_dir / fname)
|
||
|
|
renders.append(tf.to_tensor(render).unsqueeze(0)[:, :3, :, :].cuda())
|
||
|
|
gts.append(tf.to_tensor(gt).unsqueeze(0)[:, :3, :, :].cuda())
|
||
|
|
image_names.append(fname)
|
||
|
|
return renders, gts, image_names
|
||
|
|
|
||
|
|
def evaluate(model_paths):
|
||
|
|
|
||
|
|
full_dict = {}
|
||
|
|
per_view_dict = {}
|
||
|
|
full_dict_polytopeonly = {}
|
||
|
|
per_view_dict_polytopeonly = {}
|
||
|
|
|
||
|
|
for scene_dir in model_paths:
|
||
|
|
print("Scene:", scene_dir)
|
||
|
|
full_dict[scene_dir] = {}
|
||
|
|
per_view_dict[scene_dir] = {}
|
||
|
|
full_dict_polytopeonly[scene_dir] = {}
|
||
|
|
per_view_dict_polytopeonly[scene_dir] = {}
|
||
|
|
|
||
|
|
test_dir = Path(scene_dir) / "test"
|
||
|
|
|
||
|
|
for method in os.listdir(test_dir):
|
||
|
|
print("Method:", method)
|
||
|
|
|
||
|
|
full_dict[scene_dir][method] = {}
|
||
|
|
per_view_dict[scene_dir][method] = {}
|
||
|
|
full_dict_polytopeonly[scene_dir][method] = {}
|
||
|
|
per_view_dict_polytopeonly[scene_dir][method] = {}
|
||
|
|
|
||
|
|
method_dir = test_dir / method
|
||
|
|
gt_dir = method_dir/ "gt"
|
||
|
|
renders_dir = method_dir / "renders"
|
||
|
|
renders, gts, image_names = readImages(renders_dir, gt_dir)
|
||
|
|
|
||
|
|
ssims = []
|
||
|
|
psnrs = []
|
||
|
|
lpipss = []
|
||
|
|
|
||
|
|
for idx in tqdm(range(len(renders)), desc="Metric evaluation progress"):
|
||
|
|
ssims.append(ssim(renders[idx], gts[idx]))
|
||
|
|
psnrs.append(psnr(renders[idx], gts[idx]))
|
||
|
|
lpipss.append(lpips(renders[idx], gts[idx], net_type='vgg'))
|
||
|
|
|
||
|
|
print("SSIM: {}".format(torch.tensor(ssims).mean()))
|
||
|
|
print("PSNR: {}".format(torch.tensor(psnrs).mean()))
|
||
|
|
print("LPIPS: {}".format(torch.tensor(lpipss).mean()))
|
||
|
|
|
||
|
|
full_dict[scene_dir][method].update({"SSIM": torch.tensor(ssims).mean().item(),
|
||
|
|
"PSNR": torch.tensor(psnrs).mean().item(),
|
||
|
|
"LPIPS": torch.tensor(lpipss).mean().item()})
|
||
|
|
per_view_dict[scene_dir][method].update({"SSIM": {name: ssim for ssim, name in zip(torch.tensor(ssims).tolist(), image_names)},
|
||
|
|
"PSNR": {name: psnr for psnr, name in zip(torch.tensor(psnrs).tolist(), image_names)},
|
||
|
|
"LPIPS": {name: lp for lp, name in zip(torch.tensor(lpipss).tolist(), image_names)}})
|
||
|
|
|
||
|
|
with open(scene_dir + "/results.json", 'w') as fp:
|
||
|
|
json.dump(full_dict[scene_dir], fp, indent=True)
|
||
|
|
with open(scene_dir + "/per_view.json", 'w') as fp:
|
||
|
|
json.dump(per_view_dict[scene_dir], fp, indent=True)
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
device = torch.device("cuda:0")
|
||
|
|
torch.cuda.set_device(device)
|
||
|
|
|
||
|
|
# Set up command line argument parser
|
||
|
|
parser = ArgumentParser(description="Training script parameters")
|
||
|
|
parser.add_argument('--model_paths', '-m', required=True, nargs="+", type=str, default=[])
|
||
|
|
args = parser.parse_args()
|
||
|
|
evaluate(args.model_paths)
|