render.py

#
# 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
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
import cv2
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
from scene import Camera
from scene.dataset_readers import get_dataset_reader
from shutil import copy
import numpy as np

class Dummy(object):
    pass

def render_set(model_params : ModelParams, iteration, background, experiment_folder, channel, ext : str = ".png"):
    model_path = model_params.model_path
    source_path = model_params.source_path

    render_path = os.path.join(model_path, experiment_folder, "test", "ours_{}".format(iteration), channel, "renders")
    gts_path = os.path.join(model_path, experiment_folder, "test", "ours_{}".format(iteration), channel, "gt")

    makedirs(render_path, exist_ok=True)
    makedirs(gts_path, exist_ok=True)

    pipe = Dummy()
    pipe.debug = False
    pipe.antialiasing = False
    pipe.compute_cov3D_python = False
    pipe.convert_SHs_python = False

    ply_dir = os.path.join(model_params.model_path, "point_cloud", f"iteration_{iteration}", experiment_folder, channel + ".ply")
    gaussians = GaussianModel(0)
    if channel == "albedo":
        gaussians.use_color_activation = True
    scene = Scene(model_params, gaussians, init_pc_path=ply_dir, out_folder=experiment_folder)

    views = scene.getTestCameras()
    gt_source_dir = os.path.join(source_path, channel)
    use_colors_activation = channel == "albedo"
    for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
        rendering = render(view, gaussians, pipe, background, colors_activation=use_colors_activation)["render"]
        # gt = view.original_image[0:3, :, :] 
        gt_source_path = os.path.join(gt_source_dir, view.image_name + ext)
        render_out_path = os.path.join(render_path, '{0:05d}'.format(idx) + ext)
        
        rendering = rendering.detach().permute(1, 2, 0).cpu().numpy()
        if ext == ".png":
            rendering = (rendering / 255.0).astype(np.uint8)

        cv2.imwrite(render_out_path, rendering[:, :, ::-1]) # cv2 BGR convention but our rendering is in RGB

        copy(gt_source_path, os.path.join(gts_path, '{0:05d}'.format(idx) + ext))
        # torchvision.utils.save_image(gt, os.path.join(gts_path, '{0:05d}'.format(idx) + ext))

def render_full(model_params : ModelParams, iteration, background, experiment_folder, ext : str = ".png"):
    model_path = model_params.model_path
    source_path = model_params.source_path

    render_path = os.path.join(model_path, experiment_folder, "test", "ours_{}".format(iteration), "full", "renders")
    gts_path = os.path.join(model_path, experiment_folder, "test", "ours_{}".format(iteration), "full", "gt")

    makedirs(render_path, exist_ok=True)
    makedirs(gts_path, exist_ok=True)

    pipe = Dummy()
    pipe.debug = False
    pipe.antialiasing = False
    pipe.compute_cov3D_python = False
    pipe.convert_SHs_python = False

    albedo_ply_path = os.path.join(model_params.model_path, "point_cloud", f"iteration_{iteration}", experiment_folder, "albedo.ply")
    shading_ply_path = os.path.join(model_params.model_path, "point_cloud", f"iteration_{iteration}", experiment_folder, "shading.ply")
    residual_path = os.path.join(model_params.model_path, "point_cloud", f"iteration_{iteration}", experiment_folder, "residual.safetensors")
    
    albedo_gaussians = GaussianModel(0)
    albedo_gaussians.use_color_activation = True
    albedo_gaussians.load_ply(albedo_ply_path)
    
    shading_gaussians = GaussianModel(0)
    shading_gaussians.load_ply(shading_ply_path)

    residual_gaussians = GaussianModel(sh_degree=3, grayscale_dc=True)
    residual_gaussians.load_safetensors(residual_path)

    convert_images_to_linear = ext == ".png"

    dataset_reader = get_dataset_reader(model_params.source_path, images_dir="images", depths_dir=model_params.depths, eval=model_params.eval, llffhold=model_params.llffhold, convert_images_to_linear=convert_images_to_linear)
    scene_info = dataset_reader.read()

    views = scene_info.test_cameras
    # dummy_scene = Scene(model_params, albedo_gaussians, out_folder=experiment_folder)

    # views = dummy_scene.getTestCameras()
    gt_source_dir = os.path.join(source_path, "images")

    for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
        rend_albedo = render(view, albedo_gaussians, pipe, background, colors_activation=True)["render"]
        rend_shading = render(view, shading_gaussians, pipe, background, colors_activation=False)["render"]
        rend_residual = render(view, residual_gaussians, pipe, background, colors_activation=False, enable_shs = True)["render"]

        rendering = (rend_albedo * rend_shading) + rend_residual

        # gt = view.original_image[0:3, :, :] 
        gt_source_path = os.path.join(gt_source_dir, view.image_name + ext)
        render_out_path = os.path.join(render_path, '{0:05d}'.format(idx) + ext)
        
        rendering = rendering.detach().permute(1, 2, 0).cpu().numpy()
        if ext == ".png":
            rendering = (rendering / 255.0).astype(np.uint8)

        cv2.imwrite(render_out_path, rendering[:, :, ::-1]) # cv2 BGR convention but our rendering is in RGB

        copy(gt_source_path, os.path.join(gts_path, '{0:05d}'.format(idx) + ext))
        # torchvision.utils.save_image(gt, os.path.join(gts_path, '{0:05d}'.format(idx) + ext))

def render_sets(dataset : ModelParams, iteration : int, experiment_folder: str):
    with torch.no_grad():

        # always save imags in the same format as gt images to compare
        ext = os.listdir(os.path.join(dataset.source_path, dataset.images))[0][-4:]
        # no need for scene in theory, just 3 sets of gaussians. Then write same extension of images as the input (png or OpenExr)
        # scene = Scene(dataset, gaussians, use_depths_params_file=False, convert_images_to_linear = srgb_input, out_folder=experiment_folder)
        #TODO implement rendering (add parameters to render albedo, shading, diffuse and glossy maybe?)
        bg_color = [1,1,1] if dataset.white_background else [0, 0, 0]
        background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
        ply_dir = os.path.join(args.model_path, "point_cloud", f"iteration_{iteration}", experiment_folder)
        render_set(dataset, iteration, background, experiment_folder, channel="albedo", ext=ext)

        if "shading.ply" in os.listdir(ply_dir):
            render_set(dataset, iteration, background, experiment_folder, channel="shading", ext=ext)
            render_full(dataset, iteration, background, experiment_folder, ext=ext)

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("--experiment_folder", "-e", type=str)
    parser.add_argument("--quiet", action="store_true")
    args = get_combined_args(parser)
    print("Rendering " + args.model_path)

    # Initialize system state (RNG)
    safe_state(args.quiet)

    render_sets(model.extract(args), args.iteration, args.experiment_folder)