|
- import configargparse
- import glob
- import torch
- import os.path
-
- import numpy as np
- # from utils.utils import rgb_lpips, rgb_ssim
- from skimage import metrics
- from PIL import Image
- import lpips
-
-
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-
- def evaluate_replica_tensorf(gt_dir_path, pred_dir_path, scene_id, mode='test'):
- scene_id_simple = scene_id.replace('_', '')
-
- gt_rgbs_dir_path = os.path.join(gt_dir_path, scene_id, 'rgb')
- # pred_rgbs_dir_path = os.path.join(pred_dir_path, scene_id, 'rawfeas',
- # f'{scene_id_simple}_sem_MLPFea_rawfeas_tinyvoxel', # room_1, room_2
- # # f'{scene_id_simple}_sem_MLPFea_rawfeas_pca96_512voxel', # office_0, office_2, office_4, room_0
- # # f'{scene_id_simple}_sem_MLPFea_rawfeas_sdf_512voxel', # office_3
- # 'imgs_vis')
- pred_rgbs_dir_path = os.path.join(pred_dir_path, f'{scene_id_simple}_sem_MLPFea_no_rawfeas_tinyvoxel', 'imgs_vis') # ablation
-
- img_num = len(glob.glob(os.path.join(pred_rgbs_dir_path, '019999_*.png')))
-
- psnr_total = []
- ssim_total = []
- lpips_total = []
- lpips_vgg = lpips.LPIPS(net='vgg').to(device) # version 0.1
-
- for i in range(img_num):
- pred_rgb_path = os.path.join(pred_rgbs_dir_path, f'019999_{i:03d}.png')
- # pred_rgb_path = os.path.join(pred_rgbs_dir_path, f'029999_{i:03d}.png')
- # pred_rgb_path = os.path.join(pred_rgbs_dir_path, f'039999_{i:03d}.png')
- pred_rgb = Image.open(pred_rgb_path)
- pred_rgb = np.array(pred_rgb)
- h, w, _ = pred_rgb.shape
-
- if mode == 'train':
- gt_rgb_path = os.path.join(gt_rgbs_dir_path, f'rgb_{5*i}.png')
- else:
- gt_rgb_path = os.path.join(gt_rgbs_dir_path, f'rgb_{5*i+2}.png')
- gt_rgb = Image.open(gt_rgb_path)
- gt_rgb = np.array(gt_rgb)
-
- pred_rgb_norm, gt_rgb_norm = (pred_rgb / 255.).astype(np.float32), (gt_rgb / 255.).astype(np.float32)
- '''psnr'''
- mse = np.mean((pred_rgb_norm - gt_rgb_norm) ** 2)
- psnr = -10. * np.log(mse) / np.log(10.)
- psnr_total.append(psnr)
-
- '''ssim'''
- ssim = metrics.structural_similarity(pred_rgb_norm, gt_rgb_norm, multichannel=True, data_range=1)
- ssim_total.append(ssim)
-
- '''lpips'''
- pred_rgb_tensor = torch.tensor(pred_rgb_norm).permute(2, 0, 1).unsqueeze(0)
- gt_rgb_tensor = torch.tensor(gt_rgb_norm).permute(2, 0, 1).unsqueeze(0)
- lpips_value = lpips_vgg.forward(pred_rgb_tensor.to(device), gt_rgb_tensor.to(device)).cpu().item()
- lpips_total.append(lpips_value)
-
- print(f'Image {i} has PSNR {psnr}, SSIM {ssim}, LPIPS {lpips_value}')
-
- psnr_avg = np.mean(np.array(psnr_total))
- ssim_avg = np.mean(np.array(ssim_total))
- lpips_avg = np.mean(np.array(lpips_total))
- print(f'TensoRF: scene {scene_id_simple} average PSNR is {psnr_avg}, SSIM is {ssim_avg}, average LPIPS is {lpips_avg}')
-
-
- def evaluate_scannet_tensorf(gt_dir_path, pred_dir_path, scene_id, mode='test', resize=True):
- scene_id_simple = scene_id.replace('_', '')
- img_indices = np.loadtxt(os.path.join(gt_dir_path, f'{scene_id_simple}_00', mode+'_split.txt')).astype(np.int16)
-
- gt_rgbs_dir_path = os.path.join(gt_dir_path, f'{scene_id_simple}_00', mode, f'{mode}_images')
- pred_rgbs_dir_path = os.path.join(pred_dir_path, scene_id, 'rawfeas', f'{scene_id_simple}_rawfeas_tinyvoxel',
- 'imgs_vis')
-
- psnr_total = []
- ssim_total = []
- lpips_total = []
- lpips_vgg = lpips.LPIPS(net='vgg').to(device)
-
- for i in range(len(img_indices)):
- # pred_rgb_path = os.path.join(pred_rgbs_dir_path, f'019999_{i:03d}.png')
- # pred_rgb_path = os.path.join(pred_rgbs_dir_path, f'029999_{i:03d}.png')
- pred_rgb_path = os.path.join(pred_rgbs_dir_path, f'039999_{i:03d}.png')
- pred_rgb = Image.open(pred_rgb_path)
- pred_rgb = np.array(pred_rgb)
- h, w, _ = pred_rgb.shape
-
- gt_rgb_path = os.path.join(gt_rgbs_dir_path, f'{img_indices[i]}.jpg')
- gt_rgb = Image.open(gt_rgb_path)
- if resize:
- gt_rgb = gt_rgb.resize((w, h), Image.LANCZOS)
- gt_rgb = np.array(gt_rgb)
-
- pred_rgb_norm, gt_rgb_norm = (pred_rgb / 255.).astype(np.float32), (gt_rgb / 255.).astype(np.float32)
- '''psnr'''
- mse = np.mean((pred_rgb_norm - gt_rgb_norm) ** 2)
- psnr = -10. * np.log(mse) / np.log(10.)
- psnr_total.append(psnr)
-
- '''ssim'''
- ssim = metrics.structural_similarity(pred_rgb_norm, gt_rgb_norm, multichannel=True, data_range=1)
- ssim_total.append(ssim)
-
- '''lpips'''
- pred_rgb_tensor = torch.tensor(pred_rgb_norm).permute(2, 0, 1).unsqueeze(0)
- gt_rgb_tensor = torch.tensor(gt_rgb_norm).permute(2, 0, 1).unsqueeze(0)
- lpips_value = lpips_vgg.forward(pred_rgb_tensor.to(device), gt_rgb_tensor.to(device)).cpu().item()
- lpips_total.append(lpips_value)
-
- print(f'Image {i} has PSNR {psnr}, SSIM {ssim}, LPIPS {lpips_value}')
-
- psnr_avg = np.mean(np.array(psnr_total))
- ssim_avg = np.mean(np.array(ssim_total))
- lpips_avg = np.mean(np.array(lpips_total))
- print(f'TensoRF: scene {scene_id_simple} average PSNR is {psnr_avg}, SSIM is {ssim_avg}, average LPIPS is {lpips_avg}')
-
-
- def evaluate_replica_semnerf(gt_dir_path, pred_dir_path, scene_id, mode='test'):
- scene_id_simple = scene_id.replace('_', '')
-
- gt_rgbs_dir_path = os.path.join(gt_dir_path, scene_id, 'rgb')
- pred_rgbs_dir_path = os.path.join(pred_dir_path, 'replica', f'{scene_id_simple}_P1.00', 'step_000000')
-
- img_num = len(glob.glob(os.path.join(pred_rgbs_dir_path, 'rgb_*.png')))
-
- psnr_total = []
- ssim_total = []
- lpips_total = []
- lpips_vgg = lpips.LPIPS(net='vgg').to(device)
-
- for i in range(img_num):
- img_idx = 3 * i
- pred_rgb_path = os.path.join(pred_rgbs_dir_path, f'rgb_{img_idx:03d}.png')
- pred_rgb = Image.open(pred_rgb_path)
- pred_rgb = np.array(pred_rgb)
- h, w, _ = pred_rgb.shape
-
- if mode == 'train':
- gt_rgb_path = os.path.join(gt_rgbs_dir_path, f'rgb_{5*img_idx}.png')
- else:
- gt_rgb_path = os.path.join(gt_rgbs_dir_path, f'rgb_{5*img_idx+2}.png')
- gt_rgb = Image.open(gt_rgb_path)
- gt_rgb = np.array(gt_rgb)
-
- pred_rgb_norm, gt_rgb_norm = (pred_rgb / 255.).astype(np.float32), (gt_rgb / 255.).astype(np.float32)
- '''psnr'''
- mse = np.mean((pred_rgb_norm - gt_rgb_norm) ** 2)
- psnr = -10. * np.log(mse) / np.log(10.)
- psnr_total.append(psnr)
-
- '''ssim'''
- ssim = metrics.structural_similarity(pred_rgb_norm, gt_rgb_norm, multichannel=True, data_range=1)
- ssim_total.append(ssim)
-
- '''lpips'''
- pred_rgb_tensor = torch.tensor(pred_rgb_norm).permute(2, 0, 1).unsqueeze(0)
- gt_rgb_tensor = torch.tensor(gt_rgb_norm).permute(2, 0, 1).unsqueeze(0)
- lpips_value = lpips_vgg.forward(pred_rgb_tensor.to(device), gt_rgb_tensor.to(device)).cpu().item()
- lpips_total.append(lpips_value)
-
- print(f'Image {img_idx} has PSNR {psnr}, SSIM {ssim}, LPIPS {lpips_value}')
-
- psnr_avg = np.mean(np.array(psnr_total))
- ssim_avg = np.mean(np.array(ssim_total))
- lpips_avg = np.mean(np.array(lpips_total))
- print(f'Sem-NeRF: scene {scene_id_simple} average PSNR is {psnr_avg}, SSIM is {ssim_avg}, average LPIPS is {lpips_avg}')
-
-
- def evaluate_scannet_semnerf(gt_dir_path, pred_dir_path, scene_id, mode='test', resize=True):
- scene_id_simple = scene_id.replace('_', '')
- img_indices = np.loadtxt(os.path.join(gt_dir_path, f'{scene_id_simple}_00', mode+'_split.txt')).astype(np.int16)
-
- gt_rgbs_dir_path = os.path.join(gt_dir_path, f'{scene_id_simple}_00', mode, f'{mode}_images')
- pred_rgbs_dir_path = os.path.join(pred_dir_path, 'scannet', f'{scene_id}_P1.00', 'step_000000')
-
- psnr_total = []
- ssim_total = []
- lpips_total = []
- lpips_vgg = lpips.LPIPS(net='vgg').to(device)
-
- for i in range(len(img_indices)):
- pred_rgb_path = os.path.join(pred_rgbs_dir_path, f'rgb_{i:03d}.png')
- pred_rgb = Image.open(pred_rgb_path)
- pred_rgb = np.array(pred_rgb)
- h, w, _ = pred_rgb.shape
-
- gt_rgb_path = os.path.join(gt_rgbs_dir_path, f'{img_indices[i]}.jpg')
- gt_rgb = Image.open(gt_rgb_path)
- if resize:
- gt_rgb = gt_rgb.resize((w, h), Image.LANCZOS)
- gt_rgb = np.array(gt_rgb)
-
- pred_rgb_norm, gt_rgb_norm = (pred_rgb / 255.).astype(np.float32), (gt_rgb / 255.).astype(np.float32)
- '''psnr'''
- mse = np.mean((pred_rgb_norm - gt_rgb_norm) ** 2)
- psnr = -10. * np.log(mse) / np.log(10.)
- psnr_total.append(psnr)
-
- '''ssim'''
- ssim = metrics.structural_similarity(pred_rgb_norm, gt_rgb_norm, multichannel=True, data_range=1)
- ssim_total.append(ssim)
-
- '''lpips'''
- pred_rgb_tensor = torch.tensor(pred_rgb_norm).permute(2, 0, 1).unsqueeze(0)
- gt_rgb_tensor = torch.tensor(gt_rgb_norm).permute(2, 0, 1).unsqueeze(0)
- lpips_value = lpips_vgg.forward(pred_rgb_tensor.to(device), gt_rgb_tensor.to(device)).cpu().item()
- lpips_total.append(lpips_value)
-
- print(f'Image {i} has PSNR {psnr}, SSIM {ssim}, LPIPS {lpips_value}')
-
- psnr_avg = np.mean(np.array(psnr_total))
- ssim_avg = np.mean(np.array(ssim_total))
- lpips_avg = np.mean(np.array(lpips_total))
- print(f'Sem-NeRF: scene {scene_id_simple} average PSNR is {psnr_avg}, SSIM is {ssim_avg}, average LPIPS is {lpips_avg}')
-
-
- def evaluate_replica_dmnerf(gt_dir_path, pred_dir_path, scene_id, mode='test'):
- scene_id_simple = scene_id.replace('_', '')
-
- gt_rgbs_dir_path = os.path.join(gt_dir_path, scene_id, 'rgb')
- pred_rgbs_dir_path = os.path.join(pred_dir_path, 'replica', f'{scene_id_simple}', 'sparse1', 'testset_200000')
-
- img_num = len(glob.glob(os.path.join(pred_rgbs_dir_path, 'instance_*.png')))
-
- psnr_total = []
- ssim_total = []
- lpips_total = []
- lpips_vgg = lpips.LPIPS(net='vgg').to(device)
-
- for i in range(img_num):
- pred_rgb_path = os.path.join(pred_rgbs_dir_path, f'{i:03d}.png')
- pred_rgb = Image.open(pred_rgb_path)
- pred_rgb = np.array(pred_rgb)
- h, w, _ = pred_rgb.shape
-
- if mode == 'train':
- gt_rgb_path = os.path.join(gt_rgbs_dir_path, f'rgb_{5*i}.png')
- else:
- gt_rgb_path = os.path.join(gt_rgbs_dir_path, f'rgb_{5*i+2}.png')
- gt_rgb = Image.open(gt_rgb_path)
- gt_rgb = np.array(gt_rgb)
-
- pred_rgb_norm, gt_rgb_norm = (pred_rgb / 255.).astype(np.float32), (gt_rgb / 255.).astype(np.float32)
- '''psnr'''
- mse = np.mean((pred_rgb_norm - gt_rgb_norm) ** 2)
- psnr = -10. * np.log(mse) / np.log(10.)
- psnr_total.append(psnr)
-
- '''ssim'''
- ssim = metrics.structural_similarity(pred_rgb_norm, gt_rgb_norm, multichannel=True, data_range=1)
- ssim_total.append(ssim)
-
- '''lpips'''
- pred_rgb_tensor = torch.tensor(pred_rgb_norm).permute(2, 0, 1).unsqueeze(0)
- gt_rgb_tensor = torch.tensor(gt_rgb_norm).permute(2, 0, 1).unsqueeze(0)
- lpips_value = lpips_vgg.forward(pred_rgb_tensor.to(device), gt_rgb_tensor.to(device)).cpu().item()
- lpips_total.append(lpips_value)
-
- print(f'Image {i} has PSNR {psnr}, SSIM {ssim}, LPIPS {lpips_value}')
-
- psnr_avg = np.mean(np.array(psnr_total))
- ssim_avg = np.mean(np.array(ssim_total))
- lpips_avg = np.mean(np.array(lpips_total))
- print(f'DM-NeRF: scene {scene_id_simple} average PSNR is {psnr_avg}, SSIM is {ssim_avg}, average LPIPS is {lpips_avg}')
-
-
- def evaluate_scannet_dmnerf(gt_dir_path, pred_dir_path, scene_id, mode='test', resize=True):
- scene_id_simple = scene_id.replace('_', '')
- img_indices = np.loadtxt(os.path.join(gt_dir_path, f'{scene_id_simple}_00', mode+'_split.txt')).astype(np.int16)
-
- gt_rgbs_dir_path = os.path.join(gt_dir_path, f'{scene_id_simple}_00', mode, f'{mode}_images')
- pred_rgbs_dir_path = os.path.join(pred_dir_path, 'scannet', f'{scene_id_simple}_00', 'sparse01',
- 'testset_300000') # 10,38
- # 'render_test_300000') # 12,24,33,88,113,192
-
- psnr_total = []
- ssim_total = []
- lpips_total = []
- lpips_vgg = lpips.LPIPS(net='vgg').to(device)
-
- for i in range(len(img_indices)):
- pred_rgb_path = os.path.join(pred_rgbs_dir_path, f'{i:03d}.png')
- pred_rgb = Image.open(pred_rgb_path)
- pred_rgb = np.array(pred_rgb)
- h, w, _ = pred_rgb.shape
-
- gt_rgb_path = os.path.join(gt_rgbs_dir_path, f'{img_indices[i]}.jpg')
- gt_rgb = Image.open(gt_rgb_path)
- if resize:
- gt_rgb = gt_rgb.resize((w, h), Image.LANCZOS)
- gt_rgb = np.array(gt_rgb)
-
- pred_rgb_norm, gt_rgb_norm = (pred_rgb / 255.).astype(np.float32), (gt_rgb / 255.).astype(np.float32)
- '''psnr'''
- mse = np.mean((pred_rgb_norm - gt_rgb_norm) ** 2)
- psnr = -10. * np.log(mse) / np.log(10.)
- psnr_total.append(psnr)
-
- '''ssim'''
- ssim = metrics.structural_similarity(pred_rgb_norm, gt_rgb_norm, multichannel=True, data_range=1)
- ssim_total.append(ssim)
-
- '''lpips'''
- pred_rgb_tensor = torch.tensor(pred_rgb_norm).permute(2, 0, 1).unsqueeze(0)
- gt_rgb_tensor = torch.tensor(gt_rgb_norm).permute(2, 0, 1).unsqueeze(0)
- lpips_value = lpips_vgg.forward(pred_rgb_tensor.to(device), gt_rgb_tensor.to(device)).cpu().item()
- lpips_total.append(lpips_value)
-
- print(f'Image {i} has PSNR {psnr}, SSIM {ssim}, LPIPS {lpips_value}')
-
- psnr_avg = np.mean(np.array(psnr_total))
- ssim_avg = np.mean(np.array(ssim_total))
- lpips_avg = np.mean(np.array(lpips_total))
- print(f'DM-NeRF: scene {scene_id_simple} average PSNR is {psnr_avg}, SSIM is {ssim_avg}, average LPIPS is {lpips_avg}')
-
-
- if __name__ == '__main__':
- parser = configargparse.ArgumentParser()
- parser.add_argument('--model', type=str, choices=['tensorf', 'semnerf', 'dmnerf'])
- parser.add_argument('--dataset', type=str, choices=['replica', 'scannet'])
- parser.add_argument('--gt_dir_path', type=str)
- parser.add_argument('--pred_dir_path', type=str)
- parser.add_argument('--scene_id', type=str, choices=['office_0', 'office_2', 'office_3', 'office_4',
- 'room_0', 'room_1', 'room_2',
- 'scene_0010', 'scene_0012', 'scene_0024', 'scene_0033',
- 'scene_0038', 'scene_0088', 'scene_0113', 'scene_0192',
- 'all'])
-
- args = parser.parse_args()
-
- if args.model == 'tensorf':
- if args.dataset == 'replica':
- evaluate_replica_tensorf(args.gt_dir_path, args.pred_dir_path, args.scene_id)
- elif args.dataset == 'scannet':
- evaluate_scannet_tensorf(args.gt_dir_path, args.pred_dir_path, args.scene_id)
-
- elif args.model == 'dmnerf':
- if args.dataset == 'replica':
- evaluate_replica_dmnerf(args.gt_dir_path, args.pred_dir_path, args.scene_id)
- elif args.dataset == 'scannet':
- evaluate_scannet_dmnerf(args.gt_dir_path, args.pred_dir_path, args.scene_id)
-
- elif args.model == 'semnerf':
- if args.dataset == 'replica':
- evaluate_replica_semnerf(args.gt_dir_path, args.pred_dir_path, args.scene_id)
- elif args.dataset == 'scannet':
- evaluate_scannet_semnerf(args.gt_dir_path, args.pred_dir_path, args.scene_id)
|