eval_nvs.py

import os

import numpy as np
from argparse import ArgumentParser

import torch
from skimage.io import imread
from tqdm import tqdm
from torchmetrics.functional import structural_similarity_index_measure
import lpips


def compute_psnr_float(img_gt, img_pr):
    img_gt = img_gt.reshape([-1, 3]).astype(np.float32)
    img_pr = img_pr.reshape([-1, 3]).astype(np.float32)
    mse = np.mean((img_gt - img_pr) ** 2, 0)
    mse = np.mean(mse)
    psnr = 10 * np.log10(1 / mse)
    return psnr

def color_map_forward(rgb):
    dim = rgb.shape[-1]
    if dim==3:
        return rgb.astype(np.float32)/255
    else:
        rgb = rgb.astype(np.float32)/255
        rgb, alpha = rgb[:,:,:3], rgb[:,:,3:]
        rgb = rgb * alpha + (1-alpha)
        return rgb

def main():
    parser = ArgumentParser()
    parser.add_argument('--gt',type=str)
    parser.add_argument('--pr',type=str)
    parser.add_argument('--name',type=str)
    parser.add_argument("--num_samples", type=int, default=4)

    args = parser.parse_args()

    num_images = 16
    gt_dir = args.gt
    pr_dir = args.pr

    lpips_fn = lpips.LPIPS(net='vgg').cuda().eval()
    psnrs, ssims, lpipss = [], [], []

    for bi in range(args.num_samples):
        name = args.name + f'-{bi}'
        img_pr_list = imread(os.path.join(pr_dir, f'{bi}.png'))
        img_pr_list = np.array_split(img_pr_list, num_images, 1)
        for k in tqdm(range(num_images)):
            img_gt_int = imread(os.path.join(gt_dir, f'{k:03}.png'))
            img_pr_int = img_pr_list[k]
            img_gt = color_map_forward(img_gt_int)
            img_pr = color_map_forward(img_pr_int)
            psnr = compute_psnr_float(img_gt, img_pr)

            with torch.no_grad():
                img_gt_tensor = torch.from_numpy(img_gt.astype(np.float32)).permute(2,0,1).unsqueeze(0).cuda()
                img_pr_tensor = torch.from_numpy(img_pr.astype(np.float32)).permute(2,0,1).unsqueeze(0).cuda()
                ssim = float(structural_similarity_index_measure(img_pr_tensor, img_gt_tensor).flatten()[0].cpu().numpy())
                gt_img_th, pr_img_th = img_gt_tensor*2-1, img_pr_tensor*2-1
                score = float(lpips_fn(gt_img_th, pr_img_th).flatten()[0].cpu().numpy())

            ssims.append(ssim)
            lpipss.append(score)
            psnrs.append(psnr)

        msg=f'{name:<15}\t{np.mean(psnrs):.5f}\t{np.mean(ssims):.5f}\t{np.mean(lpipss):.5f}'
        print(msg)
        with open(os.path.abspath(os.path.join(pr_dir, '../', 'nvs.log')), 'a') as f:
            f.write(msg+'\n')

if __name__=="__main__":
    main()