test.py

import os, sys 
import random
import numpy as np
import torch
import argparse
import cv2
import time
import datetime

from scene import GaussianModel, Scene_mica
from src.deform_model import Deform_Model
from gaussian_renderer import render
from arguments import ModelParams, PipelineParams, OptimizationParams

from utils.metrics import img_mse, img_ssim, img_psnr, perceptual


def set_random_seed(seed):
    r"""Set random seeds for everything.

    Args:
        seed (int): Random seed.
        by_rank (bool):
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

if __name__ == "__main__":
    # Set up command line argument parser
    parser = argparse.ArgumentParser(description="Training script parameters")
    lp = ModelParams(parser)
    op = OptimizationParams(parser)
    pp = PipelineParams(parser)
    parser.add_argument('--seed', type=int, default=0, help='Random seed.')
    parser.add_argument('--idname', type=str, default='id1_25', help='id name')
    parser.add_argument('--logname', type=str, default='log', help='log name')
    parser.add_argument('--image_res', type=int, default=512, help='image resolution')
    parser.add_argument("--start_checkpoint", type=str, default = None)
    args = parser.parse_args(sys.argv[1:])
    args.device = "cuda"
    lpt = lp.extract(args)
    opt = op.extract(args)
    ppt = pp.extract(args)

    batch_size = 1
    set_random_seed(args.seed)

    ## deform model
    DeformModel = Deform_Model(args.device).to(args.device)
    DeformModel.training_setup()
    DeformModel.eval()

    ## dataloader
    # data_dir = os.path.join('dataset', args.idname)
    # mica_datadir = os.path.join('metrical-tracker/output', args.idname)
    data_dir = os.path.join('data', args.idname)
    mica_datadir = os.path.join(data_dir, "tracker_output", args.idname)
    logdir = data_dir+'/'+args.logname
    scene = Scene_mica(data_dir, mica_datadir, train_type=1, white_background=lpt.white_background, device = args.device)
    
    first_iter = 0
    gaussians = GaussianModel(lpt.sh_degree)
    gaussians.training_setup(opt)

    if args.start_checkpoint:
        (model_params, gauss_params, first_iter) = torch.load(args.start_checkpoint)
        DeformModel.restore(model_params)
        gaussians.restore(gauss_params, opt)

    bg_color = [1, 1, 1] if lpt.white_background else [0, 1, 0]
    background = torch.tensor(bg_color, dtype=torch.float32, device=args.device)
    
    # fourcc = cv2.VideoWriter_fourcc(*'XVID')
    # vid_save_path = os.path.join(logdir, 'test.avi')
    fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
    vid_save_path = os.path.join(logdir, 'test.mp4')
    out = cv2.VideoWriter(vid_save_path, fourcc, 15, (args.image_res*2, args.image_res), True)

    viewpoint = scene.getCameras().copy()
    codedict = {}
    codedict['shape'] = scene.shape_param.to(args.device)
    DeformModel.example_init(codedict)

    ssim = 0
    psnr = 0
    lpips = 0

    for iteration in range(len(viewpoint)):
        viewpoint_cam = viewpoint[iteration]
        frame_id = viewpoint_cam.uid

        # deform gaussians
        codedict['expr'] = viewpoint_cam.exp_param
        codedict['eyes_pose'] = viewpoint_cam.eyes_pose
        codedict['eyelids'] = viewpoint_cam.eyelids
        codedict['jaw_pose'] = viewpoint_cam.jaw_pose
        verts_final, rot_delta, scale_coef = DeformModel.decode(codedict)
        gaussians.update_xyz_rot_scale(verts_final[0], rot_delta[0], scale_coef[0])

        # Render
        render_pkg = render(viewpoint_cam, gaussians, ppt, background)
        image= render_pkg["render"]
        image = image.clamp(0, 1)

        gt_image = viewpoint_cam.original_image
        _rmse = img_mse(image[None, ...], gt_image[None, ...], mask=None, error_type='rmse', use_mask=False)
        _ssim = img_ssim(image[None, ...], gt_image[None, ...])
        _psnr = img_psnr(image[None, ...], gt_image[None, ...], rmse=_rmse)
        _lpips = perceptual(image[None, ...], gt_image[None, ...], mask=None, use_mask=False)
        ssim += _ssim.item()
        psnr += _psnr.item()
        lpips += _lpips.item()

        save_image = np.zeros((args.image_res, args.image_res*2, 3))
        gt_image_np = (gt_image*255.).permute(1,2,0).detach().cpu().numpy()
        image_np = (image*255.).permute(1,2,0).detach().cpu().numpy()

        save_image[:, :args.image_res, :] = gt_image_np
        save_image[:, args.image_res:, :] = image_np
        save_image = save_image.astype(np.uint8)
        save_image = save_image[:,:,[2,1,0]]

        out.write(save_image)
    out.release()

    ssim /= len(viewpoint)
    psnr /= len(viewpoint)
    lpips /= len(viewpoint)
    with open(os.path.join(logdir, "metrics.txt"), 'w') as f:
        f.writelines(f"psnr: {psnr}\n")
        f.writelines(f"ssim: {ssim}\n")
        f.writelines(f"lpips: {lpips}\n")
        f.close()