render.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 json
import torch
from scene import Scene
import os
from tqdm import tqdm
from os import makedirs
from gaussian_renderer import render
import torchvision
from utils.general_utils import safe_state
from argparse import ArgumentParser
from arguments import ModelParams, PipelineParams, SparseParams, get_combined_args
from gaussian_renderer import GaussianModel
try:
    from diff_gaussian_rasterization import SparseGaussianAdam
    SPARSE_ADAM_AVAILABLE = True
except:
    SPARSE_ADAM_AVAILABLE = False


# =====================================================================================
# ADAPTIVE INFERENCE -- the last stage of the method.
#
# The 3D model is finished; one further pass of the frozen 2D fixer is run on the render
# of each target view, and the two are mixed:
#
#     y     = F(render)                      the fixer's raw output
#     D     = mean((y - render)^2)           how much the fixer wants to change
#     mse   = 6.893 * D^1.251                the fixer's own error, from D
#     beta* = mse / (mse + 0.286 * mse^0.769)
#     out   = beta* * y + (1 - beta*) * render
#
# Both power laws were MEASURED, and there is no tuned constant anywhere in the mix:
# beta* is derived per view from D alone, which is computable at inference time without
# ever seeing the ground truth.  The shape is the MMSE mix of two estimates whose error
# variances are mse and 0.286*mse^0.769 -- the render is trusted where the fixer wants a
# large change (which is where the fixer is least reliable), and the fixer is trusted
# where it only wants a small one.
# =====================================================================================

def beta_star(fixed, raw):
    """Per-view mixing weight.  fixed/raw: CHW float tensors in [0, 1]."""
    d = float(((fixed - raw) ** 2).mean())
    mse = 6.893 * d ** 1.251
    return mse / (mse + 0.286 * mse ** 0.769)


def blend_beta(fixed, raw):
    b = beta_star(fixed, raw)
    return (b * fixed + (1 - b) * raw).clamp(0, 1), b


def render_set(model_path, name, iteration, views, gaussians, pipeline, background, train_test_exp, separate_sh):
    render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders")
    gts_path = os.path.join(model_path, name, "ours_{}".format(iteration), "gt")

    makedirs(render_path, exist_ok=True)
    makedirs(gts_path, exist_ok=True)

    for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
        rendering = render(view, gaussians, pipeline, background, use_trained_exp=train_test_exp, separate_sh=separate_sh)["render"]
        gt = view.original_image[0:3, :, :]

        if args.train_test_exp:
            rendering = rendering[..., rendering.shape[-1] // 2:]
            gt = gt[..., gt.shape[-1] // 2:]

        torchvision.utils.save_image(rendering, os.path.join(render_path, '{0:05d}'.format(idx) + ".png"))
        torchvision.utils.save_image(gt, os.path.join(gts_path, '{0:05d}'.format(idx) + ".png"))

def render_sets(dataset : ModelParams, iteration : int, pipeline : PipelineParams, skip_train : bool, skip_test : bool, separate_sh: bool):
    with torch.no_grad():
        gaussians = GaussianModel(dataset.sh_degree)
        scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False)

        bg_color = [1,1,1] if dataset.white_background else [0, 0, 0]
        background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")

        if not skip_train:
             render_set(dataset.model_path, "train", scene.loaded_iter, scene.getTrainCameras(), gaussians, pipeline, background, dataset.train_test_exp, separate_sh)

        if not skip_test:
             render_set(dataset.model_path, "test", scene.loaded_iter, scene.getTestCameras(), gaussians, pipeline, background, dataset.train_test_exp, separate_sh)

def render_sets_adaptive(dataset: ModelParams, iteration: int, pipeline: PipelineParams,
                         sparse, kind: str):
    """Render the held-out views, run ONE pass of the frozen fixer on each, mix by
    beta*, and score all three stages.

    Writes, under model_path:
        adaptive_<kind>/render|fixed|blended/NNNNN.png
        final_<kind>.json        per-view and mean PSNR / SSIM / LPIPS for each stage

    The fixer is called exactly once per view; the mix afterwards costs nothing.
    """
    from metrics import evaluate_pairs, fmt

    gaussians = GaussianModel(dataset.sh_degree)
    scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False,
                  sparse=sparse)
    bg = torch.tensor([1, 1, 1] if dataset.white_background else [0, 0, 0],
                      dtype=torch.float32, device="cuda")
    support, test = scene.getTrainCameras(), scene.getTestCameras()
    print(f"adaptive inference: fixer={kind} | {len(support)} support / {len(test)} test")

    if kind == "cnfix":
        import numpy as np
        from utils.cnfix_utils import CNFixer, cam_R_C, find_two_refs, relative_pose
        cn = CNFixer(sparse.cnfix_ckpt, ddim_steps=int(sparse.cnfix_steps),
                     timestep=int(sparse.cnfix_timestep))
        sup_R = [cam_R_C(c)[0] for c in support]
        sup_C = [cam_R_C(c)[1] for c in support]
        refs = [(c.original_image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
                for c in support]

        def fix(cam, raw):
            R, C = cam_R_C(cam)
            i, j = find_two_refs(sup_R, sup_C, R, C)
            return cn.fix(raw, [refs[i], refs[j]],
                          relative_pose(R, C, sup_R[i], sup_C[i]),
                          relative_pose(R, C, sup_R[j], sup_C[j])).clamp(0, 1).cuda()
    elif kind == "difix":
        # The DiFix baseline fixer.  Same call convention as during distillation
        # (utils/fixer_utils.py): one step, one UNPOSED reference -- the nearest support
        # view by camera centre -- and the input padded to a multiple of 8.
        import numpy as np
        import torch.nn.functional as Fn
        from PIL import Image
        from utils.fixer_utils import _difix_pipeline
        pipe = _difix_pipeline(sparse)
        centres = torch.stack([c.camera_center for c in support])

        def _pil(t):
            return Image.fromarray((t.detach().clamp(0, 1).permute(1, 2, 0)
                                    * 255).to(torch.uint8).cpu().numpy())

        def _pad8(t):
            _, h, w = t.shape
            return Fn.pad(t[None], (0, (-w) % 8, 0, (-h) % 8), mode="reflect")[0], h, w

        refs = [_pil(_pad8(c.original_image)[0]) for c in support]

        def fix(cam, raw):
            px, h, w = _pad8(raw)
            j = int(torch.argmin((centres - cam.camera_center[None]).norm(dim=1)))
            gen = torch.Generator(device="cuda").manual_seed(int(sparse.fixer_seed))
            o = pipe(prompt="remove degradation", image=_pil(px), ref_image=refs[j],
                     num_inference_steps=1, timesteps=[int(sparse.fixer_timestep)],
                     guidance_scale=0.0, generator=gen).images[0]
            o = torch.from_numpy(np.asarray(o, np.float32) / 255.0).permute(2, 0, 1)
            return o[:, :h, :w].cuda().clamp(0, 1)
    else:
        raise SystemExit(f"--fixer must be cnfix or difix, got {kind!r}")

    root = os.path.join(dataset.model_path, f"adaptive_{kind}")
    for sub in ("render", "fixed", "blended"):
        makedirs(os.path.join(root, sub), exist_ok=True)

    raws, fixes, mixes, gts, names, betas = [], [], [], [], [], []
    for idx, cam in enumerate(tqdm(test, desc="Adaptive inference")):
        with torch.no_grad():
            raw = render(cam, gaussians, pipeline, bg,
                         use_trained_exp=dataset.train_test_exp,
                         separate_sh=SPARSE_ADAM_AVAILABLE)["render"].clamp(0, 1)
        fixed = fix(cam, raw)
        mixed, b = blend_beta(fixed, raw)
        for t, sub in ((raw, "render"), (fixed, "fixed"), (mixed, "blended")):
            torchvision.utils.save_image(t, os.path.join(root, sub, f"{idx:05d}.png"))
        raws.append(raw.cpu())
        fixes.append(fixed.cpu())
        mixes.append(mixed.cpu())
        gts.append(cam.original_image.clone().cpu())
        names.append(cam.image_name)
        betas.append(b)

    out = {}
    print(f"\n{'stage':>10s} " + " ".join(f"{m:>9s}" for m in
                                          ("PSNR", "SSIM", "LP-vgg", "LP-alex")))
    print("-" * 52)
    for nm, imgs in (("render", raws), ("fixed", fixes), ("beta*", mixes)):
        r = evaluate_pairs(imgs, gts, names)
        out[nm] = r
        print(f"{nm:>10s} " + fmt(r["mean"]).replace("PSNR ", "").replace("SSIM ", "")
              .replace("LPIPS-vgg ", "").replace("LPIPS-alex ", ""))
    out_path = os.path.join(dataset.model_path, f"final_{kind}.json")
    json.dump({"arms": out,
               "beta_star": {"mean": float(sum(betas) / len(betas)),
                             "min": float(min(betas)), "max": float(max(betas))}},
              open(out_path, "w"), indent=1)
    print(f"\nbeta*: mean {sum(betas)/len(betas):.3f}  min {min(betas):.3f}  "
          f"max {max(betas):.3f}\n-> {root}\n-> {out_path}")


if __name__ == "__main__":
    # Set up command line argument parser
    parser = ArgumentParser(description="Testing script parameters")
    model = ModelParams(parser, sentinel=True)
    pipeline = PipelineParams(parser)
    sparse = SparseParams(parser)
    parser.add_argument("--iteration", default=-1, type=int)
    parser.add_argument("--skip_train", action="store_true")
    parser.add_argument("--skip_test", action="store_true")
    parser.add_argument("--quiet", action="store_true")
    parser.add_argument("--fixer", default="none", choices=["none", "cnfix", "difix"],
                        help="run the adaptive-inference stage with this fixer and mix "
                             "by beta*.  'none' = plain rendering, as upstream 3DGS.")
    args = get_combined_args(parser)
    print("Rendering " + args.model_path)

    # Initialize system state (RNG)
    safe_state(args.quiet)

    if args.fixer == "none":
        render_sets(model.extract(args), args.iteration, pipeline.extract(args),
                    args.skip_train, args.skip_test, SPARSE_ADAM_AVAILABLE)
    else:
        render_sets_adaptive(model.extract(args), args.iteration,
                             pipeline.extract(args), sparse.extract(args), args.fixer)