Wonder3D with TRELLIS Integration

将 Wonder3D 多视图生成与 TRELLIS 3D 重建进行深度融合的模块化框架。

项目概述

本项目旨在利用 Wonder3D 生成的多视图 2D 特征来增强 TRELLIS 的 3D 潜在表示 (SLAT)，实现更精确的 3D 重建。

核心思想

输入图像 → TRELLIS (生成 3D SLAT) → 3D Token 提取
                                        ↓
                                   3D Adapter (MLP_3D)
                                        ↓
输入图像 → Wonder3D (生成多视图) → 2D 特征采样 → 2D Adapter (MLP_2D)
                                                    ↓
                                              Fusion-Update
                                                    ↓
                                            更新后的 3D Token

已实现模块

1. get_3d_token.py - 3D Token 提取器

从 TRELLIS 的 SLAT (Structured Latent) 中提取 3D Token。

主要类:

• Token3D: 3D Token 数据结构，包含 coords (体素坐标) 和 feats (特征向量)
• SLAT3DTokenExtractor: 从 SparseTensor 提取 Token 的工具类

使用示例:

from wonder3d_with_trellis import SLAT3DTokenExtractor, Token3D

extractor = SLAT3DTokenExtractor()
token = extractor.extract(slat)  # slat 是 TRELLIS 的 SparseTensor

print(f"Token 数量: {token.num_tokens}")
print(f"特征维度: {token.feat_dim}")
print(f"体素分辨率: {token.resolution}")

2. coords_to_world.py - 坐标转换器

将 TRELLIS SLAT 中的体素坐标转换为物理世界坐标。

转换公式:

world_pos = voxel_coords / resolution - 0.5 + offset

主要类:

• VoxelToWorldConverter: 体素坐标到物理坐标的转换器
• PositionalEncoding3D: 3D sin/cos 位置编码

使用示例:

from wonder3d_with_trellis import VoxelToWorldConverter

converter = VoxelToWorldConverter(resolution=64)
world_coords = converter.convert(voxel_coords)  # (N, 3) -> (N, 3)

# 带位置编码
from wonder3d_with_trellis.coords_to_world import PositionalEncoding3D
pe = PositionalEncoding3D(num_frequencies=6)
encoded = pe(world_coords)  # (N, 3) -> (N, 39)

3. adapter_3d.py - 3D Adapter (MLP_3D)

将 TRELLIS 3D Token 的特征投影到共享空间。

架构:

Input: feat_3d (N, C3) + pos_enc(coords) (N, P)
    → Linear(C3+P, d)
    → LayerNorm
    → GELU
    → Linear(d, d)
    → (Optional) Residual
Output: h3 (N, d)

主要类:

• Adapter3D: 基础 3D Adapter
• Adapter3DWithAttention: 带自注意力的 3D Adapter

使用示例:

from wonder3d_with_trellis import Adapter3D

adapter = Adapter3D(
    in_channels=64,      # TRELLIS SLAT 特征维度
    hidden_dim=256,      # 共享空间维度
    num_pos_frequencies=6
)

h3 = adapter(feat_3d, coords, resolution=64)  # (N, 64), (N, 3) -> (N, 256)

4. adapter_2d.py - 2D Adapter (MLP_2D)

将聚合后的多视图 2D 特征投影到与 3D 特征相同的共享空间。

架构:

Input: feat_2d_agg (N, C2) + view_embedding (可选)
    → Linear(C2, d)
    → LayerNorm
    → GELU
    → Linear(d, d)
Output: h2 (N, d)

主要类:

• ConvAdapter2D: VAE latent 升维卷积适配器
• Adapter2D: 基础 2D Adapter
• MultiViewFeatureAggregator: 多视图特征聚合器
• Adapter2DPipeline: 完整的 2D 处理流水线

使用示例:

from wonder3d_with_trellis import Adapter2D
from wonder3d_with_trellis.adapter_2d import ConvAdapter2D, MultiViewFeatureAggregator

# VAE latent 升维
conv_adapter = ConvAdapter2D(in_channels=4, out_channels=64)
enhanced = conv_adapter(vae_latent)  # (B, 4, H, W) -> (B, 64, H, W)

# 多视图聚合
aggregator = MultiViewFeatureAggregator(feat_dim=64, aggregation='attention')
aggregated = aggregator(multi_view_feats, visibility_mask)  # (N, V, C) -> (N, C)

# 投影到共享空间
adapter = Adapter2D(in_channels=64, hidden_dim=256)
h2 = adapter(aggregated)  # (N, 64) -> (N, 256)

5. fusion_update.py - Fusion-Update (MLP_FUSE)

门控残差更新模块，融合 3D 和 2D 特征，并映射回 TRELLIS 3D Token 潜空间。

架构:

1. delta = Linear([h3, h2], d) → GELU → Linear(d, d)
2. gate = sigmoid(Linear([h3, h2, geom], 1 or d))
3. h3' = h3 + gate * delta
4. feat_3d' = Linear(h3', C3)

主要类:

• FusionUpdate: 基础融合更新模块
• GeometryFeatureExtractor: 几何特征提取器 (法线、视图方向等)
• FusionUpdateWithGeometry: 带几何特征的融合模块
• FullFusionPipeline: 完整的融合流水线

使用示例:

from wonder3d_with_trellis import FusionUpdate
from wonder3d_with_trellis.fusion_update import FusionUpdateWithGeometry

# 基础融合
fusion = FusionUpdate(
    hidden_dim=256,
    out_channels=64,  # TRELLIS SLAT 特征维度
    geom_dim=16,
    gate_type='vector'
)

feat_3d_updated, gate = fusion(h3, h2, geom)  # (N, 256), (N, 256), (N, 16) -> (N, 64), (N, 256)

# 带几何特征的融合
fusion_geom = FusionUpdateWithGeometry(hidden_dim=256, out_channels=64)
feat_3d_updated, gate = fusion_geom(h3, h2, normals, view_dirs, depths)

6. token_vertex_correspondence.py - Token-Vertex 对应关系

构建 3D Token 与 Mesh 顶点之间的对应关系 (W_v→t)，是将 2D 多视图信息聚合回 3D Token 的关键桥梁。

核心思想:

实现路径: 先用 mesh 顶点做精确对齐，再把信息聚合回 token

vertex (连续、好对齐) ←→ token (离散、要用于 3D latent 融合)

三种对应关系方法:

1. 最近邻 (nearest): 对每个 vertex 找最近的 token，权重为 1
2. KNN 距离加权 (knn): 对每个 vertex 找 K 个最近 token，用 softmax(-d/τ) 做权重 (推荐)
3. 三线性插值 (trilinear): 将 vertex 坐标反算回体素坐标，用三线性插值权重

主要类:

• CorrespondenceWeights: 稀疏权重矩阵数据结构
• TokenVertexCorrespondence: 对应关系构建器
• TokenMeshAligner: 坐标对齐器
• TokenVertexCorrespondenceBuilder: 完整流程封装

使用示例:

from wonder3d_with_trellis import (
    TokenVertexCorrespondenceBuilder,
    build_token_vertex_correspondence
)

# 方法 1: 使用便捷函数
result = build_token_vertex_correspondence(
    token=token, mesh=mesh,
    method='knn', k=8, temperature=0.1,
    mesh_center=center, mesh_scale=scale, mesh_rotation=rotation
)
weights = result['weights']

# 方法 2: 使用构建器
builder = TokenVertexCorrespondenceBuilder(method='knn', k=8)
builder.set_mesh_transform(center, scale, rotation)
result = builder.build(token, mesh)

# 特征聚合
vertex_features = builder.aggregate_to_vertices(token_features, weights)
token_features_updated = builder.aggregate_to_tokens(vertex_features, weights)

详细文档见 docs/token_vertex_correspondence.md

7. render_mesh_wonder3d_views_v2.py - 渲染与对齐工具

TRELLIS mesh 到 Wonder3D 坐标系的转换、六视图渲染、Latent 对齐等功能。

主要功能:

• 将 TRELLIS mesh 渲染为 Wonder3D 格式六视图
• TRELLIS-Wonder3D 几何对齐验证
• Latent 空间对齐和特征采样
• BBox 对齐工具

主要类:

• OrthoMeshRenderer: 正交投影 Mesh 渲染器
• BBoxAligner: BBox 对齐工具
• LatentProjector: 3D 点到 Latent 特征投影器
• Wonder3DLatentExtractor: Wonder3D Latent 提取器
• MultiViewAggregator: 多视图特征聚合器
• AlignmentVerifier: 对齐验证器

命令行使用:

# 渲染六视图
python render_mesh_wonder3d_views_v2.py --mode render --image input.jpg --output ./output

# 对齐验证
python render_mesh_wonder3d_views_v2.py --mode verify --image input.jpg --output ./verify_output

# Latent 对齐
python render_mesh_wonder3d_views_v2.py --mode latent --image input.jpg --output ./latent_output

模块依赖关系

get_3d_token.py
       ↓
coords_to_world.py ←── adapter_3d.py
       ↓                     ↓
token_vertex_correspondence.py
       ↓                     ↓
       └──────→ fusion_update.py ←── adapter_2d.py
                      ↓
          render_mesh_wonder3d_views_v2.py

数据流

输入图像
    ├─→ TRELLIS Pipeline
    │       ├─→ SLAT (SparseTensor)
    │       │       └─→ get_3d_token.py → Token3D
    │       │               └─→ coords_to_world.py → 物理坐标
    │       │                       └─→ adapter_3d.py → h3
    │       │
    │       └─→ Mesh (MeshExtractResult)
    │               └─→ token_vertex_correspondence.py
    │                       └─→ W_v→t (对应关系)
    │
    └─→ Wonder3D Pipeline
            └─→ 多视图 Latent
                    └─→ render_mesh_wonder3d_views_v2.py
                            └─→ 顶点 2D 特征采样
                                    └─→ W_v→t 聚合 → Token 2D 特征
                                            └─→ adapter_2d.py → h2
                                                    └─→ fusion_update.py → feat_3d'

配置参数

默认路径

• TRELLIS 路径: /home/lyc/TRELLIS
• Wonder3D 路径: /home/lyc/Wonder3D
• TRELLIS 模型: /home/lyc/models/trellis-image-large
• Wonder3D 相机位姿: /home/lyc/Wonder3D/instant-nsr-pl/datasets/fixed_poses

默认参数

• 图像分辨率: 256
• 正交投影缩放: 1.35
• VAE 缩放因子: 8
• Latent 尺寸: 32
• SLAT 分辨率: 64
• 共享空间维度 d: 256

架构设计说明

3D-Adapter (MLP_3D)

• 输入: feat_3d (C3 维) + pos_enc(coords) (3D sin/cos 位置编码)
• 输出: h3 (d 维, 如 d=256)
• 结构: Linear(C3+P → d) → LayerNorm → GELU → Linear(d → d)
• 可选残差连接

2D-Adapter (MLP_2D)

• 输入: feat_2d_agg (C2 维, 来自 Wonder3D latent 采样+聚合) + view/几何辅助量
• 输出: h2 (d 维)
• 结构: Linear → LN → GELU → Linear
• 注意: VAE latent 通道较小，建议先用 ConvAdapter 升维到 C2=64/128

Fusion-Update (MLP_FUSE)

• 输入: h3, h2, geom (几何辅助信息)
• 输出: feat_3d' (更新后的 3D Token 特征)
• 计算过程:
  1. delta = Linear([h3, h2] → d) → GELU → Linear(d → d)
  2. gate = sigmoid(Linear([h3, h2, geom] → 1 or d))
  3. h3' = h3 + gate * delta
  4. feat_3d' = Linear(h3' → C3)
• geom 可包含: 视图方向与法线夹角、可见性、深度一致性等

版本信息

• 版本: 0.1.0
• 创建日期: 2025-01

TODO

• 完整的训练流水线
• 端到端的推理脚本
• 更多聚合策略 (如 Transformer-based)
• 多尺度特征融合
• 可学习的几何特征编码