pipelines/run_pipeline.py

#!/usr/bin/env python3
"""
ERPT Pipeline 主入口（Forward Warp，深度估计可选）

使用方法：
    # 默认：使用已有深度真值做 warp（不加载深度估计权重）
    python run_pipeline.py --stage warp_only --data_dir /path/to/scene

    # 强制完整流程（深度估计 + warp）
    python run_pipeline.py --stage all

    # 仅深度估计
    python run_pipeline.py --stage depth_only
"""

import argparse
import re
import time
from pathlib import Path
from typing import Dict, Any, List, Optional

import yaml
import numpy as np
import torch
import cv2

# 添加模块路径
import sys
sys.path.insert(0, str(Path(__file__).parent))

# Warp 相关（始终加载）
from core.erp_warp import warp_erp_to_target, WarpResult, create_comparison_image
from utils.io_utils import load_image, save_image, load_json, save_json, save_depth
from utils.pose_utils import Pose, load_pose

# 深度估计相关（延迟加载，仅 depth_only / all 模式才 import）
_depth_modules_loaded = False


def _load_depth_modules():
    """延迟加载深度估计模块（避免 warp_only 模式加载大模型权重）"""
    global _depth_modules_loaded
    if _depth_modules_loaded:
        return
    global build_icosahedron_slices, extract_all_tangents, compute_coverage_mask
    global estimate_all_tangent_depths
    global fuse_tangent_depths_to_erp, save_depth_visualization, visualize_depth

    from core.tangent_extraction import (
        build_icosahedron_slices,
        extract_all_tangents,
        compute_coverage_mask,
    )
    from core.depth_estimation import estimate_all_tangent_depths
    from core.depth_fusion import (
        fuse_tangent_depths_to_erp,
        save_depth_visualization,
        visualize_depth,
    )
    _depth_modules_loaded = True
    print("[Depth] 深度估计模块已加载")


# =============================================================================
# 数据发现
# =============================================================================

def discover_image_files(directory: Path) -> dict:
    """自动发现目录中的全景图文件"""
    image_extensions = ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG']
    image_files = []
    for ext in image_extensions:
        image_files.extend(directory.glob(f"*{ext}"))
    image_files = sorted(image_files)
    if not image_files:
        return {}

    result = {}
    for img_path in image_files:
        stem = img_path.stem
        match = re.search(r'[_-](\d+)$', stem)
        if match:
            result[int(match.group(1))] = img_path
            continue
        if stem.isdigit():
            result[int(stem)] = img_path

    if not result:
        for idx, img_path in enumerate(image_files):
            result[idx] = img_path

    return result


def discover_pose_files(directory: Path) -> dict:
    """自动发现目录中的位姿文件"""
    pose_files = sorted(directory.glob("*.json"))
    result = {}
    for pose_path in pose_files:
        stem = pose_path.stem
        if stem in ['meta', 'config', 'stats', 'cameras', 'render_meta', 'description']:
            continue
        match = re.search(r'[_-](\d+)$', stem)
        if match:
            result[int(match.group(1))] = pose_path
            continue
        if stem.isdigit():
            result[int(stem)] = pose_path

    return result


# =============================================================================
# 配置加载
# =============================================================================

def load_config(config_path: Path) -> Dict[str, Any]:
    with open(config_path, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def resolve_paths(cfg: Dict[str, Any], config_dir: Path) -> Dict[str, Any]:
    """解析相对路径为绝对路径"""
    data_cfg = cfg.get("data", {})
    for key in ["data_dir", "output_dir", "depth_dir"]:
        if key in data_cfg and data_cfg[key]:
            path = Path(data_cfg[key])
            if not path.is_absolute():
                data_cfg[key] = str(config_dir / path)

    depth_pro_cfg = cfg.get("depth_pro", {})
    if "repo_dir" in depth_pro_cfg and depth_pro_cfg["repo_dir"]:
        rp = Path(depth_pro_cfg["repo_dir"])
        if not rp.is_absolute():
            depth_pro_cfg["repo_dir"] = str(config_dir / rp)

    cfg["_project_root"] = str(config_dir)
    return cfg


# =============================================================================
# 深度估计流程（仅 all / depth_only 模式调用）
# =============================================================================

def run_depth_pipeline(
    center_rgb: np.ndarray,
    cfg: Dict[str, Any],
    device: torch.device,
    output_dir: Path,
    erp_h: int,
    erp_w: int,
    frame_id: int = 0,
) -> np.ndarray:
    """运行深度估计全流程：切片 -> 推理 -> 融合"""
    _load_depth_modules()

    depth_out_dir = output_dir / "depth_erp"
    depth_out_dir.mkdir(parents=True, exist_ok=True)

    # --- Step 1: 构建切片规格 ---
    print(f"\n{'='*60}")
    print(f"[Step 1] Building tangent slices (frame {frame_id})")
    print(f"{'='*60}")

    if "erp" not in cfg:
        cfg["erp"] = {}
    cfg["erp"]["height"] = erp_h
    cfg["erp"]["width"] = erp_w

    slices = build_icosahedron_slices(cfg)
    print(f"  Total slices: {len(slices)}")
    for s in slices:
        if s.slice_type != "face":
            print(f"    {s.slice_id}: type={s.slice_type}, fov={s.fov_deg:.1f}°")

    coverage_mask, coverage_stats = compute_coverage_mask(slices, erp_h, erp_w, device)
    print(f"  Coverage: {coverage_stats['total_coverage']:.2f}%")

    dbg_dir = output_dir / "debug"
    dbg_dir.mkdir(parents=True, exist_ok=True)
    save_image(np.stack([coverage_mask] * 3, axis=-1), dbg_dir / "coverage_mask.png")

    # --- Step 2: 提取切片 ---
    print(f"\n{'='*60}")
    print(f"[Step 2] Extracting tangent slices (frame {frame_id})")
    print(f"{'='*60}")

    t0 = time.time()
    tangent_rgbs = extract_all_tangents(center_rgb, slices, device)
    print(f"  Extracted {len(tangent_rgbs)} slices in {time.time()-t0:.2f}s")

    if cfg.get("run", {}).get("save_intermediates", False):
        tangent_dir = output_dir / "tangents"
        tangent_dir.mkdir(parents=True, exist_ok=True)
        for slice_id, rgb in tangent_rgbs.items():
            save_image(rgb, tangent_dir / f"{slice_id}_rgb.png")

    # --- Step 3: Depth Pro 推理 ---
    print(f"\n{'='*60}")
    print(f"[Step 3] Running Depth Pro inference (frame {frame_id})")
    print(f"{'='*60}")

    dp_cfg = cfg.get("depth_pro", {})
    if not bool(dp_cfg.get("enabled", True)):
        print("  [Warning] Depth Pro disabled, using dummy depth")
        tangent_depths = {}
        for sid, rgb in tangent_rgbs.items():
            tangent_depths[sid] = np.full(rgb.shape[:2], 5.0, dtype=np.float32)
    else:
        t0 = time.time()
        tangent_depths = estimate_all_tangent_depths(
            tangent_rgbs, slices, cfg, device,
        )
        print(f"  Estimated {len(tangent_depths)} depths in {time.time()-t0:.2f}s")

    if cfg.get("run", {}).get("save_intermediates", False):
        tangent_dir = output_dir / "tangents"
        for sid, depth in tangent_depths.items():
            save_depth(depth, tangent_dir / f"{sid}_depth.npy")

    # --- Step 4: 融合到 ERP ---
    print(f"\n{'='*60}")
    print(f"[Step 4] Fusing tangent depths to ERP (frame {frame_id})")
    print(f"{'='*60}")

    t0 = time.time()
    depth_erp, weight_sum, valid_mask = fuse_tangent_depths_to_erp(
        tangent_depths, slices, cfg, device,
        debug_dir=dbg_dir if cfg.get("run", {}).get("save_intermediates", False) else None,
    )
    print(f"  Fused in {time.time()-t0:.2f}s")

    valid_ratio = np.sum(valid_mask > 0) / (erp_h * erp_w)
    valid_depths = depth_erp[np.isfinite(depth_erp) & (depth_erp > 0)]
    if len(valid_depths) > 0:
        print(f"  Valid depth ratio: {valid_ratio * 100:.2f}%")
        print(f"  Depth range: [{valid_depths.min():.2f}, {valid_depths.max():.2f}] m")

    # --- Step 5: 保存结果 ---
    save_depth(depth_erp, depth_out_dir / f"depth_{frame_id:04d}.npy")
    save_depth_visualization(depth_erp, depth_out_dir / f"depth_{frame_id:04d}_vis.png")
    cv2.imwrite(str(depth_out_dir / f"depth_{frame_id:04d}_valid_mask.png"), valid_mask * 255)

    return depth_erp


# =============================================================================
# Warp 流程
# =============================================================================

def run_warp_pipeline(
    center_rgb: np.ndarray,
    depth_erp: np.ndarray,
    center_frame: int,
    image_files: dict,
    pose_files: dict,
    cfg: Dict[str, Any],
    device: torch.device,
    output_dir: Path,
    erp_h: int,
    erp_w: int,
) -> None:
    """运行 warp 全流程：遍历目标帧，执行 forward splatting"""
    warp_cfg = cfg.get("warp", {})
    output_depth = bool(warp_cfg.get("output_depth", True))

    # 确定目标帧列表
    available_targets = sorted([fid for fid in pose_files.keys() if fid != center_frame])

    cfg_targets = warp_cfg.get("target_frames", None)
    if cfg_targets is not None and cfg_targets != "auto":
        cfg_set = set(int(t) for t in cfg_targets)
        target_frames = [fid for fid in available_targets if fid in cfg_set]
    else:
        target_frames = available_targets

    print(f"\n{'='*60}")
    print(f"[Warp] Forward splatting from frame {center_frame}")
    print(f"{'='*60}")
    print(f"  Method: {warp_cfg.get('method', 'softmax_splatting')}")
    print(f"  Available targets with pose: {available_targets}")
    print(f"  Will warp: {target_frames}")

    # 加载中心帧位姿
    if center_frame not in pose_files:
        print(f"  [Error] Center pose not found for frame {center_frame}")
        return
    src_pose = load_pose(pose_files[center_frame])
    print(f"  Source pose: position={src_pose.position.tolist()}")

    # 输出目录
    warp_rgb_dir = output_dir / "warp_rgb"
    warp_rgb_dir.mkdir(parents=True, exist_ok=True)
    if output_depth:
        warp_depth_dir = output_dir / "warp_depth"
        warp_depth_dir.mkdir(parents=True, exist_ok=True)

    total_warp = len(target_frames)
    for idx, tgt_id in enumerate(target_frames):
        if tgt_id not in pose_files:
            print(f"  [{idx+1}/{total_warp}] Frame {tgt_id}: pose not found, skip")
            continue

        tgt_pose = load_pose(pose_files[tgt_id])
        print(f"  [{idx+1}/{total_warp}] Frame {center_frame} -> {tgt_id} ...", end="", flush=True)

        t0 = time.time()
        result = warp_erp_to_target(
            src_rgb=center_rgb,
            src_depth=depth_erp,
            src_pose=src_pose,
            tgt_pose=tgt_pose,
            cfg=cfg,
            device=device,
        )
        dt = time.time() - t0

        valid_pct = result.valid_mask.sum() / result.valid_mask.size * 100
        print(f" done ({dt:.2f}s, valid={valid_pct:.1f}%)")

        prefix = f"pano{center_frame:04d}_to_pano{tgt_id:04d}"

        # 保存 warped RGB
        save_image(result.warped_rgb, warp_rgb_dir / f"{prefix}_rgb.png")

        # 保存 valid mask
        cv2.imwrite(str(warp_rgb_dir / f"{prefix}_mask.png"), result.valid_mask * 255)

        # 保存 warped depth
        if output_depth and result.warped_depth is not None:
            save_depth(result.warped_depth, warp_depth_dir / f"{prefix}_depth_range.npy")

    print(f"  Warp complete. Output saved to: {warp_rgb_dir}")


# =============================================================================
# 主函数
# =============================================================================

def main():
    _script_dir = Path(__file__).parent
    _default_config = _script_dir / "config.yaml"

    parser = argparse.ArgumentParser(description="ERPT Pipeline")
    parser.add_argument("--config", type=str,
                        default=str(_default_config) if _default_config.exists() else None,
                        help="Config file path")
    parser.add_argument("--data_dir", type=str, default=None,
                        help="Data directory (overrides config)")
    parser.add_argument("--output_dir", type=str, default=None,
                        help="Output directory (overrides config)")
    parser.add_argument("--device", type=str, default="cuda")
    parser.add_argument("--stage", type=str, default="warp_only",
                        choices=["all", "depth_only", "warp_only"])
    parser.add_argument("--center_frame", type=int, default=None,
                        help="Center frame ID (overrides config)")
    args = parser.parse_args()

    # 加载配置
    if args.config:
        config_path = Path(args.config)
        cfg = load_config(config_path)
        cfg = resolve_paths(cfg, config_path.parent)
    else:
        cfg = {
            "data": {},
            "erp": {"auto_size": True},
            "tangent": {},
            "depth_pro": {"enabled": True, "precision": "fp16", "pass_f_px": True},
            "fusion": {"blend_mode": "multiband", "output_scale": 1.10, "k": 4},
            "run": {"save_intermediates": False},
        }

    # 命令行覆盖
    if args.data_dir:
        cfg["data"]["data_dir"] = str(Path(args.data_dir).resolve())
    if args.output_dir:
        cfg["data"]["output_dir"] = args.output_dir

    data_dir = Path(cfg["data"].get("data_dir", "inputs"))
    output_dir = Path(cfg["data"].get("output_dir", "outputs"))

    device = torch.device(args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu")
    print(f"Using device: {device}")

    center_frame = args.center_frame or int(cfg.get("warp", {}).get("center_frame", 0))

    print(f"\n{'='*60}")
    print("ERPT Pipeline")
    print(f"{'='*60}")
    print(f"Stage: {args.stage}")
    print(f"Data dir: {data_dir}")
    print(f"Output dir: {output_dir}")

    t_start = time.time()

    # --- 加载数据 ---
    print(f"\n{'='*60}")
    print("[Loading data]")
    print(f"{'='*60}")

    image_files = discover_image_files(data_dir)
    pose_files = discover_pose_files(data_dir)
    print(f"  Found {len(image_files)} images, {len(pose_files)} poses")

    if not image_files:
        raise FileNotFoundError(f"No image files found in: {data_dir}")

    if center_frame not in image_files:
        center_frame = sorted(image_files.keys())[0]
        print(f"  Using frame {center_frame} as center")

    center_rgb = load_image(image_files[center_frame])
    print(f"  Center image: {image_files[center_frame].name}")
    print(f"  Shape: {center_rgb.shape}")

    erp_cfg = cfg.get("erp", {})
    if bool(erp_cfg.get("auto_size", True)):
        erp_h, erp_w = center_rgb.shape[:2]
        print(f"  Auto size: {erp_w}x{erp_h}")
    else:
        erp_h = int(erp_cfg.get("height", 2048))
        erp_w = int(erp_cfg.get("width", 4096))

    # --- 深度加载 / 估计 ---
    depth_erp = None

    if args.stage == "all":
        print(f"\n  [Stage: all] 强制执行深度估计")
        depth_erp = run_depth_pipeline(
            center_rgb, cfg, device, output_dir, erp_h, erp_w, center_frame,
        )

    elif args.stage == "depth_only":
        depth_erp = run_depth_pipeline(
            center_rgb, cfg, device, output_dir, erp_h, erp_w, center_frame,
        )

    elif args.stage == "warp_only":
        # 搜索已有深度（真值 > 已估计结果），不回退到深度估计
        depth_candidates = []

        if center_frame in image_files:
            stem = image_files[center_frame].stem
            depth_candidates.append(data_dir / f"{stem}_depth.npy")
            depth_candidates.append(data_dir / f"{stem}_depth.exr")
            depth_candidates.append(data_dir / f"{stem}.npy")
        depth_candidates.append(data_dir / f"depth_{center_frame:04d}.npy")
        depth_candidates.append(output_dir / "depth_erp" / f"depth_{center_frame:04d}.npy")

        for dp in depth_candidates:
            if dp.exists():
                if dp.suffix == ".exr":
                    depth_erp = cv2.imread(str(dp), cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
                    if depth_erp is not None and depth_erp.ndim == 3:
                        depth_erp = depth_erp[:, :, 0]
                    depth_erp = depth_erp.astype(np.float32) if depth_erp is not None else None
                else:
                    depth_erp = np.load(str(dp)).astype(np.float32)
                if depth_erp is not None:
                    print(f"  Loaded depth from {dp}")
                    break

        # 尺寸校验
        if depth_erp is not None and depth_erp.shape != (erp_h, erp_w):
            old_shape = depth_erp.shape
            depth_erp = cv2.resize(depth_erp, (erp_w, erp_h), interpolation=cv2.INTER_LINEAR)
            print(f"  [Warning] Depth resized: {old_shape} -> ({erp_h}, {erp_w})")

        # 没找到深度 → 报错（不回退到深度估计）
        if depth_erp is None:
            tried = "\n    ".join(str(p) for p in depth_candidates)
            raise FileNotFoundError(
                f"[warp_only] 未找到深度文件，无法执行 warp。\n"
                f"已搜索路径:\n    {tried}\n"
                f"如需深度估计请使用 --stage all"
            )

    # --- Warp 阶段 ---
    warp_cfg = cfg.get("warp", {})
    warp_enabled = bool(warp_cfg.get("enabled", True))

    if args.stage in ("all", "warp_only") and warp_enabled:
        run_warp_pipeline(
            center_rgb, depth_erp, center_frame,
            image_files, pose_files,
            cfg, device, output_dir, erp_h, erp_w,
        )

    # --- 完成 ---
    total_time = time.time() - t_start
    print(f"\n{'='*60}")
    print("Pipeline Complete")
    print(f"{'='*60}")
    print(f"Total time: {total_time:.2f}s")
    print(f"Output saved to: {output_dir}")


if __name__ == "__main__":
    main()