| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import torch |
| from scene import Scene |
| import os |
| import sys |
| from tqdm import tqdm |
| from os import makedirs |
| from scene import Scene, GaussianModel |
| 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, Namespace |
| from utils.sh_utils import RGB2SH, SH2RGB |
| from predictor import FaceMeshPredictor |
| import torch.nn.functional as F |
| |
|
|
|
|
| def rot_mat_from_6dof(v: torch.Tensor) -> torch.Tensor: |
| assert v.shape[-1] == 6 |
| v = v.view(-1, 6) |
| vx, vy = v[..., :3].clone(), v[..., 3:].clone() |
|
|
| b1 = F.normalize(vx, dim=-1) |
| b3 = F.normalize(torch.cross(b1, vy), dim=-1) |
| b2 = - torch.cross(b1, b3) |
|
|
| return torch.stack((b1, b2, b3), dim=-1) |
| def render_set(use_sh, 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") |
| flame_predictor = FaceMeshPredictor.dad_3dnet() |
| makedirs(render_path, exist_ok=True) |
| makedirs(gts_path, exist_ok=True) |
| verts_noaling = gaussians.get_xyz.unsqueeze(0) |
| for idx, view in enumerate(tqdm(views, desc="Rendering progress")): |
| gt = view.original_image[0:3, :, :] |
| gt_input = gt.permute(1, 2, 0).cpu().numpy()[::2, ::2, :] * 255 |
| gt_parameters = flame_predictor(gt_input) |
| rotation_mat = rot_mat_from_6dof(gt_parameters['flame_params'].rotation).type(verts_noaling.dtype).cuda() |
| |
| |
| |
| |
| scale = torch.clamp(gt_parameters['flame_params'].scale[:, None] + 1.0, 1e-8).cuda() |
| gaussians._xyz = verts_noaling.squeeze(0) * scale.squeeze() |
| translation = gt_parameters['flame_params'].translation.cuda() |
| translation[..., 2] = 2183.5264 |
| view.mk_new_cam(R=rotation_mat.squeeze(), T=translation.squeeze()) |
| if not use_sh: |
| rendering = RGB2SH(render(view, gaussians, pipeline, background)["render"]) |
| else: |
| rendering = render(view, gaussians, pipeline, background)["render"] |
| |
| 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, dataset.data_device, dataset.xyz_grad, dataset.align, dataset.free_flame) |
| |
| 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=dataset.data_device) |
|
|
| if not skip_train: |
| render_set(dataset.use_sh, dataset.model_path, "train", scene.loaded_iter, scene.getTrainCameras(), gaussians, pipeline, background) |
|
|
| if not skip_test: |
| render_set(dataset.use_sh, dataset.model_path, "test", scene.loaded_iter, scene.getTestCameras(), gaussians, pipeline, background) |
|
|
| if __name__ == "__main__": |
| |
| 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) |
|
|
| |
| safe_state(args.quiet) |
|
|
| render_sets(model.extract(args), args.iteration, pipeline.extract(args), args.skip_train, args.skip_test) |
