Function eval_nerf_bacon

experiments/render_nerf.py:291–342

(scene, config, checkpoint, outdir, res, scale, chunk_size=10000, return_all=False, val_idx=None)

def eval_nerf_bacon(scene, config, checkpoint, outdir, res, scale, chunk_size=10000, return_all=False, val_idx=None):

    os.makedirs(f'./outputs/nerf/{outdir}', exist_ok=True)

    p = configargparse.DefaultConfigFileParser()
    with open(config) as f:
        opt = p.parse(f)

    opt = Options(**opt)
    dataset = load_dataset(opt, res, scale)
    models = load_model(opt, checkpoint)

    for k in models.keys():
        models[k].stop_after = scale

    # render images
    psnrs = []
    ssims = []
    dataset_generator = iter(dataset)
    for idx in range(len(dataset)):

        if val_idx is not None:
            dataset.val_idx = val_idx
            idx = val_idx

        in_dict, meta_dict, gt_dict = next(dataset_generator)

        images, psnr, ssim, elapsed = render_image(opt, models, dataset, chunk_size,
                                                   in_dict, meta_dict, gt_dict,
                                                   scale, return_all=return_all)

        tqdm.write(f'Scale: {scale} | PSNR: {psnr:.02f} dB, SSIM: {ssim:.02f}, Elapsed: {elapsed:.02f} ms')

        if return_all:
            for s in range(4):
                skimage.io.imsave(f'./outputs/nerf/{outdir}/r_{idx}_d{3-s}.png', (images[s]*255).astype(np.uint8))
        else:
            np.save(f'./outputs/nerf/{outdir}/r_{idx}_d{3-scale}.npy', {'psnr': psnr, 'ssim': ssim})
            skimage.io.imsave(f'./outputs/nerf/{outdir}/r_{idx}_d{3-scale}.png', (images*255).astype(np.uint8))

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

        if val_idx is not None:
            break

    if not return_all and val_idx is not None:
        np.save(f'./outputs/nerf/{outdir}/metrics_d{3-scale}.npy', {'psnr': psnrs, 'ssim': ssims,
                                                                    'avg_psnr': np.mean(psnrs),
                                                                    'avg_ssim': np.mean(ssims)})

        print(f'Avg. PSNR: {np.mean(psnrs):.02f}, Avg. SSIM: {np.mean(ssims):.02f}')


if __name__ == '__main__':