pipelines/run_full_pipeline.py

#!/usr/bin/env python3
"""
全流程 Pipeline: .blend/.glb/.gltf/.ply → 边渲边选 → ERPT Warp

支持六种模式:

1. 单 Blend 场景:
    python run_full_pipeline.py \
        --blender /path/to/blender \
        --blend /path/to/scene.blend \
        --scene-name my_scene \
        --output-root ./dataset

2. 批量 Blend（扫描 input-dir 下所有 .blend）:
    python run_full_pipeline.py \
        --blender /path/to/blender \
        --input-dir /path/to/blend_files/ \
        --output-root ./dataset

3. 单 GLB/GLTF 场景:
    python run_full_pipeline.py \
        --blender /path/to/blender \
        --glb /path/to/scene.glb \
        --scene-name my_scene \
        --output-root ./dataset

4. 批量 GLB（扫描 input-dir 下所有 .glb/.gltf）:
    python run_full_pipeline.py \
        --blender /path/to/blender \
        --input-dir /path/to/glb_files/ \
        --output-root ./dataset

5. 单 PLY 场景（无需 Blender）:
    python run_full_pipeline.py \
        --ply /path/to/scene.ply \
        --scene-name my_scene \
        --output-root ./dataset

6. 批量 PLY（扫描 input-dir 下所有 .ply）:
    python run_full_pipeline.py \
        --input-dir /path/to/ply_files/ \
        --output-root ./dataset

    加 --dry-run 预览要跑哪些场景
    已跑完的场景自动跳过（--no-skip-done 强制重跑）
"""

import argparse
import json
import os
import shutil
import subprocess
import sys
import time
from pathlib import Path


def run_step1_blend_pipeline(
    blender_exe: str,
    scene_path: str,
    temp_dir: str,
    num_frames: int,
    resolution: str,
    samples: int,
    engine: str,
    exposure: float,
    grid_spacing: float,
    camera_height,
    stop_gain: float,
    stop_score: float,
    stop_delta: float,
    min_frames: int,
    rotation_type: str = "random_yaw",
    gain_curve: bool = True,
    scene_flag: str = "--blend",
) -> int:
    """步骤 1 (Blend/GLB): 调 run_blend_pipeline.py 边渲边选。
    scene_flag: "--blend" 或 "--glb"
    """
    script = Path(__file__).parent / "run_blend_pipeline.py"
    if not script.exists():
        raise FileNotFoundError(f"找不到 run_blend_pipeline.py: {script}")

    cmd = [
        sys.executable, str(script),
        "--blender", blender_exe,
        scene_flag, scene_path,
        "--output-dir", temp_dir,
        "--num-frames", str(num_frames),
        "--render-depth",
        "--resolution", resolution,
        "--samples", str(samples),
        "--engine", engine,
        "--exposure", str(exposure),
        "--grid-spacing", str(grid_spacing),
        "--stop-gain", str(stop_gain),
        "--stop-score", str(stop_score),
        "--stop-delta", str(stop_delta),
        "--min-frames", str(min_frames),
        "--rotation-type", rotation_type,
    ]
    if camera_height is not None:
        cmd += ["--camera-height", str(camera_height)]
    if not gain_curve:
        cmd += ["--no-gain-curve"]

    print(f"\n{'='*60}")
    print("[Step 1] 边渲边选 (Blender Cycles)")
    print(f"{'='*60}")

    proc = subprocess.run(cmd, text=True)
    if proc.returncode != 0:
        print(f"  [Error] run_blend_pipeline 退出码: {proc.returncode}")
        return proc.returncode

    n = sum(1 for f in Path(temp_dir).glob("panorama_*.png"))
    print(f"  渲染完成: {n} 帧")
    return 0


def run_step1_ply_pipeline(
    ply_path: str,
    temp_dir: str,
    num_frames: int,
    resolution: str,
    grid_spacing: float,
    camera_height,
    stop_gain: float,
    stop_score: float,
    stop_delta: float,
    min_frames: int,
    rotation_type: str = "random_yaw",
    point_size: float = 2.0,
    z_up: bool = True,
) -> int:
    """步骤 1 (PLY): 调 run_ply_pipeline.py 边渲边选（无需 Blender）"""
    script = Path(__file__).parent / "run_ply_pipeline.py"
    if not script.exists():
        raise FileNotFoundError(f"找不到 run_ply_pipeline.py: {script}")

    cmd = [
        sys.executable, str(script),
        "--ply", ply_path,
        "--output-dir", temp_dir,
        "--num-frames", str(num_frames),
        "--resolution", resolution,
        "--grid-spacing", str(grid_spacing),
        "--stop-gain", str(stop_gain),
        "--stop-score", str(stop_score),
        "--stop-delta", str(stop_delta),
        "--min-frames", str(min_frames),
        "--rotation-type", rotation_type,
        "--point-size", str(point_size),
    ]
    if camera_height is not None:
        cmd += ["--camera-height", str(camera_height)]
    if not z_up:
        cmd += ["--no-z-up"]

    print(f"\n{'='*60}")
    print("[Step 1] 边渲边选 (PLY 点云)")
    print(f"{'='*60}")

    proc = subprocess.run(cmd, text=True)
    if proc.returncode != 0:
        print(f"  [Error] run_ply_pipeline 退出码: {proc.returncode}")
        return proc.returncode

    n = sum(1 for f in Path(temp_dir).glob("panorama_*.png"))
    print(f"  渲染完成: {n} 帧")
    return 0


def run_step1_hm3d_pipeline(
    blender_exe: str,
    scene_path: str,
    temp_dir: str,
    num_frames: int,
    resolution: str,
    samples: int,
    engine: str,
    exposure: float,
    grid_spacing: float,
    camera_height,
    stop_gain: float,
    stop_score: float,
    stop_delta: float,
    min_frames: int,
    rotation_type: str = "random_yaw",
    gain_curve: bool = True,
) -> int:
    """步骤 1 (HM3D GLB): 调 run_hm3d_pipeline.py 边渲边选。"""
    script = Path(__file__).parent / "run_hm3d_pipeline.py"
    if not script.exists():
        raise FileNotFoundError(f"找不到 run_hm3d_pipeline.py: {script}")

    cmd = [
        sys.executable, str(script),
        "--blender", blender_exe,
        "--glb", scene_path,
        "--output-dir", temp_dir,
        "--num-frames", str(num_frames),
        "--render-depth",
        "--resolution", resolution,
        "--samples", str(samples),
        "--engine", engine,
        "--exposure", str(exposure),
        "--grid-spacing", str(grid_spacing),
        "--stop-gain", str(stop_gain),
        "--stop-score", str(stop_score),
        "--stop-delta", str(stop_delta),
        "--min-frames", str(min_frames),
        "--rotation-type", rotation_type,
        "--hm3d", "True",
    ]
    if camera_height is not None:
        cmd += ["--camera-height", str(camera_height)]
    if not gain_curve:
        cmd += ["--no-gain-curve"]

    print(f"\n{'='*60}")
    print("[Step 1] 边渲边选 (HM3D GLB)")
    print(f"{'='*60}")

    proc = subprocess.run(cmd, text=True)
    if proc.returncode != 0:
        print(f"  [Error] run_hm3d_pipeline 退出码: {proc.returncode}")
        return proc.returncode

    n = sum(1 for f in Path(temp_dir).rglob("panorama_*.png"))
    print(f"  渲染完成: {n} 帧")
    return 0


def run_step2_organize_hm3d(temp_dir: str, scene_dir: str) -> int:
    """步骤 2 (HM3D): 整理多空间目录结构

    temp_dir 里有:
        frame_selection/
        space_00/ (panorama_*.png, *_depth.npy, pose_*.json)
        space_01/
        ...

    整理成（每个 space 一个独立目录）:
        scene_dir/space_00/input/   → 中心帧 RGB + depth + 所有 pose
        scene_dir/space_00/output/  → 所有帧 RGB + depth（GT 真值）
        scene_dir/space_01/input/
        scene_dir/space_01/output/
        ...
        scene_dir/frame_selection/  → 选帧信息
    """
    temp = Path(temp_dir)

    print(f"\n{'='*60}")
    print("[Step 2] 整理目录结构 (HM3D 多空间)")
    print(f"{'='*60}")

    space_dirs = sorted(
        [d for d in temp.iterdir() if d.is_dir() and d.name.startswith("space_")]
    )
    if not space_dirs:
        print("  [Error] 没有找到 space_XX 目录")
        return 1

    print(f"  共 {len(space_dirs)} 个空间")

    for space_d in space_dirs:
        space_name = space_d.name

        rgb_files = sorted(space_d.glob("panorama_*.png"))
        if not rgb_files:
            print(f"  {space_name}: 无渲染结果，跳过")
            continue

        n_frames = len(rgb_files)
        out_space_dir = Path(scene_dir) / space_name
        inp_dir = out_space_dir / "input"
        out_dir = out_space_dir / "output"
        inp_dir.mkdir(parents=True, exist_ok=True)
        out_dir.mkdir(parents=True, exist_ok=True)

        for rgb_path in rgb_files:
            shutil.copy2(str(rgb_path), str(out_dir / rgb_path.name))
            depth_path = space_d / rgb_path.name.replace(".png", "_depth.npy")
            if depth_path.exists():
                shutil.copy2(str(depth_path), str(out_dir / depth_path.name))

        center_rgb = space_d / "panorama_0000.png"
        center_depth = space_d / "panorama_0000_depth.npy"
        if center_rgb.exists():
            shutil.copy2(str(center_rgb), str(inp_dir / center_rgb.name))
        if center_depth.exists():
            shutil.copy2(str(center_depth), str(inp_dir / center_depth.name))

        n_pose = 0
        for pose_path in sorted(space_d.glob("pose_*.json")):
            shutil.copy2(str(pose_path), str(inp_dir / pose_path.name))
            n_pose += 1

        print(f"  {space_name}: {n_frames} 帧 → output/, 中心帧 + {n_pose} pose → input/")

    sel_dir = Path(scene_dir) / "frame_selection"
    sel_dir.mkdir(parents=True, exist_ok=True)

    sel_json = temp / "frame_selection" / "selected_frames.json"
    if sel_json.exists():
        shutil.copy2(str(sel_json), str(sel_dir / "selected_frames.json"))

    cand_npy = temp / "frame_selection" / "candidates_filtered.npy"
    if cand_npy.exists():
        shutil.copy2(str(cand_npy), str(sel_dir / "candidates_filtered.npy"))

    return 0


def run_step2_organize(temp_dir: str, scene_dir: str) -> int:
    """步骤 2: 整理目录结构

    temp_dir 里有:
        panorama_0000.png, panorama_0000_depth.npy, pose_0000.json, ...

    整理成:
        scene_dir/input/   → 中心帧 RGB + depth + 所有 pose（供 ERPT warp 使用）
        scene_dir/output/  → 所有帧 RGB + depth（GT 真值）
    """
    temp = Path(temp_dir)
    inp_dir = Path(scene_dir) / "input"
    out_dir = Path(scene_dir) / "output"
    inp_dir.mkdir(parents=True, exist_ok=True)
    out_dir.mkdir(parents=True, exist_ok=True)

    print(f"\n{'='*60}")
    print("[Step 2] 整理目录结构")
    print(f"{'='*60}")

    # 找所有帧
    rgb_files = sorted(temp.glob("panorama_*.png"))
    if not rgb_files:
        print("  [Error] 没有找到渲染的全景图")
        return 1

    n_frames = len(rgb_files)
    print(f"  共 {n_frames} 帧")

    # output/: 复制所有帧的 RGB + depth（GT 真值）
    for rgb_path in rgb_files:
        shutil.copy2(str(rgb_path), str(out_dir / rgb_path.name))
        # depth
        depth_path = temp / rgb_path.name.replace(".png", "_depth.npy")
        if depth_path.exists():
            shutil.copy2(str(depth_path), str(out_dir / depth_path.name))

    print(f"  output/: {n_frames} 帧 RGB + depth")

    # input/: 中心帧 RGB + depth + 所有 pose
    center_rgb = temp / "panorama_0000.png"
    center_depth = temp / "panorama_0000_depth.npy"

    if center_rgb.exists():
        shutil.copy2(str(center_rgb), str(inp_dir / center_rgb.name))
    if center_depth.exists():
        shutil.copy2(str(center_depth), str(inp_dir / center_depth.name))

    # 所有 pose
    n_pose = 0
    for pose_path in sorted(temp.glob("pose_*.json")):
        shutil.copy2(str(pose_path), str(inp_dir / pose_path.name))
        n_pose += 1

    print(f"  input/: 中心帧 + {n_pose} 个 pose")

    # 复制选帧信息（供参考）
    sel_dir = inp_dir / "frame_selection"
    sel_dir.mkdir(parents=True, exist_ok=True)

    sel_json = temp / "frame_selection" / "selected_frames.json"
    if sel_json.exists():
        shutil.copy2(str(sel_json), str(sel_dir / "selected_frames.json"))

    cand_npy = temp / "frame_selection" / "candidates_filtered.npy"
    if cand_npy.exists():
        shutil.copy2(str(cand_npy), str(sel_dir / "candidates_filtered.npy"))

    # 增益曲线（从 selected_frames.json 读数据，用 PIL 画）
    if sel_json.exists():
        try:
            draw_gain_curve(str(sel_dir / "selected_frames.json"),
                            str(sel_dir / "gain_curve.jpg"))
            print(f"  增益曲线: {sel_dir}/gain_curve.jpg")
        except Exception as e:
            print(f"  [跳过] 画增益曲线失败: {e}")

    return 0


def draw_gain_curve(json_path, output_path):
    """画增益曲线（优先 matplotlib，fallback PIL）"""
    with open(json_path) as f:
        data = json.load(f)

    frames = [fr for fr in data["frames"] if not fr.get("skipped")]
    if len(frames) < 2:
        return

    fids = [fr["frame_id"] for fr in frames]
    pred_gains = [fr["gain"] for fr in frames]
    actual_gains = [fr["actual_gain"] for fr in frames]
    scores = [fr["score"] for fr in frames]
    deltas = [fr["delta_ratio"] for fr in frames]

    try:
        import matplotlib
        matplotlib.use('Agg')
        import matplotlib.pyplot as plt

        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6), sharex=True)

        ax1.plot(fids, pred_gains, 'o-', color='#2196F3', label='predicted', markersize=3, linewidth=1.5)
        ax1.plot(fids, actual_gains, 'o-', color='#FF9800', label='actual', markersize=3, linewidth=1.5)
        ax1.axhline(y=0.05, color='red', linestyle='--', alpha=0.5, label='stop_gain=5%')
        ax1.set_ylabel('gain')
        ax1.set_ylim(-0.05, 1.05)
        ax1.legend(loc='upper right', fontsize=9)
        ax1.set_title(f'Gain Curve ({len(frames)} frames)', fontsize=11)
        ax1.grid(True, alpha=0.3)
        # 标注首末值
        ax1.annotate(f'{actual_gains[0]:.0%}', (fids[0], actual_gains[0]),
                     textcoords="offset points", xytext=(5, 5), fontsize=7, color='#FF9800')
        ax1.annotate(f'{actual_gains[-1]:.0%}', (fids[-1], actual_gains[-1]),
                     textcoords="offset points", xytext=(-25, 5), fontsize=7, color='#FF9800')

        ax2.plot(fids, scores, 'D-', color='#4CAF50', label='score', markersize=3, linewidth=1.5)
        ax2.plot(fids, deltas, 's-', color='#9C27B0', label='delta', markersize=2, linewidth=1.2)
        ax2.axhline(y=-0.33, color='red', linestyle='--', alpha=0.5, label='stop_score=-0.33')
        ax2.axhline(y=0.01, color='#9C27B0', linestyle=':', alpha=0.4, label='stop_delta=1%')
        ax2.set_ylabel('value')
        ax2.set_xlabel('frame')
        ax2.legend(loc='upper right', fontsize=9)
        ax2.grid(True, alpha=0.3)
        # 标注首末值
        ax2.annotate(f'{deltas[0]:.1%}', (fids[0], deltas[0]),
                     textcoords="offset points", xytext=(5, 5), fontsize=7, color='#9C27B0')
        ax2.annotate(f'{deltas[-1]:.1%}', (fids[-1], deltas[-1]),
                     textcoords="offset points", xytext=(-25, -10), fontsize=7, color='#9C27B0')

        plt.tight_layout()
        plt.savefig(output_path, dpi=150, bbox_inches='tight')
        plt.close()
        return

    except ImportError:
        pass

    # ---- fallback: PIL ----
    try:
        from PIL import Image, ImageDraw, ImageFont

        W, H = 800, 500
        ML, MR, MT, MB = 50, 20, 30, 25
        MID = H // 2
        pw = W - ML - MR

        img = Image.new("RGB", (W, H), "white")
        draw = ImageDraw.Draw(img)
        try:
            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 9)
        except Exception:
            font = ImageFont.load_default()

        n = len(frames)

        def px(i, v, y0, y1, vmin, vmax):
            x = ML + int(i / max(n-1,1) * pw)
            y = y0 + int((1 - (v-vmin)/(vmax-vmin)) * (y1-y0))
            return x, max(y0, min(y1, y))

        def line(pts, color, y0, y1, vmin, vmax):
            for j in range(len(pts)-1):
                draw.line([px(j,pts[j],y0,y1,vmin,vmax), px(j+1,pts[j+1],y0,y1,vmin,vmax)], fill=color, width=2)

        line(pred_gains, "#2196F3", MT, MID-5, 0, 1.05)
        line(actual_gains, "#FF9800", MT, MID-5, 0, 1.05)
        line(scores, "#4CAF50", MID+10, H-MB, -0.6, 1.05)
        line(deltas, "#9C27B0", MID+10, H-MB, -0.6, 1.05)

        draw.text((ML, MT-12), f"Gain ({n} frames)", fill="black", font=font)
        draw.text((ML, MID+2), "Score / Delta", fill="black", font=font)
        img.save(output_path, quality=90)

    except ImportError:
        pass


def run_step3_erpt_warp_hm3d(scene_dir: str, device: str = "cuda") -> int:
    """步骤 3 (HM3D): 对每个空间调 run_pipeline.py 执行 ERPT warp

    遍历 scene_dir/space_XX/input/，对帧数 >= 2 的空间生成 warp 文件
    """
    script = Path(__file__).parent / "run_pipeline.py"
    if not script.exists():
        raise FileNotFoundError(f"找不到 run_pipeline.py: {script}")

    print(f"\n{'='*60}")
    print("[Step 3] ERPT Warp (HM3D 多空间)")
    print(f"{'='*60}")

    scene_path = Path(scene_dir)
    space_dirs = sorted(
        [d for d in scene_path.iterdir()
         if d.is_dir() and d.name.startswith("space_")]
    )

    if not space_dirs:
        print("  [Error] 没有找到 space_XX 目录")
        return 1

    total_ret = 0
    n_warped = 0
    n_skipped = 0

    for space_d in space_dirs:
        inp_dir = space_d / "input"
        if not inp_dir.exists():
            continue

        n_poses = len(list(inp_dir.glob("pose_*.json")))
        if n_poses < 2:
            print(f"  {space_d.name}: {n_poses} pose，跳过 warp")
            n_skipped += 1
            continue

        print(f"\n  [{space_d.name}] ERPT Warp ({n_poses} poses)...")

        cmd = [
            sys.executable, str(script),
            "--stage", "warp_only",
            "--data_dir", str(inp_dir),
            "--output_dir", str(inp_dir),
            "--device", device,
            "--center_frame", "0",
        ]

        proc = subprocess.run(cmd, text=True)
        if proc.returncode != 0:
            print(f"    [Error] {space_d.name} warp 失败 (退出码: {proc.returncode})")
            total_ret = proc.returncode
            continue

        warp_rgb_dir = inp_dir / "warp_rgb"
        warp_depth_dir = inp_dir / "warp_depth"
        keep_suffixes = ("_rgb.png", "_mask.png", "_depth_range.npy")

        n_moved = 0
        for subdir in [warp_rgb_dir, warp_depth_dir]:
            if subdir.exists():
                for f in subdir.iterdir():
                    if f.is_file() and any(f.name.endswith(s) for s in keep_suffixes):
                        shutil.move(str(f), str(inp_dir / f.name))
                        n_moved += 1
                shutil.rmtree(str(subdir), ignore_errors=True)

        n_warped += 1
        print(f"    warp 文件: {n_moved} 个")

    print(f"\n  Warp 完成: {n_warped} 个空间, 跳过 {n_skipped} 个")
    return total_ret


def run_step3_erpt_warp(scene_dir: str, device: str = "cuda") -> int:
    """步骤 3: 调 run_pipeline.py 执行 ERPT warp

    读取 scene_dir/input/ 里的中心帧 + pose → 生成 warp 文件
    warp 文件直接写到 input/ 目录
    """
    script = Path(__file__).parent / "run_pipeline.py"
    if not script.exists():
        raise FileNotFoundError(f"找不到 run_pipeline.py: {script}")

    inp_dir = Path(scene_dir) / "input"

    print(f"\n{'='*60}")
    print("[Step 3] ERPT Warp")
    print(f"{'='*60}")

    cmd = [
        sys.executable, str(script),
        "--stage", "warp_only",
        "--data_dir", str(inp_dir),
        "--output_dir", str(inp_dir),
        "--device", device,
        "--center_frame", "0",
    ]

    proc = subprocess.run(cmd, text=True)
    if proc.returncode != 0:
        print(f"  [Error] run_pipeline 退出码: {proc.returncode}")
        return proc.returncode

    # 把 warp 子目录里需要的文件提到 input/ 根目录
    # 只保留: _rgb.png, _mask.png, _depth_range.npy
    warp_rgb_dir = inp_dir / "warp_rgb"
    warp_depth_dir = inp_dir / "warp_depth"
    keep_suffixes = ("_rgb.png", "_mask.png", "_depth_range.npy")

    n_moved = 0
    for subdir in [warp_rgb_dir, warp_depth_dir]:
        if subdir.exists():
            for f in subdir.iterdir():
                if f.is_file() and any(f.name.endswith(s) for s in keep_suffixes):
                    shutil.move(str(f), str(inp_dir / f.name))
                    n_moved += 1
            # 删除整个子目录（包含不需要的 flow/weight_sum/comparison 等）
            shutil.rmtree(str(subdir), ignore_errors=True)

    print(f"  warp 文件已移到 input/: {n_moved} 个")
    return 0


def is_already_done(output_root, scene_name):
    """检查是否已经跑完"""
    sel_path = os.path.join(
        output_root, scene_name, "input", "frame_selection",
        "selected_frames.json")
    if not os.path.exists(sel_path):
        # 也检查 input/ 下直接放的
        sel_path = os.path.join(output_root, scene_name, "input",
                                "selected_frames.json")
        if not os.path.exists(sel_path):
            return False
    try:
        with open(sel_path) as f:
            data = json.load(f)
        return data.get("total_frames", 0) > 0
    except Exception:
        return False


def find_glb_files(input_dir):
    """递归查找所有 .glb / .gltf 文件，返回 [(glb_path, scene_name), ...]

    scene_name 取文件名（不含扩展名），或第一级子目录名（如有子目录）。
    """
    input_dir = os.path.abspath(input_dir)
    glb_files = []
    for root, dirs, files in os.walk(input_dir):
        for f in files:
            if f.lower().endswith(".glb") or f.lower().endswith(".gltf"):
                glb_path = os.path.join(root, f)
                rel = os.path.relpath(root, input_dir)
                if rel == ".":
                    scene_name = os.path.splitext(f)[0]
                else:
                    scene_name = rel.split(os.sep)[0]
                glb_files.append((glb_path, scene_name))
    glb_files.sort(key=lambda x: x[1])
    return glb_files


def find_blend_files(input_dir):
    """递归查找所有 .blend 文件，返回 [(blend_path, scene_name), ...]
    
    scene_name 取 input_dir 下的第一级子目录名（scene_indoor_XXXX），
    不管 .blend 文件嵌套了几层。
    
    例如:
      input_dir = /path/to/dataset/indoor
      .blend 在 /path/to/dataset/indoor/scene_indoor_0001/1407m1/xxx.blend
      → scene_name = scene_indoor_0001
    """
    input_dir = os.path.abspath(input_dir)
    blend_files = []
    for root, dirs, files in os.walk(input_dir):
        for f in files:
            if f.endswith(".blend"):
                blend_path = os.path.join(root, f)
                rel = os.path.relpath(root, input_dir)
                scene_name = rel.split(os.sep)[0]
                blend_files.append((blend_path, scene_name))
    blend_files.sort(key=lambda x: x[1])
    return blend_files


def find_ply_files(input_dir):
    """递归查找所有 .ply 文件，返回 [(ply_path, scene_name), ...]

    scene_name 取文件名（不含扩展名），或第一级子目录名（如有子目录）。
    """
    input_dir = os.path.abspath(input_dir)
    ply_files = []
    for root, dirs, files in os.walk(input_dir):
        for f in files:
            if f.lower().endswith(".ply"):
                ply_path = os.path.join(root, f)
                rel = os.path.relpath(root, input_dir)
                if rel == ".":
                    scene_name = os.path.splitext(f)[0]
                else:
                    scene_name = rel.split(os.sep)[0]
                ply_files.append((ply_path, scene_name))
    ply_files.sort(key=lambda x: x[1])
    return ply_files


def run_single_scene(args, scene_path, scene_name, scene_type="blend"):
    """跑单个场景，返回 0=成功 / 非0=失败

    scene_type: "blend" | "glb" | "hm3d" | "ply"
    """
    output_root = str(Path(args.output_root).resolve())
    scene_dir = os.path.join(output_root, scene_name)
    temp_dir = os.path.join(scene_dir, "_render_temp")
    os.makedirs(scene_dir, exist_ok=True)

    type_labels = {
        "blend": ".blend (Blender Cycles)",
        "glb":   ".glb/.gltf (Blender Cycles)",
        "hm3d":  ".glb/.gltf (HM3D 渲染)",
        "ply":   ".ply (点云渲染)",
    }
    type_label = type_labels.get(scene_type, scene_type)
    print("=" * 60)
    print(f"全流程 Pipeline: {type_label} → 边渲边选 → ERPT Warp")
    print("=" * 60)
    print(f"  Scene:  {scene_name}")
    print(f"  Input:  {scene_path}")
    print(f"  Output: {scene_dir}/")
    t_start = time.time()

    # Step 1
    if scene_type == "hm3d":
        ret = run_step1_hm3d_pipeline(
            blender_exe=args.blender,
            scene_path=scene_path,
            temp_dir=temp_dir,
            num_frames=args.num_frames,
            resolution=args.resolution,
            samples=args.samples,
            engine=args.engine,
            exposure=args.exposure,
            grid_spacing=args.grid_spacing,
            camera_height=args.camera_height,
            stop_gain=args.stop_gain,
            stop_score=args.stop_score,
            stop_delta=args.stop_delta,
            min_frames=args.min_frames,
            rotation_type=args.rotation_type,
            gain_curve=getattr(args, "gain_curve", True),
        )
    elif scene_type in ("blend", "glb"):
        scene_flag = "--blend" if scene_type == "blend" else "--glb"
        ret = run_step1_blend_pipeline(
            blender_exe=args.blender,
            scene_path=scene_path,
            temp_dir=temp_dir,
            num_frames=args.num_frames,
            resolution=args.resolution,
            samples=args.samples,
            engine=args.engine,
            exposure=args.exposure,
            grid_spacing=args.grid_spacing,
            camera_height=args.camera_height,
            stop_gain=args.stop_gain,
            stop_score=args.stop_score,
            stop_delta=args.stop_delta,
            min_frames=args.min_frames,
            rotation_type=args.rotation_type,
            gain_curve=getattr(args, "gain_curve", True),
            scene_flag=scene_flag,
        )
    else:
        ret = run_step1_ply_pipeline(
            ply_path=scene_path,
            temp_dir=temp_dir,
            num_frames=args.num_frames,
            resolution=args.resolution,
            grid_spacing=args.grid_spacing,
            camera_height=args.camera_height,
            stop_gain=args.stop_gain,
            stop_score=args.stop_score,
            stop_delta=args.stop_delta,
            min_frames=args.min_frames,
            rotation_type=args.rotation_type,
            point_size=getattr(args, "point_size", 2.0),
            z_up=getattr(args, "z_up", True),
        )

    if ret != 0:
        print(f"[Error] Step 1 失败")
        return ret

    # Step 2
    if scene_type == "hm3d":
        ret = run_step2_organize_hm3d(temp_dir, scene_dir)
    else:
        ret = run_step2_organize(temp_dir, scene_dir)
    if ret != 0:
        print(f"[Error] Step 2 失败")
        return ret

    # Step 3
    if not args.skip_warp:
        if scene_type == "hm3d":
            ret = run_step3_erpt_warp_hm3d(scene_dir, device=args.device)
        else:
            ret = run_step3_erpt_warp(scene_dir, device=args.device)
        if ret != 0:
            print(f"[Error] Step 3 失败")

    # 清理
    if os.path.exists(temp_dir):
        shutil.rmtree(temp_dir, ignore_errors=True)

    dt = time.time() - t_start
    print(f"\n{'='*60}")
    print(f"完成! {scene_name}, {dt:.1f}s ({dt/60:.1f}min)")
    print(f"{'='*60}")
    return 0


def run_single(args):
    """单场景模式（blend、glb、hm3d 或 ply）"""
    if args.ply:
        ply_path = str(Path(args.ply).resolve())
        scene_name = args.scene_name or Path(args.ply).stem
        ret = run_single_scene(args, ply_path, scene_name, scene_type="ply")
    elif args.hm3d:
        glb_path = str(Path(args.hm3d).resolve())
        scene_name = args.scene_name or Path(args.hm3d).stem
        ret = run_single_scene(args, glb_path, scene_name, scene_type="hm3d")
    elif args.glb:
        glb_path = str(Path(args.glb).resolve())
        scene_name = args.scene_name or Path(args.glb).stem
        ret = run_single_scene(args, glb_path, scene_name, scene_type="glb")
    else:
        blend_path = str(Path(args.blend).resolve())
        scene_name = args.scene_name or Path(args.blend).stem
        ret = run_single_scene(args, blend_path, scene_name, scene_type="blend")
    if ret != 0:
        sys.exit(1)


def run_batch(args):
    """批量模式（自动检测 .blend / .glb / .gltf / .ply）"""
    input_dir_abs = os.path.abspath(args.input_dir)
    if not os.path.isdir(input_dir_abs):
        print(f"[Error] --input-dir 目录不存在: {input_dir_abs}")
        print(f"        (原始参数: {args.input_dir})")
        sys.exit(1)

    # 没有 --blender → 只能跑 PLY
    if not getattr(args, "blender", None):
        scene_files = find_ply_files(args.input_dir)
        scene_type = "ply"
        ext_label = ".ply"
    else:
        # 有 blender：优先 .blend，其次 .glb/.gltf (HM3D)，最后 .ply
        scene_files = find_blend_files(args.input_dir)
        scene_type = "blend"
        ext_label = ".blend"
        if not scene_files:
            scene_files = find_glb_files(args.input_dir)
            scene_type = "hm3d"  # 默认使用 HM3D 渲染管线
            ext_label = ".glb/.gltf (HM3D)"
        if not scene_files:
            scene_files = find_ply_files(args.input_dir)
            scene_type = "ply"
            ext_label = ".ply"

    if not scene_files:
        print(f"[Error] 在 {args.input_dir} 下没找到 {ext_label} 文件")
        sys.exit(1)

    output_root = str(Path(args.output_root).resolve())

    print(f"{'='*60}")
    print(f"批量处理模式 ({ext_label})")
    print(f"{'='*60}")
    print(f"  输入目录: {args.input_dir}")
    print(f"  输出目录: {output_root}")
    print(f"  找到 {len(scene_files)} 个 {ext_label} 文件")
    # input(f"  按 Enter 键继续，或 Ctrl+C 取消...")
    
    to_run = []
    skipped = []
    for scene_path, scene_name in scene_files:
        if args.skip_done and is_already_done(output_root, scene_name):
            skipped.append((scene_path, scene_name))
        else:
            to_run.append((scene_path, scene_name))

    if skipped:
        print(f"  跳过 {len(skipped)} 个已完成:")
        for _, sn in skipped:
            print(f"    ✓ {sn}")

    print(f"  待处理 {len(to_run)} 个:")
    for bp, sn in to_run:
        print(f"    → {sn}  ({os.path.basename(bp)})")

    if args.dry_run:
        print(f"\n[Dry run] 不实际运行")
        return

    if not to_run:
        print(f"\n全部已完成!")
        return

    t_all = time.time()
    success = []
    failed = []

    for idx, (scene_path, scene_name) in enumerate(to_run):
        print(f"\n{'='*60}")
        print(f"[{idx+1}/{len(to_run)}] {scene_name}")
        print(f"{'='*60}")

        t_scene = time.time()
        try:
            ret = run_single_scene(args, scene_path, scene_name, scene_type)
            dt = time.time() - t_scene
            if ret == 0:
                success.append((scene_name, dt))
                print(f"\n  ✓ {scene_name} ({dt:.0f}s)")
            else:
                failed.append((scene_name, f"exit code {ret}"))
                print(f"\n  ✗ {scene_name} 失败 ({dt:.0f}s)")
        except Exception as e:
            dt = time.time() - t_scene
            failed.append((scene_name, str(e)))
            print(f"\n  ✗ {scene_name} 异常: {e} ({dt:.0f}s)")

    dt_all = time.time() - t_all
    print(f"\n{'='*60}")
    print(f"批量处理完成")
    print(f"{'='*60}")
    print(f"  总耗时: {dt_all:.0f}s ({dt_all/60:.1f}min = {dt_all/3600:.1f}h)")
    print(f"  成功: {len(success)} 个")
    for sn, dt in success:
        print(f"    ✓ {sn} ({dt:.0f}s)")
    if failed:
        print(f"  失败: {len(failed)} 个")
        for sn, reason in failed:
            print(f"    ✗ {sn}: {reason}")
    if skipped:
        print(f"  跳过: {len(skipped)} 个 (已完成)")


def main():
    parser = argparse.ArgumentParser(
        description="全流程: .blend/.glb/.ply → 边渲边选 → ERPT Warp"
    )

    # 输入（四种模式互斥）
    input_group = parser.add_mutually_exclusive_group()
    input_group.add_argument("--blend", type=str, default=None,
                             help=".blend 场景文件路径（单 Blend 场景模式）")
    input_group.add_argument("--glb", type=str, default=None,
                             help=".glb / .gltf 场景文件路径（单 GLB 场景模式）")
    input_group.add_argument("--hm3d", type=str, default=None,
                             help=".glb / .gltf 场景文件路径（单 HM3D 场景模式）")
    input_group.add_argument("--ply", type=str, default=None,
                             help=".ply 场景文件路径（单 PLY 场景模式）")

    parser.add_argument("--input-dir", type=str, default=None,
                        help="包含场景文件的根目录（批量模式，自动检测 .blend/.glb/.ply）")
    parser.add_argument("--scene-name", type=str, default=None,
                        help="场景名（默认从文件名提取）")
    parser.add_argument("--output-root", type=str, default="./dataset",
                        help="输出根目录（默认 ./dataset）")

    # Blender 参数（仅 Blend/GLB/HM3D 模式需要）
    parser.add_argument("--blender", type=str, default=None,
                        help="Blender 可执行文件路径（Blend/GLB/HM3D 模式必须）")
    parser.add_argument("--samples", type=int, default=128)
    parser.add_argument("--engine", type=str, default="CYCLES")
    parser.add_argument("--exposure", type=float, default=0.0)
    parser.add_argument("--gain-curve", action="store_true", default=True,
                        help="画增益曲线 (默认开启)")
    parser.add_argument("--no-gain-curve", dest="gain_curve", action="store_false")

    # PLY 参数（仅 PLY 模式）
    parser.add_argument("--point-size", type=float, default=2.0,
                        help="点云渲染点径（像素），PLY 模式有效（默认 2.0）")
    parser.add_argument("--z-up", action="store_true", default=True,
                        help="PLY 坐标系为 Z-up（默认 True）")
    parser.add_argument("--no-z-up", dest="z_up", action="store_false",
                        help="PLY 坐标系为 Y-up（已是 ERPT_native，不转换）")

    # 通用渲染参数
    parser.add_argument("--num-frames", type=int, default=30)
    parser.add_argument("--resolution", type=str, default="2048,1024")

    # 选帧参数
    parser.add_argument("--grid-spacing", type=float, default=0.5)
    parser.add_argument("--camera-height", type=float, default=None)
    parser.add_argument("--stop-gain", type=float, default=0.08)
    parser.add_argument("--stop-score", type=float, default=-0.3)
    parser.add_argument("--stop-delta", type=float, default=0.08)
    parser.add_argument("--min-frames", type=int, default=5)
    parser.add_argument("--rotation-type", type=str, default="random_yaw",
                        choices=["none", "rotate_x_90", "rotate_x_180",
                                 "rotate_z_90", "random_yaw"])

    # ERPT 参数
    parser.add_argument("--device", type=str, default="cuda")
    parser.add_argument("--skip-warp", action="store_true",
                        help="只做步骤 1+2，跳过 ERPT warp")

    # 批量模式参数
    parser.add_argument("--skip-done", action="store_true", default=True,
                        help="跳过已跑完的场景（默认开启）")
    parser.add_argument("--no-skip-done", action="store_true",
                        help="强制重跑所有场景")
    parser.add_argument("--dry-run", action="store_true",
                        help="只列出要跑的场景，不实际运行")

    args = parser.parse_args()

    if args.no_skip_done:
        args.skip_done = False

    # 校验 Blend/GLB/HM3D 模式必须提供 --blender
    if (args.blend or args.glb or args.hm3d) and not args.blender:
        parser.error("--blend / --glb / --hm3d 模式必须同时提供 --blender 可执行文件路径")

    # 模式判定
    if args.input_dir:
        run_batch(args)
    elif args.blend or args.glb or args.hm3d or args.ply:
        run_single(args)
    else:
        parser.error("必须指定 --blend / --glb / --hm3d / --ply（单场景）或 --input-dir（批量）")


if __name__ == "__main__":
    main()