evaluate.py

#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import os
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
import cv2
import json
import numpy as np
import torch
from argparse import ArgumentParser
from os import makedirs
from shutil import copy
from tqdm import tqdm

from arguments import ModelParams, get_combined_args
from gaussian_renderer import render
from lpipsPyTorch import lpips
from scene.dataset_readers import get_dataset_reader
from scene.gaussian_model import GaussianModel
from utils.general_utils import safe_state
from utils.image_utils import psnr, read_images
from utils.loss_utils import ssim


class Dummy:
    pass


def make_pipe():
    pipe = Dummy()
    pipe.debug = False
    pipe.antialiasing = False
    pipe.compute_cov3D_python = False
    pipe.convert_SHs_python = False
    return pipe


def save_rendering(rendering, out_path, ext):
    rendering = rendering.detach().permute(1, 2, 0).cpu().numpy()
    if ext != ".exr":
        rendering = (rendering.clip(0, 1) * 255.0).astype(np.uint8)
    cv2.imwrite(out_path, rendering[:, :, ::-1])  # render is RGB; cv2 writes BGR


def render_channel(views, gaussians, pipe, background, gt_source_dir, out_dir, render_ext, gt_ext, colors_activation):
    render_path = os.path.join(out_dir, "renders")
    makedirs(render_path, exist_ok=True)
    gt_path = None
    if gt_source_dir is not None:
        gt_path = os.path.join(out_dir, "gt")
        makedirs(gt_path, exist_ok=True)

    for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
        rendering = render(view, gaussians, pipe, background, colors_activation=colors_activation)["render"]
        save_rendering(rendering, os.path.join(render_path, "{:05d}{}".format(idx, render_ext)), render_ext)
        if gt_path is not None:
            copy(os.path.join(gt_source_dir, view.image_name + gt_ext),
                 os.path.join(gt_path, "{:05d}{}".format(idx, gt_ext)))


def render_full_channel(views, albedo_g, shading_g, residual_g, pipe, background, gt_source_dir, out_dir, ext):
    render_path = os.path.join(out_dir, "renders")
    gt_path = os.path.join(out_dir, "gt")
    makedirs(render_path, exist_ok=True)
    makedirs(gt_path, exist_ok=True)

    for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
        a = render(view, albedo_g, pipe, background, colors_activation=True)["render"]
        s = render(view, shading_g, pipe, background, colors_activation=False)["render"]
        r = render(view, residual_g, pipe, background, colors_activation=False, enable_shs=True)["render"]
        rendering = a * s + r
        save_rendering(rendering, os.path.join(render_path, "{:05d}{}".format(idx, ext)), ext)
        copy(os.path.join(gt_source_dir, view.image_name + ext),
             os.path.join(gt_path, "{:05d}{}".format(idx, ext)))


def render_test_set(dataset, iteration):
    bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]
    background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
    pipe = make_pipe()
    input_images_ext = os.listdir(os.path.join(dataset.source_path, dataset.images))[0][-4:]
    convert_images_to_linear = input_images_ext != ".exr" 

    dataset_reader = get_dataset_reader(
        path=dataset.source_path,
        images_dir=dataset.images,
        albedo_dir=dataset.albedo,
        depths_dir=dataset.depths,
        eval=dataset.eval,
        llffhold=dataset.llffhold,
        convert_images_to_linear=convert_images_to_linear,
    )
    views = dataset_reader.read().test_cameras

    ply_dir = os.path.join(dataset.model_path, "point_cloud", f"iteration_{iteration}")
    out_root = os.path.join(dataset.model_path, "test", f"ours_{iteration}")
    available = set(os.listdir(ply_dir))

    for channel in ("albedo", "shading"):
        if f"{channel}.ply" not in available:
            continue
        print(f"\nRendering channel: {channel}")
        gaussians = GaussianModel(0)
        gaussians.use_color_activation = (channel == "albedo")
        gaussians.load_ply(os.path.join(ply_dir, f"{channel}.ply"))

        channel_src_dir = os.path.join(dataset.source_path, channel)
        has_gt = os.path.isdir(channel_src_dir)
        gt_ext = os.listdir(channel_src_dir)[0][-4:] if has_gt else None
        # shading renders are always PNG; other channels match GT when available, else default to PNG
        render_ext = ".png" if (channel == "shading" or not has_gt) else gt_ext

        render_channel(
            views, gaussians, pipe, background,
            gt_source_dir=channel_src_dir if has_gt else None,
            out_dir=os.path.join(out_root, channel),
            render_ext=render_ext,
            gt_ext=gt_ext,
            colors_activation=(channel == "albedo"),
        )

    if {"albedo.ply", "shading.ply", "residual.safetensors"}.issubset(available):
        print("\nRendering channel: full")
        albedo_g = GaussianModel(0)
        albedo_g.use_color_activation = True
        albedo_g.load_ply(os.path.join(ply_dir, "albedo.ply"))

        shading_g = GaussianModel(0)
        shading_g.load_ply(os.path.join(ply_dir, "shading.ply"))

        residual_g = GaussianModel(sh_degree=3, grayscale_dc=True)
        residual_g.load_safetensors(os.path.join(ply_dir, "residual.safetensors"))

        render_full_channel(
            views, albedo_g, shading_g, residual_g, pipe, background,
            gt_source_dir=os.path.join(dataset.source_path, dataset.images),
            out_dir=os.path.join(out_root, "full"),
            ext=input_images_ext,
        )


def compute_metrics(model_path, gamma_correct=False, srgb_input=False):
    full_dict, per_view_dict = {}, {}
    test_dir = os.path.join(model_path, "test")
    print("\nScene:", model_path)

    for method in os.listdir(test_dir):
        method_dir = os.path.join(test_dir, method)
        print("Method:", method)
        for channel in os.listdir(method_dir):
            gt_dir = os.path.join(method_dir, channel, "gt")
            if not os.path.isdir(gt_dir):
                print(f"\n Channel: {channel} — no GT, skipping")
                continue
            print("\n Channel:", channel)
            key = f"{method}_{channel}"
            full_dict[key] = {}
            per_view_dict[key] = {}

            renders, gts, names = read_images(
                os.path.join(method_dir, channel, "renders"),
                gt_dir,
                gamma_correct=gamma_correct,
                convert_gt_to_linear=srgb_input,
            )

            ssims, psnrs, lpipss = [], [], []
            for r, g in tqdm(zip(renders, gts), total=len(renders), desc="Metric evaluation progress"):
                ssims.append(ssim(r, g))
                psnrs.append(psnr(r, g))
                lpipss.append(lpips(r, g, net_type="vgg"))

            print("  SSIM : {:>12.7f}".format(torch.tensor(ssims).mean()))
            print("  PSNR : {:>12.7f}".format(torch.tensor(psnrs).mean()))
            print("  LPIPS: {:>12.7f}".format(torch.tensor(lpipss).mean()))
            print("")

            full_dict[key].update({
                "SSIM": torch.tensor(ssims).mean().item(),
                "PSNR": torch.tensor(psnrs).mean().item(),
                "LPIPS": torch.tensor(lpipss).mean().item(),
            })
            per_view_dict[key].update({
                "SSIM": {n: v for v, n in zip(torch.tensor(ssims).tolist(), names)},
                "PSNR": {n: v for v, n in zip(torch.tensor(psnrs).tolist(), names)},
                "LPIPS": {n: v for v, n in zip(torch.tensor(lpipss).tolist(), names)},
            })

    with open(os.path.join(model_path, "results.json"), "w") as fp:
        json.dump(full_dict, fp, indent=True)
    with open(os.path.join(model_path, "per_view.json"), "w") as fp:
        json.dump(per_view_dict, fp, indent=True)


if __name__ == "__main__":
    parser = ArgumentParser(description="Render the test set and compute metrics on it")
    model = ModelParams(parser, sentinel=True)
    parser.add_argument("--iteration", default=-1, type=int)
    parser.add_argument("--gamma_correct", action="store_true")
    parser.add_argument("--skip_render", action="store_true", help="Reuse existing renders and only compute metrics")
    parser.add_argument("--skip_metrics", action="store_true", help="Only render, do not compute metrics")
    parser.add_argument("--quiet", action="store_true")
    args = get_combined_args(parser)
    print("Evaluating", args.model_path)

    safe_state(args.quiet)
    device = torch.device("cuda:0")
    torch.cuda.set_device(device)

    dataset = model.extract(args)

    if not args.skip_render:
        with torch.no_grad():
            render_test_set(dataset, args.iteration)

    if not args.skip_metrics:
        compute_metrics(dataset.model_path, args.gamma_correct, dataset.srgb_input)