mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-12 19:04:03 +08:00
163 lines
6.4 KiB
Python
Executable File
163 lines
6.4 KiB
Python
Executable File
from ezexr import imread
|
|
import os
|
|
import numpy as np
|
|
from tqdm import tqdm
|
|
from math import log10
|
|
import json
|
|
import cv2
|
|
import imageio
|
|
import argparse
|
|
|
|
|
|
def wrmse(gt, est, mask):
|
|
if mask is None:
|
|
gt = gt.flatten()
|
|
est = est.flatten()
|
|
else:
|
|
gt = gt[mask].flatten()
|
|
est = est[mask].flatten()
|
|
error = np.sqrt(np.mean(np.power(gt - est, 2)))
|
|
|
|
return error
|
|
|
|
def si_wrmse(gt, est, mask):
|
|
if mask is None:
|
|
gt_c = gt.flatten()
|
|
est_c = est.flatten()
|
|
else:
|
|
gt_c = gt[mask].flatten()
|
|
est_c = est[mask].flatten()
|
|
alpha = (np.dot(np.transpose(gt_c), est_c)) / (np.dot(np.transpose(est_c), est_c))
|
|
error = wrmse(gt, est * alpha, mask)
|
|
|
|
return error
|
|
|
|
def tonemap(img, gamma=2.4):
|
|
"""Apply gamma, then clip between 0 and 1, finally convert to uint8 [0,255]"""
|
|
return (np.clip(np.power(img,1/gamma), 0.0, 1.0)*255).astype('uint8')
|
|
|
|
|
|
def angular_error(gt_render, pred_render, mask=None):
|
|
# The error need to be computed with the normalized rgb image.
|
|
# Normalized RGB is r = R / (R+G+B), g = G / (R+G+B), b = B / (R+G+B)
|
|
# The angular distance is the distance between pixel 1 and pixel 2.
|
|
# It's computed with cos^-1(p1·p2 / ||p1||*||p2||)
|
|
gt_norm = np.empty((gt_render.shape))
|
|
pred_norm = np.empty(pred_render.shape)
|
|
|
|
for i in range(3):
|
|
gt_norm[:,:,i] = gt_render[:,:,i] / np.sum(gt_render, axis=2, keepdims=True)[:,:,0]
|
|
pred_norm[:,:,i] = pred_render[:,:,i] / (np.sum(pred_render, axis=2, keepdims=True)[:,:,0] + 1e-8)
|
|
|
|
angular_error_arr = np.arccos( np.sum(gt_norm*pred_norm, axis=2, keepdims=True)[:,:,0] /
|
|
((np.sqrt(np.sum(gt_norm*gt_norm, axis=2, keepdims=True)[:,:,0])*np.sqrt(np.sum(pred_norm*pred_norm, axis=2, keepdims=True)[:,:,0]))) )
|
|
|
|
if mask is not None:
|
|
angular_error_arr = angular_error_arr[mask[:,:,0]]
|
|
else:
|
|
angular_error_arr = angular_error_arr.flatten()
|
|
angular_error_arr = angular_error_arr[~np.isnan(angular_error_arr)]
|
|
mean = np.mean(angular_error_arr)
|
|
# convert to degree
|
|
mean = mean * 180 / np.pi
|
|
return mean
|
|
|
|
|
|
def psnr(original, compressed):
|
|
original_tone_mapped = tonemap(original)
|
|
compressed_tone_mapped = tonemap(compressed)
|
|
|
|
cv2.imwrite('./out.png', cv2.cvtColor(compressed_tone_mapped, cv2.COLOR_RGB2BGR))
|
|
cv2.imwrite('./gt.png', cv2.cvtColor(original_tone_mapped, cv2.COLOR_RGB2BGR))
|
|
|
|
mse = wrmse(original_tone_mapped, compressed_tone_mapped, None)
|
|
if(mse == 0): # MSE is zero means no noise is present in the signal .
|
|
# Therefore PSNR have no importance.
|
|
return 100
|
|
max_pixel = 255 #max(original.max(), compressed.max())
|
|
psnr = 20 * log10(max_pixel / mse)
|
|
return psnr
|
|
|
|
|
|
def remove_nan_and_sort(x):
|
|
x = np.array(x)
|
|
x = x[~np.isnan(x)]
|
|
x = np.sort(x)
|
|
return x
|
|
|
|
def calculate_metrics(results_dataset_roots, gt_dataset_root, output_folder, dataset_names):
|
|
results_dataset_roots = [results_dataset_roots]
|
|
dataset_names = [dataset_names]
|
|
mse = {}
|
|
si_mse = {}
|
|
angular = {}
|
|
psnr_error = {}
|
|
for key in dataset_names:
|
|
mse[key] = []
|
|
si_mse[key] = []
|
|
angular[key] = []
|
|
psnr_error[key] = []
|
|
|
|
gt_dataset_files = sorted([os.path.join(gt_dataset_root, f) for f in os.listdir(gt_dataset_root) if f.endswith('.exr')])
|
|
|
|
# create output folder output_folder, 'metrics'
|
|
os.system('mkdir -p ' + output_folder)
|
|
|
|
for results_dataset_root, dataset_name in zip(results_dataset_roots, dataset_names):
|
|
|
|
for gt_dataset_file in tqdm(gt_dataset_files):
|
|
gt_dataset_img_exr = imread(gt_dataset_file)[:, :, :3].astype(np.float32)
|
|
|
|
result_dataset_file = gt_dataset_file.replace(gt_dataset_root, results_dataset_root)
|
|
result_dataset_img_exr = imread(result_dataset_file)[:, :, :3].astype(np.float32)
|
|
|
|
mask = np.float32(imageio.imread('lantern_scripts/render_mask.png')) / 255.0
|
|
mask = mask == 1
|
|
mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
|
|
gt_dataset_img_exr_for_PSNR = gt_dataset_img_exr[mask]
|
|
result_dataset_img_exr_for_PSNR = result_dataset_img_exr[mask]
|
|
|
|
mse_result = wrmse(gt_dataset_img_exr, result_dataset_img_exr, mask)
|
|
scale_invariant_mse_result = si_wrmse(gt_dataset_img_exr, result_dataset_img_exr, mask)
|
|
angular_result = angular_error(gt_dataset_img_exr, result_dataset_img_exr, mask)
|
|
psnr_result = psnr(gt_dataset_img_exr_for_PSNR, result_dataset_img_exr_for_PSNR)
|
|
mse[dataset_name].append(mse_result)
|
|
si_mse[dataset_name].append(scale_invariant_mse_result)
|
|
angular[dataset_name].append(angular_result)
|
|
psnr_error[dataset_name].append(psnr_result)
|
|
si_mse[dataset_name].append(scale_invariant_mse_result)
|
|
|
|
mse[dataset_name] = remove_nan_and_sort(mse[dataset_name])
|
|
|
|
mse_ours = np.mean(mse[dataset_name])
|
|
si_mse_ours = np.mean(si_mse[dataset_name])
|
|
angular_ours = np.mean(angular[dataset_name])
|
|
psnr_ours = np.mean(psnr_error[dataset_name])
|
|
|
|
# print(dataset_name + ', si_RMSE: {}'.format(si_mse_ours))
|
|
# print(dataset_name + ', RMSE: {}'.format(mse_ours))
|
|
# print(dataset_name + ', angular: {}'.format(angular_ours))
|
|
# print(dataset_name + ', PSNR: {}'.format(psnr_ours))
|
|
|
|
output_dict = {}
|
|
output_dict['model_name'] = dataset_name
|
|
output_dict['si_RMSE'] = str(si_mse_ours)
|
|
output_dict['RMSE'] = str(mse_ours)
|
|
output_dict['angular'] = str(angular_ours)
|
|
output_dict['PSNR'] = str(psnr_ours)
|
|
|
|
output_addr = output_folder + 'render_results.json'
|
|
|
|
with open(output_addr, 'w') as f:
|
|
json.dump(output_dict, f, indent=4)
|
|
|
|
if __name__ == '__main__':
|
|
argparser = argparse.ArgumentParser()
|
|
argparser.add_argument("--results_dataset_roots", type=str, default="/mnt/data/coffee_room2/metrics/renders/unaligned_lin_renders/")
|
|
argparser.add_argument("--gt_dataset_root", type=str, default="/mnt/data/coffee_room2/GT/GT_exr_renders/")
|
|
argparser.add_argument("--output_folder", type=str, default="/mnt/data/coffee_room2/metrics/results_table/")
|
|
argparser.add_argument("--dataset_names", type=str, default="/mnt/data/coffee_room2/metrics/results_table/")
|
|
|
|
args = argparser.parse_args()
|
|
|
|
calculate_metrics(args.results_dataset_roots, args.gt_dataset_root, args.output_folder, args.dataset_names) |