scripts/phase1_validation/probe_tnt3.py

#!/usr/bin/env python
"""Round-3 probe — pinpoint whether Native/Ours 挑的是同一批图."""
import os, sys, inspect, json
import numpy as np
from PIL import Image

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

DATASET_ROOT = "/root/autodl-tmp/dataset/tnt"
OUTPUT_ROOT  = "/root/autodl-tmp/SplatAtlas/outputs"
SCENES = [("truck", "PASS"), ("lighthouse", "FAIL")]
THUMB = 256  # reverse-lookup 降采样尺寸


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


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


def load_thumb(path):
    img = Image.open(path).convert('RGB').resize((THUMB, THUMB), Image.LANCZOS)
    return np.asarray(img, dtype=np.float32) / 255.0


# ---------- Probe 11: Reverse lookup Native GT → images_2 ----------
def probe_reverse_lookup(scene):
    sec(f"PROBE 11 — Reverse-lookup Native GT in full images_2 [{scene}]")
    cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}")
    gt_dir = os.path.join(cell, "gt_test_30000")
    img_dir = os.path.join(DATASET_ROOT, scene, "images_2")

    native_gts = sorted([f for f in os.listdir(gt_dir) if f.lower().endswith(('.png','.jpg'))])
    i2_files = sorted([f for f in os.listdir(img_dir) if f.lower().endswith(('.png','.jpg','.jpeg'))])
    print(f"Native GT count : {len(native_gts)}")
    print(f"images_2 count  : {len(i2_files)}")
    print(f"Building {THUMB}x{THUMB} thumb cache for vectorized PSNR match...")

    i2_arr = np.stack([load_thumb(os.path.join(img_dir, f)) for f in i2_files])  # (N, H, W, 3)
    print(f"  cache shape: {i2_arr.shape}")

    print(f"\n{'idx':>3} {'native_gt':>12} -> {'best_match':>15} {'best_psnr':>10} {'margin':>8}")
    print("-" * 60)
    matches = []
    for idx, ng in enumerate(native_gts):
        ng_arr = load_thumb(os.path.join(gt_dir, ng))
        diff = i2_arr - ng_arr[None]
        mse = (diff * diff).mean(axis=(1, 2, 3))
        psnrs = 10 * np.log10(1.0 / np.maximum(mse, 1e-10))
        order = np.argsort(-psnrs)
        best_idx, best_psnr = order[0], psnrs[order[0]]
        margin = best_psnr - psnrs[order[1]]
        best_name = i2_files[best_idx]
        matches.append((ng, best_name, float(best_psnr), float(margin)))
        if idx < 10:
            flag = "" if best_psnr > 30 else " [LOW]"
            print(f"{idx:>3} {ng:>12} -> {best_name:>15} {best_psnr:>10.2f} {margin:>8.2f}{flag}")

    if len(matches) > 10:
        print(f"... (showing first 10 of {len(matches)})")

    match_names = [m[1] for m in matches]
    match_set = set(match_names)
    low_quality = [m for m in matches if m[2] < 30]

    print(f"\n--- Reverse-lookup summary ---")
    print(f"Unique matches : {len(match_set)} / {len(matches)} Native GT")
    print(f"Low-quality (<30 dB) matches: {len(low_quality)}   "
          f"→ {'[WARN] Native GT 可能来自一个 Ours 看不到的图源' if low_quality else '[OK]'}")

    # 对比 Ours split
    from scene.dataset_readers import readColmapSceneInfo
    scene_info = readColmapSceneInfo(*build_args(
        os.path.join(DATASET_ROOT, scene), "images_2"))
    ours_set = set(c.image_name for c in scene_info.test_cameras)

    overlap = match_set & ours_set
    n_only  = match_set - ours_set
    o_only  = ours_set - match_set
    print(f"\n--- Compare to Ours eval=True test split ---")
    print(f"Native真挑 ∩ Ours挑  : {len(overlap):>3}")
    print(f"Native only          : {len(n_only):>3}  e.g. {sorted(n_only)[:5]}")
    print(f"Ours only            : {len(o_only):>3}  e.g. {sorted(o_only)[:5]}")

    if not o_only and not n_only:
        print(f"\n[VERDICT] 两边挑的是完全同一批图 → 纯粹 enumerate 顺序不同")
        # 检查是否 ordering 差异
        our_test_names = [c.image_name for c in scene_info.test_cameras]
        if list(match_names) == our_test_names:
            print(f"         且顺序也完全一致 — baseline 差异不来自 GT 端")
        else:
            print(f"         但 enumerate 顺序不同 — 会影响 render↔GT 配对:")
            print(f"         Native order (first 5): {match_names[:5]}")
            print(f"         Ours   order (first 5): {our_test_names[:5]}")
    elif len(overlap) == 0:
        print(f"\n[VERDICT] 两边完全不相交 → Native 训练时的 dataset_readers 和我们用的不是同一个版本")
    else:
        pct = 100 * len(overlap) / max(len(ours_set), 1)
        print(f"\n[VERDICT] 部分重合 ({pct:.0f}%) → split 逻辑在某维度上漂移")

    return matches, scene_info


# ---------- Probe 12: Baseline JSON 内容 ----------
def probe_baseline_json(scene):
    sec(f"PROBE 12 — Native baseline JSON content [{scene}]")
    cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}")
    for name in ["metrics_test_iter30000.json", "results.json",
                 "cfg_args"]:
        p = os.path.join(cell, name)
        if os.path.exists(p):
            print(f"\n  <{name}>")
            with open(p) as f:
                content = f.read()
            print(content[:800])
            if len(content) > 800:
                print(f"  ... ({len(content)} bytes total)")


# ---------- Probe 13: enumerate 顺序与 LLFF hold 对齐 ----------
def probe_enumerate_order(scene, matches, scene_info):
    sec(f"PROBE 13 — enumerate alignment: Native GT idx i <-> which images_2 file [{scene}]")
    our_names = [c.image_name for c in scene_info.test_cameras]
    print(f"{'idx':>3} {'Native→真图':>18} {'Ours test_cam[i]':>18} {'same?':>6}")
    print("-" * 55)
    n_align = 0
    for i, (ng, bn, bp, _) in enumerate(matches[:min(15, len(matches))]):
        on = our_names[i] if i < len(our_names) else "—"
        same = (bn == on)
        if same: n_align += 1
        print(f"{i:>3} {bn:>18} {on:>18} {'Y' if same else 'N':>6}")
    # full sweep
    full_align = sum(1 for i, (_, bn, _, _) in enumerate(matches)
                     if i < len(our_names) and bn == our_names[i])
    print(f"\nFull enumerate alignment: {full_align} / {len(matches)} positions match")
    if full_align == len(matches):
        print("[OK] Native 和 Ours 的 enumerate 顺序完全一致 — 问题不在 GT/split 层")
    else:
        print(f"[!!] Enumerate 顺序不同 — 我们的 render[i] 对的是 Native GT[j]，i≠j")


def main():
    for scene, label in SCENES:
        print(f"\n\n{'#'*70}\n# SCENE: {scene}  [{label}]\n{'#'*70}")
        matches, scene_info = probe_reverse_lookup(scene)
        probe_baseline_json(scene)
        probe_enumerate_order(scene, matches, scene_info)


if __name__ == "__main__":
    main()