scripts/phase1_validation/probe_tnt4.py

#!/usr/bin/env python
"""Round-4 — renderer-level diagnosis."""
import os, sys, inspect, json, math, struct
import numpy as np
import torch
from PIL import Image
from plyfile import PlyData

sys.path.insert(0, '/root/autodl-tmp/3dgsAtlas_official')
import gsplat

DATASET_ROOT = "/root/autodl-tmp/dataset/tnt"
OUTPUT_ROOT  = "/root/autodl-tmp/SplatAtlas/outputs"
SCENES = [("truck", "PASS"), ("lighthouse", "FAIL")]


def sec(t):
    print("\n" + "=" * 70); print(f"  {t}"); print("=" * 70)


def build_args(source_path, img_dir):
    from scene.dataset_readers import readColmapSceneInfo
    sig = inspect.signature(readColmapSceneInfo)
    a = []
    for i, (k, p) in enumerate(sig.parameters.items()):
        if i == 0:                  a.append(source_path)
        elif i == 1:                a.append(img_dir)
        elif k == "eval":           a.append(True)
        elif k == "train_test_exp": a.append(False)
        else:                       a.append(p.default if p.default != inspect.Parameter.empty else "")
    return a


# ---------- Probe 14: cfg_args (训练超参) ----------
def probe_cfg_args(scene):
    sec(f"PROBE 14 — Training hyperparameters (cfg_args) [{scene}]")
    cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}")
    for name in ["cfg_args", "args.json", "config.json", "training_args.json",
                 "hparams.json"]:
        p = os.path.join(cell, name)
        if os.path.exists(p):
            print(f"  <{name}>")
            content = open(p).read()
            print(content)
            for key in ['white_background', 'sh_degree', 'resolution', 'background',
                        'convert_SHs_python', 'images']:
                if key in content:
                    print(f"    [!] contains key: '{key}'")
            return
    print("  (no cfg_args-like file found)")
    print("  listing cell dir contents:")
    for f in sorted(os.listdir(cell))[:30]:
        p = os.path.join(cell, f)
        t = "DIR" if os.path.isdir(p) else "FIL"
        print(f"    {t}  {f}")


# ---------- Probe 15: PLY SH 完整性 ----------
def probe_ply_struct(scene):
    sec(f"PROBE 15 — PLY structure & SH completeness [{scene}]")
    ply = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}",
                       "point_cloud", "iteration_30000", "point_cloud.ply")
    pd = PlyData.read(ply)
    v = pd['vertex']
    props = v.data.dtype.names
    print(f"  N gaussians: {len(v)}")
    print(f"  N properties: {len(props)}")
    f_rest_count = sum(1 for p in props if p.startswith('f_rest_'))
    print(f"  f_dc_*   count: {sum(1 for p in props if p.startswith('f_dc_'))}")
    print(f"  f_rest_* count: {f_rest_count}  "
          f"(SH full 3-band needs 45; implied sh_degree="
          f"{'3' if f_rest_count == 45 else '?'})")
    # opacity/scale/rot 分布
    for k in ['opacity', 'scale_0', 'f_dc_0', 'f_rest_0']:
        if k in props:
            vals = np.asarray(v[k])
            print(f"  {k:<10} min={vals.min():.3f}  max={vals.max():.3f}  "
                  f"mean={vals.mean():.3f}  std={vals.std():.3f}")


# ---------- Probe 16: Per-image PSNR breakdown ----------
def probe_per_image(scene):
    """直接复刻 render_single.py 的渲染循环，但这次按 Ours order 逐张
    输出 PSNR + 和对应 Native render 同 cam 的 PSNR，找到崩最狠的几张。"""
    sec(f"PROBE 16 — Per-image PSNR breakdown & delta per-cam [{scene}]")
    from scene.dataset_readers import readColmapSceneInfo
    from utils.graphics_utils import getWorld2View2
    from utils.general_utils import PILtoTorch
    from PIL import Image, ImageOps

    source_path = os.path.join(DATASET_ROOT, scene)
    cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}")
    img_dir = "images_2"
    resolution = 2

    scene_info = readColmapSceneInfo(*build_args(source_path, img_dir))
    test_cams = scene_info.test_cameras
    print(f"  test cams: {len(test_cams)}")

    # 构建 Native renders index:  native_render_name → image
    native_render_dir = os.path.join(cell, "renders_test_30000")
    if os.path.isdir(os.path.join(native_render_dir, "renders")):
        native_render_dir = os.path.join(native_render_dir, "renders")
    native_renders = sorted([f for f in os.listdir(native_render_dir)
                             if f.lower().endswith(('.png', '.jpg'))])
    native_gt_dir = os.path.join(cell, "gt_test_30000")
    native_gts = sorted([f for f in os.listdir(native_gt_dir)
                         if f.lower().endswith(('.png', '.jpg'))])

    # 用 reverse-lookup (thumb) 找出 Native idx i 对应的真图名，从而把
    # Native render 按真图名索引
    THUMB = 128
    i2_files = sorted([f for f in os.listdir(os.path.join(source_path, img_dir))
                       if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
    i2_thumbs = np.stack([
        np.asarray(Image.open(os.path.join(source_path, img_dir, f)).convert('RGB')
                   .resize((THUMB, THUMB), Image.LANCZOS), dtype=np.float32) / 255.0
        for f in i2_files])

    native_name_for = {}  # native_idx → real image_name
    for idx, ng in enumerate(native_gts):
        t = np.asarray(Image.open(os.path.join(native_gt_dir, ng)).convert('RGB')
                       .resize((THUMB, THUMB), Image.LANCZOS), dtype=np.float32) / 255.0
        mse = ((i2_thumbs - t) ** 2).mean(axis=(1, 2, 3))
        native_name_for[idx] = i2_files[int(np.argmin(mse))]

    real_to_native_idx = {v: k for k, v in native_name_for.items()}

    # 加载 PLY
    pd = PlyData.read(os.path.join(cell, "point_cloud", "iteration_30000",
                                    "point_cloud.ply"))
    vv = pd['vertex']
    device = torch.device("cuda")

    def t32(x): return torch.tensor(x, dtype=torch.float32, device=device)
    means     = t32(np.stack((vv['x'], vv['y'], vv['z']), -1))
    quats     = t32(np.stack((vv['rot_0'], vv['rot_1'], vv['rot_2'], vv['rot_3']), -1))
    scales    = torch.exp(t32(np.stack((vv['scale_0'], vv['scale_1'], vv['scale_2']), -1)))
    opacities = torch.sigmoid(t32(np.asarray(vv['opacity'])))
    f_dc      = t32(np.stack((vv['f_dc_0'], vv['f_dc_1'], vv['f_dc_2']), -1)).unsqueeze(1)
    f_rest    = t32(np.stack([vv[f'f_rest_{i}'] for i in range(45)], -1))
    f_rest    = f_rest.view(-1, 3, 15).transpose(1, 2)
    shs       = torch.cat([f_dc, f_rest], dim=1)

    bg_color = torch.tensor([0., 0., 0.], device=device)

    print(f"\n  {'i':>3} {'cam_name':>15} {'ours_psnr':>10} {'nat_psnr':>10} "
          f"{'delta':>8} {'render_vs_render':>18}")
    print("  " + "-" * 75)

    ours_psnrs, native_psnrs_this = [], []
    render_vs_render = []
    for i, c in enumerate(test_cams):
        w = int(round(c.width  / resolution))
        h = int(round(c.height / resolution))
        viewmat = t32(getWorld2View2(np.array(c.R), np.array(c.T))).unsqueeze(0)
        fx = w / (2 * math.tan(c.FovX / 2))
        fy = h / (2 * math.tan(c.FovY / 2))
        K = t32(np.array([[fx, 0, w/2], [0, fy, h/2], [0, 0, 1]])).unsqueeze(0)

        with torch.no_grad():
            colors, _, _ = gsplat.rasterization(
                means=means, quats=quats, scales=scales, opacities=opacities,
                colors=shs, viewmats=viewmat, Ks=K, width=w, height=h,
                sh_degree=3, packed=True, render_mode='RGB',
                backgrounds=bg_color.unsqueeze(0))
        ours_rgb = colors[0].clamp(0, 1)

        # 读对应 GT
        pil = Image.open(os.path.join(source_path, img_dir, c.image_name))
        pil = ImageOps.exif_transpose(pil)
        gt_chw = PILtoTorch(pil, (w, h)).to(device)
        if gt_chw.shape[0] == 4: gt_chw = gt_chw[:3]
        gt_rgb = gt_chw.permute(1, 2, 0)

        mse = ((ours_rgb - gt_rgb) ** 2).mean().item()
        ours_p = 10 * math.log10(1.0 / max(mse, 1e-10))
        ours_psnrs.append(ours_p)

        # 对应的 Native render
        nat_p_str, vs_str = "—", "—"
        if c.image_name in real_to_native_idx:
            nidx = real_to_native_idx[c.image_name]
            nr_path = os.path.join(native_render_dir, native_renders[nidx])
            ng_path = os.path.join(native_gt_dir,     native_gts[nidx])
            nr = np.asarray(Image.open(nr_path).convert('RGB'), dtype=np.float32) / 255.0
            ng = np.asarray(Image.open(ng_path).convert('RGB'), dtype=np.float32) / 255.0
            if nr.shape == ng.shape:
                nmse = ((nr - ng) ** 2).mean()
                nat_p = 10 * math.log10(1.0 / max(nmse, 1e-10))
                native_psnrs_this.append(nat_p)
                nat_p_str = f"{nat_p:.2f}"
                # render vs render
                ours_np = (ours_rgb.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
                nr_u8 = (nr * 255).astype(np.uint8)
                if ours_np.shape == nr_u8.shape:
                    vs_mse = ((ours_np.astype(np.float32)/255 - nr_u8.astype(np.float32)/255) ** 2).mean()
                    vs = 10 * math.log10(1.0 / max(vs_mse, 1e-10))
                    render_vs_render.append(vs)
                    vs_str = f"{vs:.2f}"

        d = ours_p - (native_psnrs_this[-1] if native_psnrs_this else ours_p)
        if i < 10 or (nat_p_str != "—" and abs(ours_p - float(nat_p_str)) > 2):
            print(f"  {i:>3} {c.image_name:>15} {ours_p:>10.2f} {nat_p_str:>10} "
                  f"{d:>+8.2f} {vs_str:>18}")

    print(f"\n  Ours mean PSNR                 : {np.mean(ours_psnrs):.4f} dB")
    if native_psnrs_this:
        print(f"  Native mean PSNR (recomputed)  : {np.mean(native_psnrs_this):.4f} dB")
        print(f"  Mean delta                     : "
              f"{np.mean(ours_psnrs) - np.mean(native_psnrs_this):+.4f} dB")
    if render_vs_render:
        arr = np.array(render_vs_render)
        print(f"\n  gsplat_render vs Native_render (per-cam PSNR):")
        print(f"    mean={arr.mean():.2f}  min={arr.min():.2f}  max={arr.max():.2f}")
        print(f"    >35 dB (near identical): {(arr>35).sum()}/{len(arr)}")
        print(f"    25-35 dB                : {((arr>=25)&(arr<=35)).sum()}/{len(arr)}")
        print(f"    <25 dB (visible diff)   : {(arr<25).sum()}/{len(arr)}")
        if arr.mean() > 35:
            print("    [结论] 两个 renderer 输出几乎一致 → 锅不在 renderer")
        elif arr.mean() > 25:
            print("    [结论] renderer 有明显差异，但在合理范围")
        else:
            print("    [结论] renderer 输出差异巨大 → gsplat vs diff-gaussian-rasterization 不等价")


def main():
    for scene, label in SCENES:
        print(f"\n\n{'#'*70}\n# SCENE: {scene}  [{label}]\n{'#'*70}")
        probe_cfg_args(scene)
        probe_ply_struct(scene)
        probe_per_image(scene)


if __name__ == "__main__":
    main()