-
Notifications
You must be signed in to change notification settings - Fork 3
/
render.py
executable file
·108 lines (88 loc) · 4.46 KB
/
render.py
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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use
# under the terms of the LICENSE.md file.
#
# For inquiries contact george.drettakis@inria.fr
#
import os
import sys
import yaml
from os import makedirs
import torch
import numpy as np
import subprocess
cmd = 'nvidia-smi -q -d Memory |grep -A4 GPU|grep Used'
result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE).stdout.decode().split('\n')
os.environ['CUDA_VISIBLE_DEVICES']=str(np.argmin([int(x.split()[2]) for x in result[:-1]]))
os.system('echo $CUDA_VISIBLE_DEVICES')
from scene import Scene
import json
import time
import torchvision
from tqdm import tqdm
from utils.general_utils import safe_state, parse_cfg, get_render_func
from argparse import ArgumentParser
def render_set(base_model, model_path, name, iteration, views, gaussians, pipe, background, render_mode, ape_code):
render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders")
makedirs(render_path, exist_ok=True)
gts_path = os.path.join(model_path, name, "ours_{}".format(iteration), "gt")
makedirs(gts_path, exist_ok=True)
t_list = []
per_view_dict = {}
per_view_level_dict = {}
modules = __import__('gaussian_renderer')
for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
torch.cuda.synchronize(); t0 = time.time()
if ape_code==-1:
render_pkg = getattr(modules, get_render_func(base_model))(view, gaussians, pipe, background, iteration, render_mode)
else:
render_pkg = getattr(modules, get_render_func(base_model))(view, gaussians, pipe, background, iteration, render_mode, ape_code)
torch.cuda.synchronize(); t1 = time.time()
t_list.append(t1-t0)
rendering = torch.clamp(render_pkg["render"], 0.0, 1.0)
visible_count = render_pkg["visibility_filter"].sum()
per_view_dict['{0:05d}'.format(idx)+".png"] = visible_count.item()
gt = view.original_image[0:3, :, :]
torchvision.utils.save_image(rendering, os.path.join(render_path, '{0:05d}'.format(idx) + ".png"))
torchvision.utils.save_image(gt, os.path.join(gts_path, '{0:05d}'.format(idx) + ".png"))
t = np.array(t_list[5:])
fps = 1.0 / t.mean()
print(f'Test FPS: \033[1;35m{fps:.5f}\033[0m')
with open(os.path.join(model_path, name, "ours_{}".format(iteration), "per_view_count.json"), 'w') as fp:
json.dump(per_view_dict, fp, indent=True)
def render_sets(dataset, opt, pipe, iteration, skip_train, skip_test, ape_code):
with torch.no_grad():
modules = __import__('scene.gs_model_'+dataset.base_model, fromlist=[''])
model_config = dataset.model_config
gaussians = getattr(modules, model_config['name'])(**model_config['kwargs'])
scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False, resolution_scales=dataset.resolution_scales)
gaussians.eval()
gaussians.set_coarse_interval(opt)
if not os.path.exists(dataset.model_path):
os.makedirs(dataset.model_path)
if not skip_train:
render_set(dataset.base_model, dataset.model_path, "train", scene.loaded_iter, scene.getTrainCameras(), gaussians, pipe, scene.background, dataset.render_mode, ape_code)
if not skip_test:
render_set(dataset.base_model, dataset.model_path, "test", scene.loaded_iter, scene.getTestCameras(), gaussians, pipe, scene.background, dataset.render_mode, ape_code)
if __name__ == "__main__":
# Set up command line argument parser
parser = ArgumentParser(description="Testing script parameters")
parser.add_argument('-m', '--model_path', type=str, required=True)
parser.add_argument("--iteration", default=-1, type=int)
parser.add_argument("--ape", default=-1, type=int)
parser.add_argument("--skip_train", action="store_true")
parser.add_argument("--skip_test", action="store_true")
parser.add_argument("--quiet", action="store_true")
args = parser.parse_args(sys.argv[1:])
with open(os.path.join(args.model_path, "config.yaml")) as f:
cfg = yaml.load(f, Loader=yaml.FullLoader)
lp, op, pp = parse_cfg(cfg)
lp.model_path = args.model_path
print("Rendering " + args.model_path)
# Initialize system state (RNG)
safe_state(args.quiet)
render_sets(lp, op, pp, args.iteration, args.skip_train, args.skip_test, args.ape)