File size: 1,583 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 | import argparse, json
from pathlib import Path
import torch
from arguments import ModelParams, PipelineParams, OptimizationParams
from scene import Scene, FourierModel
parser = argparse.ArgumentParser()
lp = ModelParams(parser); op = OptimizationParams(parser); pp = PipelineParams(parser)
parser.add_argument('--load_iteration', type=int, default=6000)
args = parser.parse_args()
dataset = lp.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)
out = Path(dataset.model_path) / 'lod_models'
out.mkdir(parents=True, exist_ok=True)
full_r = gaussians._radii_logits.detach().clone()
full_p = gaussians._phases_raw.detach().clone()
full_k = gaussians.fourier_K
rows = []
for k in range(1, full_k+1):
gaussians._radii_logits = torch.nn.Parameter(full_r[:, :k].clone().requires_grad_(False))
gaussians._phases_raw = torch.nn.Parameter(full_p[:, :k].clone().requires_grad_(False))
gaussians.fourier_K = k
gaussians.active_K = k
path = out / f'fourier_k{k}.ply'
gaussians.save_ply(str(path))
rows.append({'K':k,'path':str(path),'bytes':path.stat().st_size,'primitives':int(gaussians.get_xyz.shape[0])})
gaussians._radii_logits = torch.nn.Parameter(full_r)
gaussians._phases_raw = torch.nn.Parameter(full_p)
gaussians.fourier_K = full_k
gaussians.active_K = full_k
(out/'sizes.json').write_text(json.dumps(rows,indent=2))
print(json.dumps(rows,indent=2))
|