scripts/phase1_validation/run_full_batch.py

import os
import sys
import subprocess
import json
import csv
import glob
from tqdm import tqdm
import numpy as np

BASE_DIR = "/root/autodl-tmp/SplatAtlas"
OUT_BASE = os.path.join(BASE_DIR, "outputs")
DATASET_ROOT = "/root/autodl-tmp/dataset"
CSV_PATH = os.path.join(BASE_DIR, "scripts/phase1_validation/validation_results_final.csv")
ERR_LOG  = os.path.join(BASE_DIR, "scripts/phase1_validation/final_batch_errors.log")

# 全量 20 种方法
ALL_METHODS = [
    "3dgsmcmc", "analyticsplatting", "atomgs", "coadaptgs", "conegs", 
    "erankgs", "ges", "ghap", "gof", "gslpm", "lightgaussian", "opti3dgs", 
    "pgsr", "reactgs", "steepgs", "vanilla_3dgs", "absgs", "gaussianpro", 
    "minisplatting", "pixelgs"
]

def build_scene_database():
    """绝对白名单机制，屏蔽衍生目录"""
    db = {}
    allowed_roots = ["360", "tnt", "deepblending_clean"]
    
    for root in allowed_roots:
        search_path = os.path.join(DATASET_ROOT, root, "*")
        for path in glob.glob(search_path):
            if not os.path.isdir(path): continue
            
            scene_name = os.path.basename(path).lower()
            
            if root == "tnt":
                res, dataset_type = 2, "tnt"
            elif root == "deepblending_clean":
                res, dataset_type = 1, "db"
            elif root == "360":
                outdoor = ['garden', 'stump', 'treehill', 'bicycle', 'flowers']
                res = 4 if any(o in scene_name for o in outdoor) else 2
                dataset_type = "mip360_outdoor" if res == 4 else "mip360_indoor"
                
            db[scene_name] = {"source": path, "res": res, "dataset": dataset_type, "bg": "0,0,0"}
    return db

def log_error(cell_name, phase, stdout, stderr):
    with open(ERR_LOG, "a") as f:
        f.write(f"\n{'='*50}\n[{phase} ERROR] Cell: {cell_name}\n")
        if stdout: f.write(f"--- STDOUT ---\n" + "\n".join(stdout.strip().split("\n")[-20:]) + "\n")
        if stderr: f.write(f"--- STDERR ---\n" + "\n".join(stderr.strip().split("\n")[-20:]) + "\n")

def get_actual_cell_dir(method, scene):
    target = f"{method}_{scene}".lower()
    if os.path.exists(OUT_BASE):
        for d in os.listdir(OUT_BASE):
            if d.lower() == target:
                return os.path.join(OUT_BASE, d)
    return None

def phase_a_native(method, scene_key, scene_cfg, cell_dir):
    json_path = os.path.join(cell_dir, "metrics_test_iter30000.json")
    if os.path.exists(json_path):
        try:
            with open(json_path, 'r') as f:
                data = json.load(f)
                if "PSNR" in data or ("photometric" in data and "PSNR" in data["photometric"]):
                    return True
        except: pass

    cmd_render = [
        "python", "scripts/main_render.py", "--method", method,
        "--source_path", scene_cfg["source"], "--model_path", cell_dir,
        "--iteration", "30000", "--resolution", str(scene_cfg["res"])
    ]
    res = subprocess.run(cmd_render, cwd=BASE_DIR, capture_output=True, text=True)
    if res.returncode != 0:
        log_error(f"{method}_{scene_key}", "PHASE_A_RENDER", res.stdout, res.stderr)
        return False

    py_metric = f"""
import sys, json, os
sys.path.append("ufd_evalkit")
from photometric import compute_photometric_metrics
r_dir = "{cell_dir}/renders_test_30000/renders"
g_dir = "{cell_dir}/renders_test_30000/gt"
if not os.path.exists(r_dir):
    r_dir = "{cell_dir}/renders_test_30000"
    g_dir = "{cell_dir}/gt_test_30000"
metrics = compute_photometric_metrics(r_dir, g_dir)
with open("{json_path}", "w") as f:
    json.dump({{"PSNR": metrics.get("PSNR", 0)}}, f, indent=4)
"""
    res2 = subprocess.run(["python", "-c", py_metric], cwd=BASE_DIR, capture_output=True, text=True)
    if res2.returncode != 0:
        log_error(f"{method}_{scene_key}", "PHASE_A_METRIC", res2.stdout, res2.stderr)
        return False
    return True

def phase_b_gsplat(ply_path, scene_cfg, cell_dir, cell_name):
    out_dir = os.path.join(OUT_BASE, "phase1_validation", cell_name)
    os.makedirs(out_dir, exist_ok=True)
    cmd = [
        "python", "scripts/phase1_validation/render_single.py",
        "--ply_path", ply_path, "--source_path", scene_cfg["source"],
        "--model_path", cell_dir, "--output_dir", out_dir,
        "--resolution", str(scene_cfg["res"]), "--bg_color", scene_cfg["bg"]
    ]
    res = subprocess.run(cmd, cwd=BASE_DIR, capture_output=True, text=True)
    
    if res.returncode != 0:
        log_error(cell_name, "PHASE_B_GSPLAT", res.stdout, res.stderr)
        return None, None, None, "ERROR_B"

    native_psnr, gsplat_psnr, delta = None, None, None
    for line in res.stdout.split('\n'):
        if "Mean PSNR (ours/gsplat)" in line:
            try: gsplat_psnr = float(line.split(":")[1].strip().split()[0])
            except: pass
        if "Mean PSNR (native baseline)" in line:
            try: native_psnr = float(line.split(":")[1].strip().split()[0])
            except: pass
        if "Delta" in line and "dB" in line and "STATUS" not in line:
            try: delta = float(line.split(":")[1].strip().split()[0])
            except: pass
            
    if gsplat_psnr is None:
        try:
            with open(os.path.join(out_dir, "psnr_results.json"), "r") as f:
                d = json.load(f)
                gsplat_psnr = float(np.mean(list(d.values())))
        except: pass

    notes = ""
    if gsplat_psnr is not None and gsplat_psnr < 15.0:
        notes = "low PSNR"

    return native_psnr, gsplat_psnr, delta, notes

def main():
    print("=== SplatAtlas Phase 1 Final Full Batch ===")
    SCENES = build_scene_database()
    print(f"[Info] Loaded {len(SCENES)} verified scenes.")
    
    completed = set()
    write_header = not os.path.exists(CSV_PATH)
    if not write_header:
        with open(CSV_PATH, "r") as f:
            reader = csv.DictReader(f)
            for row in reader:
                completed.add((row['method'], row['scene']))

    valid_cells = []
    for method in ALL_METHODS:
        for scene, cfg in SCENES.items():
            cell_dir = get_actual_cell_dir(method, scene)
            if cell_dir:
                valid_cells.append((method, scene, cfg, cell_dir))

    print(f"[Info] Found {len(valid_cells)} pre-trained cells across {len(ALL_METHODS)} methods.\n")

    pbar = tqdm(total=len(valid_cells), desc="Validating Cells", unit="cell")

    with open(CSV_PATH, "a", newline='') as f:
        writer = csv.writer(f)
        if write_header:
            writer.writerow(["method", "scene", "dataset", "resolution", "has_ply", "native_psnr", "gsplat_psnr", "delta_db", "status", "notes"])

        for method, scene, cfg, cell_dir in valid_cells:
            if (method, scene) in completed:
                pbar.update(1)
                continue

            ply_path = os.path.join(cell_dir, "point_cloud/iteration_30000/point_cloud.ply")
            if not os.path.exists(ply_path):
                writer.writerow([method, scene, cfg["dataset"], cfg["res"], "False", "", "", "", "SKIP", "PLY missing"])
                f.flush()
                pbar.set_postfix_str(f"[{method}/{scene}] SKIP")
                pbar.update(1)
                continue

            cell_name = os.path.basename(cell_dir)
            
            if not phase_a_native(method, scene, cfg, cell_dir):
                writer.writerow([method, scene, cfg["dataset"], cfg["res"], "True", "", "", "", "ERROR_A", "Native render failed"])
                f.flush()
                pbar.set_postfix_str(f"[{method}/{scene}] ERROR_A")
                pbar.update(1)
                continue

            native, gsplat, delta, notes = phase_b_gsplat(ply_path, cfg, cell_dir, cell_name)
            
            if native is None and gsplat is None:
                writer.writerow([method, scene, cfg["dataset"], cfg["res"], "True", "", "", "", "ERROR_B", notes])
                f.flush()
                pbar.set_postfix_str(f"[{method}/{scene}] ERROR_B")
            else:
                if delta is None and native is not None and gsplat is not None:
                    delta = gsplat - native
                
                status = "PASS" # 基于你的要求，弱化硬阈值，仅记录数据
                if delta is not None and abs(delta) >= 1.0:
                    status = "FAIL"
                    notes = f"Systematic Bias | {notes}".strip(" | ")
                    
                writer.writerow([method, scene, cfg["dataset"], cfg["res"], "True", 
                                 f"{native:.4f}" if native else "", 
                                 f"{gsplat:.4f}" if gsplat else "", 
                                 f"{delta:.4f}" if delta else "", 
                                 status, notes])
                f.flush()
                pbar.set_postfix_str(f"[{method}/{scene}] {status}")
            
            pbar.update(1)
    
    pbar.close()
    print("\n========================================")
    print(f"FINAL BATCH COMPLETED. Results saved to: {CSV_PATH}")
    print("========================================")

if __name__ == "__main__":
    main()