jmunkberg
commented
1 year ago Hello @JiuTongBro ,
This is the script we used for the relighting comparison. Let me know if this works for you.
import os
import imageio
import numpy as np
import torch
def get_errors(file_o, it, x, ref):
mse = torch.nn.functional.mse_loss(x, ref, size_average=None, reduce=None, reduction='mean')
psnr = -10. / np.log(10.) * torch.log(mse)
file_o.write("%d, %1.8f, %1.8f\n" % (it, mse, psnr))
return torch.stack((mse, psnr), dim=0)
def get_ref(ds, frame):
ref_file = os.path.join("references_test", ds, "val_%03d" % frame, "albedo.png")
ref = torch.tensor(imageio.imread(ref_file).astype(np.float32) / 255.0, dtype=torch.float32, device="cuda")
return ref[..., 0:3], ref[..., 3:4]
def get_our(ds, frame):
our_file = os.path.join("our_test", ds, "validate", "val_%06d_kd.png" % (frame))
mask_file = os.path.join("our_test", ds, "imgs_blender", "city_r_%03d.png" % (frame))
our = torch.tensor(imageio.imread(our_file).astype(np.float32) / 255.0, dtype=torch.float32, device="cuda")
mask = torch.tensor(imageio.imread(mask_file).astype(np.float32) / 255.0, dtype=torch.float32, device="cuda")
return our[..., :3], mask[..., 3:4]
def get_nerfactor(ds, frame):
nerfactor_file = os.path.join("nerfactor_test", ds, "batch%09d" % f, "pred_albedo.png")
nerfactor_mask = os.path.join("nerfactor_test", ds, "batch%09d" % f, "gt_alpha.png")
nerfactor = torch.tensor(imageio.imread(nerfactor_file).astype(np.float32) / 255.0, dtype=torch.float32, device="cuda")
nerfactor_mask = torch.tensor(imageio.imread(nerfactor_mask).astype(np.float32) / 255.0, dtype=torch.float32, device="cuda")[..., None]
return nerfactor, nerfactor_mask
datasets = ['drums_3072', 'ficus_2188', 'hotdog_2163', 'lego_3072']
env_probes = ['city', 'courtyard', 'forest', 'interior', 'night', 'studio', 'sunrise', 'sunset']
for ds in datasets:
# Load all images and compute an average scale factor
scale_our = torch.tensor([0, 0, 0], dtype=torch.float32, device="cuda")
scale_nerfactor = torch.tensor([0, 0, 0], dtype=torch.float32, device="cuda")
ref, our, nerfactor = [], [], []
for f in range(8):
ref += [get_ref(ds, f)]
our += [get_our(ds, f)]
nerfactor += [get_nerfactor(ds, f)]
for i in range(3):
scale_our[i] += (torch.sum(ref[f][0][..., i] * ref[f][1][..., 0]) / torch.sum(ref[f][1][..., 0])) / (torch.sum(our[f][0][..., i] * our[f][1][..., 0]) / torch.sum(our[f][1][..., 0]))
scale_nerfactor[i] += (torch.sum(ref[f][0][..., i] * ref[f][1][..., 0]) / torch.sum(ref[f][1][..., 0])) / (torch.sum(nerfactor[f][0][..., i] * nerfactor[f][1][..., 0]) / torch.sum(nerfactor[f][1][..., 0]))
scale_our /= 8
scale_nerfactor /= 8
#print("Scale_our :", scale_our.detach().cpu().numpy())
#print("Scale_nerfactor:", scale_nerfactor.detach().cpu().numpy())
print("Dataset: ", ds)
print(" MSE, PSNR")
os.makedirs("metrics", exist_ok=True)
fout_our = open('metrics/metrics_kd_our_%s.txt' % ds, 'w')
fout_our.write('ID, MSE, PSNR\n')
fout_nerfactor = open('metrics/metrics_kd_nerfactor_%s.txt' % ds, 'w')
fout_nerfactor.write('ID, MSE, PSNR\n')
errors_our = torch.tensor([0, 0], dtype=torch.float32, device="cuda")
errors_nerfactor = torch.tensor([0, 0], dtype=torch.float32, device="cuda")
itr = 0
for f in range(8):
# Scale our / nerfactor images by global scale factor
_ref = ref[f][0]
_our = our[f][0] * scale_our[None, None, :]
_nerfactor = nerfactor[f][0] * scale_nerfactor[None, None, :]
# Make sure all images comform by adding white bckground based on coverage masks
_ref = torch.lerp(torch.ones_like(_ref), _ref, ref[f][1])
_our = torch.lerp(torch.ones_like(_our), _our, our[f][1])
_nerfactor = torch.lerp(torch.ones_like(_nerfactor), _nerfactor, nerfactor[f][1])
# Compute image metrics
errors_our = errors_our + get_errors(fout_our, itr, _our, _ref)
errors_nerfactor = errors_nerfactor + get_errors(fout_nerfactor, itr, _nerfactor, _ref)
itr += 1
errors_our /= itr
errors_nerfactor /= itr
fout_our.write("AVERAGES: %1.4f, %2.3f\n" % (errors_our[0].item(), errors_our[1].item()))
fout_our.close()
fout_nerfactor.write("AVERAGES: %1.4f, %2.3f\n" % (errors_nerfactor[0].item(), errors_nerfactor[1].item()))
fout_nerfactor.close()
print("nerfactor: %1.8f, %2.3f" % (errors_nerfactor[0].item(), errors_nerfactor[1].item()))
print(" our: %1.8f, %2.3f" % (errors_our[0].item(), errors_our[1].item()))
JiuTongBro
Hi. Thanks for your impressive work firstly.
Recently I am trying to reimplement the relighting results in your paper. However, although my predicted albedo is almost the same as in your paper, as shown below, I can't get the corresponding metric results as yout paper provided. I followed your relighting code in https://github.com/NVlabs/nvdiffrecmc/blob/main/blender/blender.py, and I suppose the problem lies in the RGB-channel scale processing.
Since your prediceted kd is in SRGB space, but the ground truth albedo generated by Blender is seems to be in linear space. I wonder, did you scale your albedo texture in SRGB space or in linear space. And, in relighting shall I transfer the texture_kd image to linear space before I put it into Bledenr use your code？Moreover, in https://github.com/NVlabs/nvdiffrecmc/issues/14 you mentioned we can "compare raw PSNR to the reference, normalize by average intensity, or normalize color channels individually", does that mean we need to scale the relighted image, but not like in NeRFactor where we scale the albedo of each view, to get the paper results? Thans.