File size: 2,079 Bytes
d62394f |
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 53 54 55 56 57 58 |
import argparse
import torch
from model import GaussianSplatting2D
from utils.misc_utils import load_cfg
def get_gaussian_cfg(args):
gaussian_cfg = f"num-{args.num_gaussians:d}"
if args.disable_inverse_scale:
gaussian_cfg += f"_scale-{args.init_scale:.1f}"
else:
gaussian_cfg += f"_inv-scale-{args.init_scale:.1f}"
if not args.quantize:
args.pos_bits, args.scale_bits, args.rot_bits, args.feat_bits = 32, 32, 32, 32
min_bits = min(args.pos_bits, args.scale_bits, args.rot_bits, args.feat_bits)
max_bits = max(args.pos_bits, args.scale_bits, args.rot_bits, args.feat_bits)
if min_bits < 4 or max_bits > 32:
raise ValueError(
f"Bit precision must be between 4 and 32 but got: {args.pos_bits:d}, {args.scale_bits:d}, {args.rot_bits:d}, {args.feat_bits:d}"
)
gaussian_cfg += f"_bits-{args.pos_bits:d}-{args.scale_bits:d}-{args.rot_bits:d}-{args.feat_bits:d}"
if not args.disable_topk_norm:
gaussian_cfg += f"_top-{args.topk:d}"
gaussian_cfg += f"_{args.init_mode[0]}-{args.init_random_ratio:.1f}"
return gaussian_cfg
def get_log_dir(args):
gaussian_cfg = get_gaussian_cfg(args)
loss_cfg = f"l1-{args.l1_loss_ratio:.1f}_l2-{args.l2_loss_ratio:.1f}_ssim-{args.ssim_loss_ratio:.1f}"
folder = f"{gaussian_cfg}_{loss_cfg}"
if args.downsample:
folder += f"_ds-{args.downsample_ratio:.1f}"
if not args.disable_lr_schedule:
folder += f"_decay-{args.max_decay_times:d}-{args.decay_ratio:.1f}"
if not args.disable_prog_optim:
folder += "_prog"
return f"{args.log_root}/{args.exp_name}/{folder}"
def main(args):
args.log_dir = get_log_dir(args)
ImageGS = GaussianSplatting2D(args)
if args.eval:
ImageGS.render(render_height=args.render_height)
else:
ImageGS.optimize()
if __name__ == "__main__":
torch.hub.set_dir("models/torch")
parser = argparse.ArgumentParser()
parser = load_cfg(cfg_path="cfgs/default.yaml", parser=parser)
arguments = parser.parse_args()
main(arguments)
|