File size: 2,712 Bytes
8e02a28 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | import argparse, json, os, sys, time
from pathlib import Path
import numpy as np
import torch
from PIL import Image
sys.path.insert(0, str(Path(__file__).resolve().parent))
from arguments import ModelParams, PipelineParams, OptimizationParams
from scene import Scene, FourierModel
from gaussian_renderer import render
from utils.image_utils import psnr
from utils.loss_utils import ssim
from lpipsPyTorch import lpips as compute_lpips
parser = argparse.ArgumentParser(description='Fourier Splatting LOD sweep')
lp = ModelParams(parser)
op = OptimizationParams(parser)
pp = PipelineParams(parser)
parser.add_argument('--load_iteration', type=int, default=6000)
parser.add_argument('--lod_values', type=str, default='1,2,3,4,5,6')
args = parser.parse_args()
dataset, opt, pipe = lp.extract(args), op.extract(args), pp.extract(args)
gaussians = FourierModel(dataset.sh_degree, fourier_K=dataset.fourier_K,
use_beta_color=dataset.beta_color, sb_lobes=dataset.sb_lobes)
scene = Scene(dataset, gaussians, load_iteration=args.load_iteration, shuffle=False)
bg = torch.tensor([1,1,1] if dataset.white_background else [0,0,0], dtype=torch.float32, device='cuda')
lod_values = [int(x) for x in args.lod_values.split(',')]
results = []
vis_dir = Path(dataset.model_path) / 'lod_visuals'
vis_dir.mkdir(parents=True, exist_ok=True)
with torch.inference_mode():
for k in lod_values:
gaussians.active_K = min(k, gaussians.fourier_K)
pvals, svals, lvals = [], [], []
for j, cam in enumerate(scene.getTestCameras()):
pred = torch.clamp(render(cam, gaussians, pipe, bg)['render'], 0, 1).unsqueeze(0)
gt = torch.clamp(cam.original_image.cuda(), 0, 1).unsqueeze(0)
pvals.append(float(psnr(pred, gt).item()))
svals.append(float(ssim(pred, gt).item()))
lvals.append(float(compute_lpips(pred, gt, net_type='vgg').item()))
if j in (0, len(scene.getTestCameras())//2):
side = torch.cat([gt[0], pred[0]], dim=2).permute(1,2,0).cpu().numpy()
Image.fromarray((side*255).clip(0,255).astype(np.uint8)).save(vis_dir / f'k{k}_view{j}.png')
row = {'K': k, 'PSNR': float(np.mean(pvals)), 'SSIM': float(np.mean(svals)),
'LPIPS_VGG': float(np.mean(lvals)), 'primitives': int(gaussians.get_xyz.shape[0]),
'heldout_views': len(scene.getTestCameras())}
print(json.dumps(row), flush=True)
results.append(row)
out = {'checkpoint_iteration': args.load_iteration, 'scene': 'vrnerf_apartment', 'results': results}
path = Path(dataset.model_path) / 'lod_metrics.json'
path.write_text(json.dumps(out, indent=2))
print('WROTE', path, flush=True)
|