eval_orb_relight.py

import os
os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
import random
import argparse
from utils.base_utils import load_cfg
import torch
import numpy as np
from tqdm import tqdm
import imageio

import cv2
from omegaconf import OmegaConf

from skimage.io import imread, imsave
from lpips import LPIPS
from kornia.losses import ssim_loss
_lpips = None

def rgb_to_srgb(f: np.ndarray):
    # f is loaded from .exr
    # output is NOT clipped to [0, 1]
    assert len(f.shape) == 3, f.shape
    assert f.shape[2] == 3, f.shape
    f = np.where(f > 0.0031308, np.power(np.maximum(f, 0.0031308), 1.0 / 2.4) * 1.055 - 0.055, 12.92 * f)
    return f


def srgb_to_rgb(f: np.ndarray):
    # f is LDR
    assert len(f.shape) == 3, f.shape
    assert f.shape[2] == 3, f.shape
    f = np.where(f <= 0.04045, f / 12.92, np.power((np.maximum(f, 0.04045) + 0.055) / 1.055, 2.4))
    return f

def erode_mask(mask, target_size):
    if mask.ndim == 3:
        mask = mask[...,0]
    if mask.dtype == np.float32:
        mask = (mask*255).clip(0, 255).astype(np.uint8)
    # shrink mask by a small margin to prevent inaccurate mask boundary.
    kernel = np.ones((5, 5), np.uint8)
    mask = cv2.erode(mask, kernel)
    if target_size is not None:
        mask = cv2.resize(mask, (target_size, target_size))
    return (mask > 127).astype(np.float32)

def lpips(inputs: np.ndarray, target: np.ndarray, mask: np.ndarray):
    if LPIPS is None:
        return np.nan
    global _lpips
    if _lpips is None:
        _lpips = LPIPS(net='vgg', verbose=False).cuda()
    inputs = rgb_to_srgb(inputs)
    target = rgb_to_srgb(target)

    mask = erode_mask(mask, None)
    inputs = inputs * mask[:, :, None]
    target = target * mask[:, :, None]

    inputs = torch.tensor(inputs, dtype=torch.float32, device='cuda').permute(2, 0, 1).unsqueeze(0)
    target = torch.tensor(target, dtype=torch.float32, device='cuda').permute(2, 0, 1).unsqueeze(0)
    return _lpips(inputs, target, normalize=True).item()

def mse_to_psnr(mse):
    """Compute PSNR given an MSE (we assume the maximum pixel value is 1)."""
    return -10. / np.log(10.) * np.log(mse)


def calc_PSNR(img_pred, img_gt, mask_gt,
              max_value, tonemapping, scale_invariant,
              divide_mask=True):
    # make sure img_pred, img_gt are linear
    '''
        calculate the PSNR between the predicted image and ground truth image.
        a scale is optimized to get best possible PSNR.
        images are clip by max_value_ratio.
        params:
        img_pred: numpy.ndarray of shape [H, W, 3]. predicted HDR image.
        img_gt: numpy.ndarray of shape [H, W, 3]. ground truth HDR image.
        mask_gt: numpy.ndarray of shape [H, W]. ground truth foreground mask.
        max_value: Float. the maximum value of the ground truth image clipped to.
            This is designed to prevent the result being affected by too bright pixels.
        tonemapping: Bool. Whether the images are tone-mapped before comparion.
        divide_mask: Bool. Whether the mse is divided by the foreground area.
    '''
    if mask_gt.ndim == 3:
        mask_gt = mask_gt[..., 0]
    if mask_gt.dtype == np.float32:
        mask_gt = (mask_gt * 255).clip(0, 255).astype(np.uint8)
    else:
        import ipdb; ipdb.set_trace()
    # shrink mask by a small margin to prevent inaccurate mask boundary.
    kernel = np.ones((5, 5), np.uint8)
    mask_gt = cv2.erode(mask_gt, kernel)
    mask_gt = (mask_gt > 127).astype(np.float32)

    img_pred = img_pred * mask_gt[..., None]
    img_gt = img_gt * mask_gt[..., None]
    img_gt[img_gt < 0] = 0

    if scale_invariant:
        img_pred_pixels = img_pred[np.where(mask_gt > 0.5)]
        img_gt_pixels = img_gt[np.where(mask_gt > 0.5)]
        for c in range(3):
            if (img_pred_pixels[:, c] ** 2).sum() <= 1e-6:
                img_pred_pixels[:, c] = np.ones_like(img_pred_pixels[:, c])
                # import ipdb; ipdb.set_trace()
        scale = (img_gt_pixels * img_pred_pixels).sum(axis=0) / (img_pred_pixels ** 2).sum(axis=0)
        assert scale.shape == (3,), scale.shape
        if (scale < 0).any():
            import ipdb; ipdb.set_trace()
        if (img_pred < 0).any():
            import ipdb; ipdb.set_trace()
        if (img_gt < 0).any():
            import ipdb; ipdb.set_trace()
        img_pred = scale * img_pred
        # if not tonemapping:
    #     imageio.imsave("./rescaled.exr", img_pred)
    #     imageio.imsave("./rescaled_gt.exr", img_gt)

    # clip the prediction and the gt img by the maximum_value
    img_pred = np.clip(img_pred, 0, max_value)
    img_gt = np.clip(img_gt, 0, max_value)

    if tonemapping:
        img_pred = rgb_to_srgb(img_pred)
        img_gt = rgb_to_srgb(img_gt)
        # imageio.imsave("./rescaled.png", (img_pred*255).clip(0,255).astype(np.uint8))
        # imageio.imsave("./rescaled_gt.png", (img_gt*255).clip(0,255).astype(np.uint8))

    if not divide_mask:
        mse = ((img_pred - img_gt) ** 2).mean()
        lb = ((np.ones_like(img_gt) * .5 * mask_gt[:, :, None] - img_gt) ** 2).mean()
    else:
        mse = ((img_pred - img_gt) ** 2).sum() / mask_gt.sum()
        lb = ((np.ones_like(img_gt) * .5 * mask_gt[:, :, None] - img_gt) ** 2).sum() / mask_gt.sum()
    out = mse_to_psnr(mse)
    lb = mse_to_psnr(lb)
    out = max(out, lb)
    return out, img_pred, img_gt


def ssim(inputs: np.ndarray, target: np.ndarray, mask: np.ndarray):
    if ssim_loss is None:
        return np.nan

    mask = erode_mask(mask, None)
    inputs = inputs * mask[:, :, None]
    target = target * mask[:, :, None]

    # image_pred and image_gt: (1, 3, H, W) in range [0, 1]
    inputs = torch.tensor(inputs, dtype=torch.float32, device='cuda').permute(2, 0, 1).unsqueeze(0)
    target = torch.tensor(target, dtype=torch.float32, device='cuda').permute(2, 0, 1).unsqueeze(0)
    dssim_ = ssim_loss(inputs, target, 3).item()  # dissimilarity in [0, 1]
    return 1 - 2 * dssim_  # in [-1, 1]


def load_mask_png(path: str):
    f = imread(path).astype(np.float32)
    f = f / 255.
    assert len(f.shape) == 2, f.shape
    return f

def img_read_rgba(path):
    im = imread(path).astype(np.float32)
    im_rgb = np.array(im)[..., :3] / 255.
    im_alpha = np.array(im)[..., [-1]] / 255.
    im = (im_rgb * im_alpha + (1 - im_alpha))
    return im

def img_read_rgb(path):
    im = imread(path).astype(np.float32)
    im_rgb = np.array(im)[..., :3] / 255.
    return im_rgb

def eval_relight(model_name, relight_dir, gt_dir):
    relight_env_name = gt_dir.split('/')[-1]
    save_dir = f'data/relight/orb/{model_name}/afterScale/{relight_env_name}'
    os.makedirs(save_dir, exist_ok=True)
    avg_PSNR, avg_SSIM, avg_LPIPS = 0, 0, 0
    msg = ''
    num = len(os.listdir(relight_dir))
    for name in tqdm(os.listdir(relight_dir)):
        relit_path = os.path.join(relight_dir, name)
        mask_path = os.path.join(gt_dir, 'test_mask', name)
        gt_path = os.path.join(gt_dir, 'test', name)
        relit_im = img_read_rgba(relit_path)
        gt_im = img_read_rgb(gt_path)
        mask = load_mask_png(mask_path)
 
        cur_psnr, img_pred, img_gt = calc_PSNR(relit_im, gt_im, mask, max_value=1, tonemapping=False, divide_mask=False, scale_invariant=True)
        cur_lpips = lpips(img_pred, img_gt, mask)
        cur_ssim = ssim(img_gt, img_pred, mask)   
        
        msg += f'{name}, psnr: {cur_psnr}, ssim: {cur_ssim}, lpips: {cur_lpips}\n'
        avg_PSNR += cur_psnr
        avg_SSIM += cur_ssim
        avg_LPIPS += cur_lpips
        
        img_pred_save = np.concatenate([img_pred, mask[..., None]], axis=-1)
        img_gt_save = np.concatenate([img_gt, mask[..., None]], axis=-1)
        
        imageio.imsave(f"{save_dir}/{name}", (img_pred_save*255).clip(0,255).astype(np.uint8))
        imageio.imsave(f"{save_dir}/gt_{name}", (img_gt_save*255).clip(0,255).astype(np.uint8))        
    
    avg_PSNR /= num
    avg_SSIM /= num
    avg_LPIPS /= num
    avg_info = f'Avg_psnr: {avg_PSNR}, avg_ssim: {avg_SSIM}, avg_lpips: {avg_LPIPS}\n'
    msg += avg_info
    print(avg_info)
    with open(f'{save_dir}/metrics_record.txt', 'a') as f:
        f.write(msg)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    
    parser.add_argument('--cfg', type=str, default='configs/mat/syn/compressor.yaml')
    parser.add_argument('-e', '--extra_name', type=str, default=None)
    parser.add_argument('--relight_dir', type=str, default='data/relight/orb/teapot_scene006_relighting_teapot_scene001',
                        help='relighted images from blender')
    parser.add_argument('--gt_dir', type=str, default='nerf_data/blender_LDR/teapot_scene001',
                        help='ground truth relighting images from orb LDR datasets')
    # flags = parser.parse_args()

    flags, unknown = parser.parse_known_args()
    cfg = load_cfg(flags.cfg)
    cfg = OmegaConf.create(cfg)
    unknown_cfg = OmegaConf.from_dotlist(unknown)
    cfg = OmegaConf.merge(cfg, unknown_cfg)
    cfg = OmegaConf.to_container(cfg, resolve=True)
    if flags.extra_name is not None:
        cfg['name'] = '_'.join([cfg['name'], flags.extra_name])

    relight_dir = f'data/relight/orb/{cfg["name"]}/noScale/{flags.relight_dir}'
    eval_relight(cfg['name'], relight_dir, flags.gt_dir)