| 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) |
|
|
