Temporary repository used to save branch code
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 
 
 

649 lines
28 KiB

#!/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()