Spaces:
Sleeping
Sleeping
| # | |
| # 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 torch | |
| from scene import Scene | |
| import os | |
| from tqdm import tqdm | |
| from os import makedirs | |
| from gaussian_renderer import render | |
| import torchvision | |
| from utils.general_utils import safe_state | |
| from argparse import ArgumentParser | |
| from arguments import ModelParams, PipelineParams, get_combined_args | |
| from gaussian_renderer import GaussianModel | |
| def render_set(model_path, name, iteration, views, gaussians, pipeline, background): | |
| render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders") | |
| gts_path = os.path.join(model_path, name, "ours_{}".format(iteration), "gt") | |
| makedirs(render_path, exist_ok=True) | |
| makedirs(gts_path, exist_ok=True) | |
| for idx, view in enumerate(tqdm(views, desc="Rendering progress")): | |
| rendering = render(view, gaussians, pipeline, background)["render"] | |
| 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")) | |
| def render_sets(dataset : ModelParams, iteration : int, pipeline : PipelineParams, skip_train : bool, skip_test : bool): | |
| with torch.no_grad(): | |
| gaussians = GaussianModel(dataset.sh_degree) | |
| scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False) | |
| bg_color = [1,1,1] if dataset.white_background else [0, 0, 0] | |
| background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") | |
| if not skip_train: | |
| render_set(dataset.model_path, "train", scene.loaded_iter, scene.getTrainCameras(), gaussians, pipeline, background) | |
| if not skip_test: | |
| render_set(dataset.model_path, "test", scene.loaded_iter, scene.getTestCameras(), gaussians, pipeline, background) | |
| if __name__ == "__main__": | |
| # Set up command line argument parser | |
| parser = ArgumentParser(description="Testing script parameters") | |
| model = ModelParams(parser, sentinel=True) | |
| pipeline = PipelineParams(parser) | |
| parser.add_argument("--iteration", 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 = get_combined_args(parser) | |
| print("Rendering " + args.model_path) | |
| # Initialize system state (RNG) | |
| render_sets(model.extract(args), args.iteration, pipeline.extract(args), args.skip_train, args.skip_test) |