ImplicitSurfaceNetwork-xj/application/SDF_Visualize/visualize_sdf.py

#!/usr/bin/env python3
"""
visualize_sdf.py  (Enhanced Version)
────────────────────────────────────────────────────────────────────────
功能：
  · 【二维】各 subface / 合成 SDF 的 XZ 截面热力图
  · 【三维】将 SDF 值映射到 OBJ 模型顶点，输出带颜色的 PLY 文件
             并生成 matplotlib 3D 散点图预览

用法：
  # 仅二维（原有功能）
  python visualize_sdf.py sdf_slice.txt

  # 二维 + 三维（新增功能）
  python visualize_sdf.py sdf_slice.txt --obj path/to/model.obj

  # 更多选项
  python visualize_sdf.py sdf_slice.txt --obj model.obj --subface -1 --no-2d
  python visualize_sdf.py sdf_slice.txt --obj model.obj --subface 0

选项：
  --obj FILE        OBJ 模型路径（启用三维可视化）
  --subface  INT    指定用于三维着色的 subface 索引（-1 = 合成 SDF，默认 -1）
  --no-2d           跳过二维热力图，仅生成三维可视化
  --grid    INT     二维可视化网格分辨率覆盖（不影响 txt 文件中已有数据）
  --out-dir  DIR    输出目录（默认与输入文件相同目录）

依赖：
  pip install numpy matplotlib scipy trimesh
  （trimesh 可选，仅三维功能需要；若未安装则自动跳过 PLY 导出）
────────────────────────────────────────────────────────────────────────
"""

import sys
import argparse
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import matplotlib.ticker as mticker
from pathlib import Path
from scipy.ndimage import gaussian_filter
from scipy.interpolate import RegularGridInterpolator
from matplotlib.gridspec import GridSpec

# ── surface_type 枚举（与 C++ 侧对齐）────────────────────────────────
SURFACE_NAMES = {
    0: "Plane (cap)",
    1: "Sphere",
    2: "Cylinder",
    3: "Cone",
    4: "ExtrudePolyline\nSide Face",
    5: "ExtrudeHelixline\nSide Face",
}

# 全局共享色彩归一化（二维与三维保持一致）
_SHARED_NORM = None
_SHARED_VMAX = None

# 设置全局样式
plt.rcParams.update({
    'font.size': 9,
    'axes.titlesize': 10,
    'axes.labelsize': 9,
    'xtick.labelsize': 8,
    'ytick.labelsize': 8,
    'legend.fontsize': 8,
    'figure.titlesize': 12,
    'figure.autolayout': False,
    'savefig.bbox': 'tight',
    'savefig.pad_inches': 0.1,
})


# ══════════════════════════════════════════════════════════════════════
# 1. 数据读取
# ══════════════════════════════════════════════════════════════════════

def load_sdf_data(path: str) -> dict:
    """解析 dump_sdf_slice 输出文件，返回结构化字典。"""
    lines = Path(path).read_text().splitlines()
    idx = 0

    tok = lines[idx].split(); idx += 1
    grid_res = int(tok[0])
    y_slice  = float(tok[1])

    tok = lines[idx].split(); idx += 1
    x0, x1, z0, z1 = float(tok[0]), float(tok[1]), float(tok[2]), float(tok[3])

    ns = int(lines[idx]); idx += 1

    N = grid_res + 1
    subfaces = []
    for _ in range(ns):
        stype = int(lines[idx]); idx += 1
        rows = []
        for _ in range(N):
            rows.append(list(map(float, lines[idx].split())))
            idx += 1
        subfaces.append({"type": stype, "grid": np.array(rows, dtype=np.float64)})

    return dict(grid_res=grid_res, y_slice=y_slice,
                x0=x0, x1=x1, z0=z0, z1=z1, subfaces=subfaces)


def load_obj_vertices(obj_path: str):
    """
    轻量级 OBJ 顶点读取（不依赖 trimesh）。
    返回 (vertices: np.ndarray [N,3], faces: list[list[int]])
    """
    vertices, faces = [], []
    with open(obj_path, 'r', encoding='utf-8', errors='ignore') as f:
        for line in f:
            line = line.strip()
            if line.startswith('v '):
                coords = list(map(float, line.split()[1:4]))
                vertices.append(coords)
            elif line.startswith('f '):
                # 支持 f v, f v/vt, f v/vt/vn 格式
                tokens = line.split()[1:]
                face = [int(t.split('/')[0]) - 1 for t in tokens]  # OBJ 下标从 1 开始
                faces.append(face)
    return np.array(vertices, dtype=np.float64), faces


# ══════════════════════════════════════════════════════════════════════
# 2. SDF 插值（2D grid → 3D 顶点）
# ══════════════════════════════════════════════════════════════════════

def build_combined_grid(data: dict) -> np.ndarray:
    """计算合成 SDF（所有 subface 取 max，即 CSG 求交）。"""
    subfaces = data['subfaces']
    if not subfaces:
        raise ValueError("No subface data available.")
    combined = subfaces[0]['grid'].copy()
    for sf in subfaces[1:]:
        combined = np.maximum(combined, sf['grid'])
    return combined


def compute_vertex_sdf(vertices: np.ndarray, data: dict, subface_idx: int = -1) -> np.ndarray:
    """
    对三维顶点数组用双线性插值计算 SDF 值。

    参数
    ----
    vertices   : shape (N, 3)，世界坐标
    data       : load_sdf_data 返回的字典
    subface_idx: -1 = 合成 SDF；>= 0 = 指定 subface

    返回
    ----
    sdf_values : shape (N,)
    """
    x0, x1 = data['x0'], data['x1']
    z0, z1 = data['z0'], data['z1']
    N = data['grid_res'] + 1

    if subface_idx == -1:
        grid = build_combined_grid(data)
    else:
        grid = data['subfaces'][subface_idx]['grid']

    # 构建 RegularGridInterpolator（Z 为行，X 为列）
    X_coords = np.linspace(x0, x1, N)
    Z_coords = np.linspace(z0, z1, N)
    # grid[iz, ix]，所以 points 顺序为 (Z, X)
    interp = RegularGridInterpolator(
        (Z_coords, X_coords), grid,
        method='linear', bounds_error=False,
        fill_value=None  # 超出范围外推最近值
    )

    # 使用 (x, z) 坐标插值，忽略 y（截面投影）
    query_points = np.column_stack([vertices[:, 2], vertices[:, 0]])  # (Z, X)
    sdf_values = interp(query_points)

    return sdf_values


# ══════════════════════════════════════════════════════════════════════
# 3. 共享色彩归一化（二维与三维一致）
# ══════════════════════════════════════════════════════════════════════

def build_shared_norm(data: dict, vertex_sdf: np.ndarray = None):
    """
    计算全局色彩归一化，覆盖二维网格数据与三维顶点 SDF，
    保证两端可视化颜色映射完全一致。
    """
    all_values = []
    for sf in data['subfaces']:
        g = sf['grid'].ravel()
        all_values.append(g[np.isfinite(g)])
    # 合成 SDF
    comb = build_combined_grid(data).ravel()
    all_values.append(comb[np.isfinite(comb)])
    if vertex_sdf is not None:
        all_values.append(vertex_sdf[np.isfinite(vertex_sdf)])

    all_values = np.concatenate(all_values)
    vmax = float(np.percentile(np.abs(all_values), 98))
    vmax = max(vmax, 1e-6)
    norm = mcolors.TwoSlopeNorm(vmin=-vmax, vcenter=0.0, vmax=vmax)
    return norm, vmax


# ══════════════════════════════════════════════════════════════════════
# 4. 二维热力图（保留原有功能，微调以使用共享 norm）
# ══════════════════════════════════════════════════════════════════════

def draw_panel(ax, XX, ZZ, grid: np.ndarray, title: str, y_slice: float,
               cmap=None, norm=None, highlight_sign_flip=False):
    if cmap is None:
        cmap = plt.cm.RdBu_r

    finite = grid[np.isfinite(grid)]
    if finite.size == 0:
        ax.set_title(title + "\n[no finite data]")
        return

    if norm is None:
        vmax = float(np.percentile(np.abs(finite), 98))
        vmax = max(vmax, 1e-6)
        norm = mcolors.TwoSlopeNorm(vmin=-vmax, vcenter=0.0, vmax=vmax)

    if grid.shape[0] > 10 and grid.shape[1] > 10:
        grid_smoothed = gaussian_filter(grid, sigma=0.7, mode='nearest')
    else:
        grid_smoothed = grid

    cf = ax.contourf(XX, ZZ, grid_smoothed, levels=50, cmap=cmap, norm=norm,
                     alpha=0.85, extend='both')
    cbar = plt.colorbar(cf, ax=ax, pad=0.01, fraction=0.045, shrink=0.9)
    cbar.set_label("SDF", fontsize=7, labelpad=2)
    cbar.ax.tick_params(labelsize=6, pad=1)

    try:
        cs0 = ax.contour(XX, ZZ, grid_smoothed, levels=[0.0],
                         colors=["lime"], linewidths=1.8, zorder=5, alpha=0.9)
        if len(cs0.collections) > 0:
            ax.clabel(cs0, fmt=" SDF=0 ", fontsize=7, inline=True,
                      colors="lime", zorder=6, inline_spacing=5)
    except Exception:
        pass

    if highlight_sign_flip and finite.size > 4:
        from scipy.ndimage import generic_filter
        def local_range(vals):
            return vals.max() - vals.min() if vals.size > 0 else 0
        local_range_grid = generic_filter(grid, local_range, size=3)
        mean_range = float(np.percentile(local_range_grid, 85))
        sign_flip_mask = (local_range_grid > mean_range * 2.5) & np.isfinite(grid)
        if sign_flip_mask.any():
            ax.contourf(XX, ZZ, sign_flip_mask.astype(float),
                        levels=[0.5, 1.5], colors=["yellow"], alpha=0.25, zorder=4)
            ax.contour(XX, ZZ, sign_flip_mask.astype(float), levels=[0.5],
                       colors=["yellow"], linewidths=1.0, linestyles="dotted",
                       zorder=5, alpha=0.7)
            sign_flip_pct = sign_flip_mask.sum() / sign_flip_mask.size * 100
            if sign_flip_pct > 0.5:
                ax.text(0.02, 0.98, f"⚠ {sign_flip_pct:.1f}% sign-flip",
                        transform=ax.transAxes, fontsize=6, color="black",
                        bbox=dict(boxstyle="round,pad=0.2", fc="yellow",
                                  alpha=0.7, ec="black", lw=0.5),
                        zorder=10, verticalalignment='top')

    neg_pct = np.sum(grid < 0) / grid.size * 100
    info = (f"Neg: {neg_pct:.1f}%\nPos: {100-neg_pct:.1f}%\n"
            f"Min: {np.min(grid):.2e}\nMax: {np.max(grid):.2e}")
    ax.text(0.02, 0.02, info, transform=ax.transAxes, fontsize=6,
            color="white", zorder=7,
            bbox=dict(boxstyle="round,pad=0.2", fc="black", alpha=0.6, ec="gray", lw=0.5),
            verticalalignment='bottom')

    ax.set_xlabel("X", fontsize=8, labelpad=2)
    ax.set_ylabel("Z", fontsize=8, labelpad=2)
    ax.set_title(title, fontsize=9, pad=6, fontweight='medium')
    ax.set_aspect("equal", adjustable="box")
    ax.tick_params(labelsize=7, pad=2)
    ax.xaxis.set_minor_locator(mticker.AutoMinorLocator(2))
    ax.yaxis.set_minor_locator(mticker.AutoMinorLocator(2))
    ax.grid(True, which="major", color="gray", alpha=0.2, lw=0.3, linestyle='-')
    ax.grid(True, which="minor", color="gray", alpha=0.1, lw=0.1, linestyle=':')


def plot_sdf_heatmap(data: dict, out_path: str, shared_norm=None):
    subfaces = data["subfaces"]
    ns = len(subfaces)
    y_slice = data["y_slice"]
    x0, x1, z0, z1 = data["x0"], data["x1"], data["z0"], data["z1"]
    N = data["grid_res"] + 1

    X = np.linspace(x0, x1, N)
    Z = np.linspace(z0, z1, N)
    XX, ZZ = np.meshgrid(X, Z)

    if ns > 0:
        combined = build_combined_grid(data)

    cmap = plt.cm.RdBu_r

    if ns <= 1:
        fig = plt.figure(figsize=(12, 5.5))
        gs = GridSpec(1, 3, width_ratios=[1, 0.05, 1.2], wspace=0.15, hspace=0.1)
        axes = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 2])]
    elif ns == 2:
        fig = plt.figure(figsize=(15, 5))
        gs = GridSpec(1, 4, width_ratios=[1, 1, 0.05, 1.2], wspace=0.15, hspace=0.1)
        axes = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1]), fig.add_subplot(gs[0, 3])]
    elif ns == 3:
        fig = plt.figure(figsize=(12, 9))
        gs = GridSpec(2, 3, width_ratios=[1, 1, 0.05], height_ratios=[1, 1], wspace=0.15, hspace=0.2)
        axes = [
            fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1]),
            fig.add_subplot(gs[1, 0]), fig.add_subplot(gs[1, 1]),
            fig.add_subplot(gs[0:2, 2])
        ]
    else:
        ncols = min(3, int(np.ceil(np.sqrt(ns + 1))))
        nrows = int(np.ceil((ns + 1) / ncols))
        fig, axes = plt.subplots(nrows, ncols, figsize=(min(6 * ncols, 18), 5 * nrows), squeeze=False)
        axes = axes.flatten()

    for s, sf in enumerate(subfaces):
        if s < len(axes) - 1:
            stype = sf["type"]
            name = SURFACE_NAMES.get(stype, f"Type {stype}")
            draw_panel(axes[s], XX, ZZ, sf["grid"],
                       f"Subface {s}: {name}", y_slice, cmap,
                       norm=shared_norm, highlight_sign_flip=(stype == 4))

    if ns > 0 and ns < len(axes):
        draw_panel(axes[ns], XX, ZZ, combined,
                   "Combined SDF\nmax(subfaces) ≈ CSG Intersection",
                   y_slice, cmap, norm=shared_norm, highlight_sign_flip=True)
        for i in range(ns + 1, len(axes)):
            axes[i].set_visible(False)

    fig.suptitle(
        f"SDF Cross-Section Analysis - Y = {y_slice:.6f}\n"
        "Blue = Inside (SDF < 0) | Red = Outside (SDF > 0) | Lime = Surface (SDF = 0)",
        fontsize=11, y=0.98, fontweight='bold'
    )
    fig.text(0.02, 0.02,
             f"Data: {Path(out_path).stem}.txt | Grid: {data['grid_res']}×{data['grid_res']} | "
             f"Subfaces: {ns} | X: [{x0:.3f}, {x1:.3f}] | Z: [{z0:.3f}, {z1:.3f}]",
             fontsize=7, style='italic', alpha=0.7)

    fig.savefig(out_path, dpi=200, bbox_inches="tight", facecolor='white')
    print(f"[SDF_VIZ] 2D Heatmap → {out_path}")
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
    plt.show()


# ══════════════════════════════════════════════════════════════════════
# 5. 三维 SDF 可视化
# ══════════════════════════════════════════════════════════════════════

def sdf_to_rgba(sdf_values: np.ndarray, norm, cmap) -> np.ndarray:
    """将 SDF 值转换为 RGBA uint8 颜色数组。"""
    normed = norm(sdf_values)
    rgba_float = cmap(normed)                    # (N, 4) float [0,1]
    rgba_uint8 = (rgba_float * 255).astype(np.uint8)
    return rgba_uint8


def export_ply_with_colors(vertices: np.ndarray,
                           faces: list,
                           colors: np.ndarray,
                           ply_path: str):
    """
    手动写出带顶点颜色的 PLY 文件（ASCII 格式，无需 trimesh）。
    colors: (N, 4) uint8 RGBA
    """
    n_verts = len(vertices)
    n_faces = len(faces)
    with open(ply_path, 'w', encoding='utf-8') as f:
        # Header
        f.write("ply\n")
        f.write("format ascii 1.0\n")
        f.write(f"element vertex {n_verts}\n")
        f.write("property float x\nproperty float y\nproperty float z\n")
        f.write("property uchar red\nproperty uchar green\nproperty uchar blue\nproperty uchar alpha\n")
        f.write(f"element face {n_faces}\n")
        f.write("property list uchar int vertex_indices\n")
        f.write("end_header\n")
        # Vertices
        for i, v in enumerate(vertices):
            r, g, b, a = colors[i]
            f.write(f"{v[0]:.6f} {v[1]:.6f} {v[2]:.6f} {r} {g} {b} {a}\n")
        # Faces
        for face in faces:
            f.write(f"{len(face)} " + " ".join(map(str, face)) + "\n")
    print(f"[SDF_VIZ] PLY export → {ply_path}")


def plot_3d_sdf_on_mesh(vertices: np.ndarray,
                        faces: list,
                        sdf_values: np.ndarray,
                        norm,
                        cmap,
                        y_slice: float,
                        out_path: str,
                        obj_name: str,
                        subface_label: str):
    """
    用 matplotlib 生成三维 SDF 颜色图预览（散点云 + 三角面片着色）。
    对大型模型自动降采样以保证渲染速度。
    """
    n = len(vertices)
    # 最多绘制 50000 个顶点散点
    MAX_SCATTER = 50_000
    if n > MAX_SCATTER:
        idx = np.random.choice(n, MAX_SCATTER, replace=False)
        verts_draw = vertices[idx]
        sdf_draw = sdf_values[idx]
        print(f"[SDF_VIZ] Large mesh ({n} verts) — sampled {MAX_SCATTER} for 3D scatter preview.")
    else:
        verts_draw = vertices
        sdf_draw = sdf_values

    colors_scatter = cmap(norm(sdf_draw))

    fig = plt.figure(figsize=(16, 7))
    fig.patch.set_facecolor('#0d0d0d')

    # ── 左图：三维散点预览 ────────────────────────────────────────────
    ax3d = fig.add_subplot(1, 2, 1, projection='3d')
    ax3d.set_facecolor('#0d0d0d')
    sc = ax3d.scatter(verts_draw[:, 0], verts_draw[:, 1], verts_draw[:, 2],
                      c=colors_scatter, s=1.2, alpha=0.85, linewidths=0)

    # Y = y_slice 参考平面
    x_range = [vertices[:, 0].min(), vertices[:, 0].max()]
    z_range = [vertices[:, 2].min(), vertices[:, 2].max()]
    px, pz = np.meshgrid(np.linspace(*x_range, 2), np.linspace(*z_range, 2))
    py = np.full_like(px, y_slice)
    ax3d.plot_surface(px, py, pz, alpha=0.07, color='cyan', linewidth=0)

    ax3d.set_xlabel("X", color='white', fontsize=8, labelpad=4)
    ax3d.set_ylabel("Y", color='white', fontsize=8, labelpad=4)
    ax3d.set_zlabel("Z", color='white', fontsize=8, labelpad=4)
    ax3d.tick_params(colors='#888888', labelsize=7)
    for pane in [ax3d.xaxis.pane, ax3d.yaxis.pane, ax3d.zaxis.pane]:
        pane.set_facecolor((0.08, 0.08, 0.08, 0.9))
        pane.set_edgecolor('#333333')
    ax3d.set_title(f"3D SDF on Mesh\n{subface_label}", color='white', fontsize=9, pad=8)

    # Colorbar
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    cbar = plt.colorbar(sm, ax=ax3d, pad=0.08, fraction=0.035, shrink=0.8)
    cbar.set_label("SDF Value", color='white', fontsize=7)
    cbar.ax.yaxis.set_tick_params(color='white', labelsize=6)
    plt.setp(plt.getp(cbar.ax.axes, 'yticklabels'), color='white')

    # ── 右图：SDF 值直方图（三维顶点分布）────────────────────────────
    ax_hist = fig.add_subplot(1, 2, 2)
    ax_hist.set_facecolor('#111111')

    # 按颜色分段绘制直方图
    bins = np.linspace(sdf_values.min(), sdf_values.max(), 80)
    bin_centers = 0.5 * (bins[:-1] + bins[1:])
    hist, _ = np.histogram(sdf_values, bins=bins)
    bar_colors = cmap(norm(bin_centers))

    for i, (h, bc) in enumerate(zip(hist, bar_colors)):
        ax_hist.bar(bin_centers[i], h, width=(bins[1] - bins[0]) * 0.95,
                    color=bc, alpha=0.85, linewidth=0)

    ax_hist.axvline(0, color='lime', linewidth=1.5, linestyle='--', label='SDF = 0 (Surface)', zorder=5)
    ax_hist.axvline(sdf_values.mean(), color='yellow', linewidth=1.0, linestyle=':', alpha=0.7,
                    label=f'Mean = {sdf_values.mean():.3f}', zorder=5)

    # 统计信息
    neg_pct = np.sum(sdf_values < 0) / len(sdf_values) * 100
    pos_pct = 100 - neg_pct
    on_surface_pct = np.sum(np.abs(sdf_values) < 0.05) / len(sdf_values) * 100
    stats_text = (f"Vertices: {n:,}\n"
                  f"Inside  (SDF<0): {neg_pct:.1f}%\n"
                  f"Outside (SDF>0): {pos_pct:.1f}%\n"
                  f"Surface (|SDF|<0.05): {on_surface_pct:.1f}%\n"
                  f"Min: {sdf_values.min():.4f}\n"
                  f"Max: {sdf_values.max():.4f}\n"
                  f"Std: {sdf_values.std():.4f}")
    ax_hist.text(0.97, 0.97, stats_text, transform=ax_hist.transAxes,
                 fontsize=7.5, color='white', verticalalignment='top', horizontalalignment='right',
                 bbox=dict(boxstyle='round,pad=0.4', fc='#1a1a1a', ec='#555555', lw=0.8))

    ax_hist.set_xlabel("SDF Value", color='white', fontsize=9)
    ax_hist.set_ylabel("Vertex Count", color='white', fontsize=9)
    ax_hist.set_title("SDF Distribution on Mesh Vertices", color='white', fontsize=9)
    ax_hist.tick_params(colors='#888888', labelsize=8)
    ax_hist.spines[:].set_edgecolor('#444444')
    ax_hist.legend(fontsize=8, facecolor='#1a1a1a', edgecolor='#555555',
                   labelcolor='white', loc='upper left')
    ax_hist.grid(True, axis='y', color='#333333', linewidth=0.4, alpha=0.6)

    fig.suptitle(
        f"3D SDF Visualization  |  Model: {Path(obj_name).name}  |  {subface_label}\n"
        "Blue = Inside (SDF < 0)  |  Red = Outside (SDF > 0)  |  White/Green = Surface (SDF ≈ 0)",
        color='white', fontsize=10, fontweight='bold', y=0.99
    )
    fig.text(0.5, 0.01,
             f"Y-slice reference: {y_slice:.4f}  |  SDF interpolated from 2D XZ grid onto 3D vertices",
             color='#888888', fontsize=7, ha='center', style='italic')

    plt.tight_layout(rect=[0, 0.03, 1, 0.96])
    fig.savefig(out_path, dpi=200, bbox_inches='tight', facecolor='#0d0d0d')
    print(f"[SDF_VIZ] 3D Preview → {out_path}")
    plt.show()


def visualize_3d(data: dict, obj_path: str, out_dir: Path,
                 subface_idx: int = -1, shared_norm=None):
    """三维可视化主流程。"""
    print(f"[SDF_VIZ] Loading OBJ: {obj_path}")
    try:
        vertices, faces = load_obj_vertices(obj_path)
    except Exception as e:
        print(f"[ERROR] Failed to load OBJ: {e}")
        return

    if len(vertices) == 0:
        print("[ERROR] OBJ file contains no vertices.")
        return

    print(f"[SDF_VIZ] OBJ: {len(vertices):,} vertices, {len(faces):,} faces")

    # 计算顶点 SDF
    print(f"[SDF_VIZ] Computing SDF at vertices (subface_idx={subface_idx}) ...")
    sdf_values = compute_vertex_sdf(vertices, data, subface_idx)
    print(f"[SDF_VIZ] SDF range on mesh: [{sdf_values.min():.4f}, {sdf_values.max():.4f}]")

    # 建立共享色彩归一化（若尚未建立）
    if shared_norm is None:
        shared_norm, _ = build_shared_norm(data, sdf_values)

    cmap = plt.cm.RdBu_r

    # 确定 subface 标签
    if subface_idx == -1:
        subface_label = "Combined SDF (max of all subfaces)"
    else:
        stype = data['subfaces'][subface_idx]['type']
        subface_label = f"Subface {subface_idx}: {SURFACE_NAMES.get(stype, f'Type {stype}')}"

    # ── 导出带颜色的 PLY ─────────────────────────────────────────────
    colors_rgba = sdf_to_rgba(sdf_values, shared_norm, cmap)
    stem = Path(obj_path).stem
    sf_tag = "combined" if subface_idx == -1 else f"sf{subface_idx}"
    ply_out = out_dir / f"{stem}_sdf_{sf_tag}.ply"
    export_ply_with_colors(vertices, faces, colors_rgba, str(ply_out))

    # ── 生成 3D 预览图 ────────────────────────────────────────────────
    preview_out = out_dir / f"{stem}_sdf_{sf_tag}_3d.png"
    plot_3d_sdf_on_mesh(
        vertices, faces, sdf_values,
        shared_norm, cmap,
        data['y_slice'],
        str(preview_out),
        obj_path, subface_label
    )


# ══════════════════════════════════════════════════════════════════════
# 6. CLI 入口
# ══════════════════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(
        description="SDF Visualizer: 2D heatmap + 3D mesh coloring",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("sdf_file", help="Path to sdf_slice.txt generated by dump_sdf_slice()")
    parser.add_argument("--obj",      metavar="FILE",  default=None,
                        help="Path to OBJ model for 3D SDF visualization")
    parser.add_argument("--subface",  metavar="INT",   type=int, default=-1,
                        help="Subface index for 3D coloring (-1 = combined SDF, default: -1)")
    parser.add_argument("--no-2d",    action="store_true",
                        help="Skip 2D heatmap, only generate 3D visualization")
    parser.add_argument("--out-dir",  metavar="DIR",   default=None,
                        help="Output directory (default: same as sdf_file)")
    args = parser.parse_args()

    # 检查输入文件
    if not Path(args.sdf_file).exists():
        print(f"[ERROR] File not found: {args.sdf_file}")
        sys.exit(1)

    # 确定输出目录
    out_dir = Path(args.out_dir) if args.out_dir else Path(args.sdf_file).parent
    out_dir.mkdir(parents=True, exist_ok=True)
    stem = Path(args.sdf_file).stem

    # 读取 SDF 数据
    print(f"[SDF_VIZ] Loading {args.sdf_file} ...")
    data = load_sdf_data(args.sdf_file)
    print(f"[SDF_VIZ] Grid: {data['grid_res']}×{data['grid_res']} | "
          f"Y-slice: {data['y_slice']:.6f} | "
          f"Subfaces: {len(data['subfaces'])}")

    # 预加载 OBJ 顶点（若有），用于建立全局共享归一化
    vertex_sdf_for_norm = None
    if args.obj and Path(args.obj).exists():
        try:
            verts_tmp, _ = load_obj_vertices(args.obj)
            if len(verts_tmp) > 0:
                vertex_sdf_for_norm = compute_vertex_sdf(verts_tmp, data, args.subface)
        except Exception:
            pass

    # 建立共享色彩归一化
    shared_norm, shared_vmax = build_shared_norm(data, vertex_sdf_for_norm)
    print(f"[SDF_VIZ] Shared color norm: vmax = ±{shared_vmax:.4f}")

    # ── 二维热力图 ────────────────────────────────────────────────────
    if not args.no_2d:
        out_2d = str(out_dir / f"{stem}_heatmap.png")
        try:
            plot_sdf_heatmap(data, out_2d, shared_norm=shared_norm)
            if Path(out_2d).exists():
                print(f"[SDF_VIZ] 2D output: {out_2d} ({Path(out_2d).stat().st_size/1024:.1f} KB)")
        except Exception as e:
            print(f"[ERROR] 2D visualization failed: {e}")
            import traceback; traceback.print_exc()

    # ── 三维可视化 ────────────────────────────────────────────────────
    if args.obj:
        if not Path(args.obj).exists():
            print(f"[ERROR] OBJ file not found: {args.obj}")
        else:
            try:
                visualize_3d(data, args.obj, out_dir, args.subface, shared_norm)
            except Exception as e:
                print(f"[ERROR] 3D visualization failed: {e}")
                import traceback; traceback.print_exc()
    else:
        if args.no_2d:
            print("[WARN] --no-2d specified but --obj not provided. Nothing to do.")


if __name__ == "__main__":
    # 兼容旧调用方式：python visualize_sdf.py sdf_slice.txt（无 --obj 参数）
    # 此时等同于仅二维模式
    if len(sys.argv) >= 2 and not sys.argv[1].startswith('-'):
        # 如果第一个参数不是选项，视为 sdf_file，其余透传给 argparse
        pass
    main()