| """fix3 tokenizer(已融合 numba 加速实现, 原 *_fast.py 并入本文件)。 |
| |
| encode_mesh / decode_tokens 走 numpy 向量化 + numba 枚举, 比纯 python 版快约 5x。 |
| quantized_mesh 与 decode 输出和纯 python 版完全一致; token 的 face 块集合一致, |
| 仅少数平局块顺序不同(本实现补了 (f2.y,f2.x) 决定性平局键, 比原版更可复现)。 |
| |
| 原文件头: |
| |
| fix3: 连接矩阵改为「octant 绝对原位 8x8 + 因子分解(按行)」编码，保留 F。 |
| |
| 相对 fix2 的核心变化 |
| -------------------- |
| 1. 不再按 F 排序、不再受 max_children 限制。每个连接矩阵(intra / inter)统一是 |
| 8x8 的「octant 绝对原位」矩阵：行/列 = octant 0..7 的固定语义。 |
| → 天然 mc=8 无损，"超标 / 合并 / 退化分摊" 这套概念全部消失。 |
| 2. 8x8 矩阵不再编码成单个 token(那是 2^64 词表)，而是按行拆成 8 个「行子码」， |
| 每个行子码 = 该行 8 个 bit ∈ [0, 256)。intra / inter 各 8 个子码。 |
| 3. **逻辑序列长度不变**：每个 face 仍是 18 个「位置」(9 coord + 3 F + 3 intra + 3 inter)， |
| 只是 intra / inter 这 6 个位置每个携带 8 个子码(给因子分解头：同一位置 8 个并行 256-way 头)。 |
| coord / F 仍是单子码位置。tokenizer.encode 返回 List[List[int]]：每个内层 list = 一个位置的子码。 |
| 4. 保留 F(标识 active octant)。重建时用 F mask 掉非 active 行列(M &= F⊗F)，消除歧义。 |
| |
| 词表(极小)：pad(1) + coord(resolution) + F(256) + intra_row(256) + inter_row(256)。 |
| """ |
| import numpy as np |
| import trimesh |
| from tqdm import tqdm |
| import os |
| import sys |
| import glob |
| sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| from pathlib import Path |
| from utils.data_utils import undiscretize, discretize, normalize, clean_mesh |
| from utils.mesh_metric import compute_mesh_metrics |
| from modules.fast_octree_level import _triangles_csr |
| import itertools |
| from dataclasses import dataclass |
|
|
| NUM_OCTANTS = 8 |
|
|
| |
| _TRI_INDICES = np.array(list(itertools.combinations(range(24), 3))) |
| _IDX_I = _TRI_INDICES[:, 0] |
| _IDX_J = _TRI_INDICES[:, 1] |
| _IDX_K = _TRI_INDICES[:, 2] |
|
|
|
|
| |
| |
| |
| class DynamicMeshTokenizer: |
| """ |
| encode 输出格式(fine)：List[List[int]]，每个内层 list 是一个「逻辑位置」的子码。 |
| 每 face 18 个位置： [v0x][v0y][v0z][v1x][v1y][v1z][v2x][v2y][v2z] (9 个 coord, 各 1 子码) |
| [F0][F1][F2] (3 个 F, 各 1 子码) |
| [M00 的 8 行][M11..][M22..] (3 个 intra, 各 8 子码) |
| [M01 的 8 行][M02..][M12..] (3 个 inter, 各 8 子码) |
| coord / F 位置长度 1；intra / inter 位置长度 8(行子码)。逻辑序列长度 = 18/face，与 fix2 相同。 |
| """ |
| def __init__(self, resolution, max_children=NUM_OCTANTS): |
| self.resolution = resolution |
| self.max_children = NUM_OCTANTS |
|
|
| self.pad_id = 0 |
| self.coord_offset = 1 |
| self.f_offset = self.coord_offset + self.resolution |
| self.row_vocab = 1 << NUM_OCTANTS |
| self.f_vocab = 1 << NUM_OCTANTS |
| self.intra_row_offset = self.f_offset + self.f_vocab |
| self.inter_row_offset = self.intra_row_offset + self.row_vocab |
| self.vocab_size = self.inter_row_offset + self.row_vocab |
|
|
| |
| def _row_to_int(self, bits): |
| v = 0 |
| for i, b in enumerate(bits): |
| if b: |
| v |= (1 << i) |
| return v |
|
|
| def _int_to_row(self, val): |
| return [(val >> i) & 1 for i in range(NUM_OCTANTS)] |
|
|
| def _mat_to_rows(self, M, offset): |
| return [self._row_to_int(M[r]) + offset for r in range(NUM_OCTANTS)] |
|
|
| def _rows_to_mat(self, rows, offset): |
| return [self._int_to_row(rows[r] - offset) for r in range(NUM_OCTANTS)] |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| TOKENS_PER_FACE = 9 + 3 + 6 * NUM_OCTANTS |
| REFINE_POS_PER_FACE = 9 |
|
|
| def encode(self, faces_data, use_hourglass=True): |
| seq = [] |
| for face in faces_data: |
| v0, v1, v2, F0, F1, F2, M00, M11, M22, M01, M02, M12 = face |
| for v in (v0, v1, v2): |
| seq.extend([v[0] + self.coord_offset, v[1] + self.coord_offset, v[2] + self.coord_offset]) |
| for F in (F0, F1, F2): |
| seq.append(self._row_to_int(F) + self.f_offset) |
| for M in (M00, M11, M22): |
| seq.extend(self._mat_to_rows(M, self.intra_row_offset)) |
| for M in (M01, M02, M12): |
| seq.extend(self._mat_to_rows(M, self.inter_row_offset)) |
| return seq |
|
|
| def decode(self, seq): |
| faces_data = [] |
| B = self.TOKENS_PER_FACE |
| clean_seq = [t for t in seq if t != self.pad_id] |
| co = self.coord_offset |
| for i in range(0, len(clean_seq) - B + 1, B): |
| ch = clean_seq[i:i + B] |
| cv = [ch[k] - co for k in range(9)] |
| v0 = (cv[0], cv[1], cv[2]); v1 = (cv[3], cv[4], cv[5]); v2 = (cv[6], cv[7], cv[8]) |
| F0 = self._int_to_row(ch[9] - self.f_offset) |
| F1 = self._int_to_row(ch[10] - self.f_offset) |
| F2 = self._int_to_row(ch[11] - self.f_offset) |
| M00 = self._rows_to_mat(ch[12:20], self.intra_row_offset) |
| M11 = self._rows_to_mat(ch[20:28], self.intra_row_offset) |
| M22 = self._rows_to_mat(ch[28:36], self.intra_row_offset) |
| M01 = self._rows_to_mat(ch[36:44], self.inter_row_offset) |
| M02 = self._rows_to_mat(ch[44:52], self.inter_row_offset) |
| M12 = self._rows_to_mat(ch[52:60], self.inter_row_offset) |
| faces_data.append((v0, v1, v2, F0, F1, F2, M00, M11, M22, M01, M02, M12)) |
| return faces_data |
|
|
| |
| |
| |
| def encode_raw(self, raw_faces_data, use_hourglass=True): |
| seq = [] |
| for face in raw_faces_data: |
| v0, v1, v2 = face |
| for v in (v0, v1, v2): |
| seq.append([v[0] + self.coord_offset]) |
| seq.append([v[1] + self.coord_offset]) |
| seq.append([v[2] + self.coord_offset]) |
| return seq |
|
|
| def decode_raw(self, seq): |
| faces_data = [] |
| clean_seq = [p for p in seq if not (len(p) == 1 and p[0] == self.pad_id)] |
| for i in range(0, len(clean_seq) - 9 + 1, 9): |
| chunk = clean_seq[i:i + 9] |
| if len(chunk) < 9: |
| break |
| cv = [chunk[k][0] - self.coord_offset for k in range(9)] |
| faces_data.append(((cv[0], cv[1], cv[2]), (cv[3], cv[4], cv[5]), (cv[6], cv[7], cv[8]))) |
| return faces_data |
|
|
|
|
| |
| |
| |
| class OctreeFaceTokenBuilder: |
| def __init__(self, tree, max_children=NUM_OCTANTS): |
| self.tree = tree |
| self.max_children = NUM_OCTANTS |
|
|
| def get_level_info(self, depth): |
| nodes_at_depth = {} |
|
|
| def collect_nodes(current, target_depth): |
| if current.depth == target_depth: |
| nodes_at_depth[current.coord] = current |
| return |
| for child in current.children: |
| if child: |
| collect_nodes(child, target_depth) |
|
|
| collect_nodes(self.tree.root, depth) |
|
|
| coords = sorted(nodes_at_depth.keys(), key=lambda x: (x[2], x[1], x[0])) |
| coord_to_idx = {c: i for i, c in enumerate(coords)} |
|
|
| adj = {i: set() for i in range(len(coords))} |
| for i, c in enumerate(coords): |
| node = nodes_at_depth[c] |
| all_neighbors = node.intra_neighbors | node.inter_neighbors |
| for neighbor_coord in all_neighbors: |
| if neighbor_coord in coord_to_idx: |
| j = coord_to_idx[neighbor_coord] |
| if i < j: |
| adj[i].add(j) |
|
|
| raw_faces_set = set() |
| covered_edges = set() |
|
|
| |
| for u in range(len(coords)): |
| for v in adj[u]: |
| common_neighbors = adj[u] & adj[v] |
| for w in common_neighbors: |
| face_coords = tuple(sorted([coords[u], coords[v], coords[w]])) |
| raw_faces_set.add(face_coords) |
| covered_edges.add(tuple(sorted([face_coords[0], face_coords[1]]))) |
| covered_edges.add(tuple(sorted([face_coords[1], face_coords[2]]))) |
| covered_edges.add(tuple(sorted([face_coords[0], face_coords[2]]))) |
|
|
| |
| for u in range(len(coords)): |
| for v in adj[u]: |
| edge = tuple(sorted([coords[u], coords[v]])) |
| if edge not in covered_edges: |
| raw_faces_set.add((edge[0], edge[1], edge[1])) |
|
|
| |
| for u in range(len(coords)): |
| has_edges = len(adj[u]) > 0 or any(u in adj[i] for i in range(u)) |
| if not has_edges: |
| raw_faces_set.add((coords[u], coords[u], coords[u])) |
|
|
| raw_faces_at_depth = list(raw_faces_set) |
| raw_faces_at_depth.sort(key=lambda f: (f[0][2], f[0][1], f[0][0], |
| f[1][2], f[1][1], f[1][0], |
| f[2][2])) |
| return raw_faces_at_depth, nodes_at_depth |
|
|
| def _active_octants(self, coord, n, nodes_map): |
| """该节点哪些 octant 子节点存在(occupancy)。 |
| 用 children 占据而非"有连接"标记 —— 因为「有连接 ⟹ 必有子节点」，occupancy 是 |
| "有连接"的超集，且正好等于 raw 的 get_level_info 收集的 depth+1 节点集； |
| 这样「有几何但无边」的孤立子节点也会被 F 标记，fine 重构能补回孤立点，与 raw 完全一致。 |
| M(连接矩阵)不变，故三角形/边重构不受影响。""" |
| return set(i for i in range(NUM_OCTANTS) if n.children[i] is not None) |
|
|
| def encode_level_tokens(self, depth): |
| raw_faces, nodes_map = self.get_level_info(depth) |
| empty = np.zeros((NUM_OCTANTS, NUM_OCTANTS), dtype=np.int8) |
|
|
| |
| node_F = {} |
| for coord, n in nodes_map.items(): |
| active = self._active_octants(coord, n, nodes_map) |
| node_F[coord] = [1 if i in active else 0 for i in range(NUM_OCTANTS)] |
|
|
| intra_cache = {} |
|
|
| def intra8(coord): |
| if coord not in intra_cache: |
| M = (nodes_map[coord].intra_matrix != 0).astype(np.int8) |
| intra_cache[coord] = M.tolist() |
| return intra_cache[coord] |
|
|
| def inter8(a, b): |
| if a == b: |
| return intra8(a) |
| na = nodes_map[a] |
| M = na.inter_matrices.get(b) |
| if M is None: |
| nb = nodes_map[b] |
| Mt = nb.inter_matrices.get(a) |
| M = Mt.T if Mt is not None else empty |
| return (np.asarray(M) != 0).astype(np.int8).tolist() |
|
|
| raw_tokens = [] |
| output_tokens = [] |
| for face in raw_faces: |
| v0_c, v1_c, v2_c = face |
| raw_tokens.append([v0_c, v1_c, v2_c]) |
|
|
| F0, F1, F2 = node_F[v0_c], node_F[v1_c], node_F[v2_c] |
| M00, M11, M22 = intra8(v0_c), intra8(v1_c), intra8(v2_c) |
| M01 = inter8(v0_c, v1_c) |
| M02 = inter8(v0_c, v2_c) |
| M12 = inter8(v1_c, v2_c) |
|
|
| output_tokens.append((v0_c, v1_c, v2_c, F0, F1, F2, M00, M11, M22, M01, M02, M12)) |
| return raw_tokens, output_tokens |
|
|
|
|
| class OctreeNode: |
| def __init__(self, x, y, z, depth): |
| self.coord = (x, y, z) |
| self.depth = depth |
| self.children = [None] * 8 |
| self.intra_neighbors = set() |
| self.inter_neighbors = set() |
| self.intra_matrix = np.zeros((8, 8), dtype=np.int8) |
| self.inter_matrices = {} |
|
|
|
|
| class MatrixOctree: |
| def __init__(self, max_depth=7): |
| self.max_depth = max_depth |
| self.root = None |
|
|
| def _get_parent_coord(self, coord): |
| return (coord[0] >> 1, coord[1] >> 1, coord[2] >> 1) |
|
|
| def _get_local_index(self, coord): |
| return ((coord[2] & 1) << 2) | ((coord[1] & 1) << 1) | (coord[0] & 1) |
|
|
| def build(self, vertices, faces): |
| resolution = 2 ** self.max_depth |
| quantized = discretize(vertices, num_discrete=resolution) |
|
|
| leaf_nodes_map = {} |
| for i in range(len(vertices)): |
| coord = tuple(quantized[i]) |
| if coord not in leaf_nodes_map: |
| leaf_nodes_map[coord] = OctreeNode(coord[0], coord[1], coord[2], self.max_depth) |
|
|
| for face in faces: |
| coords = [tuple(quantized[i]) for i in face] |
| for i in range(3): |
| u, v = coords[i], coords[(i + 1) % 3] |
| if u == v: |
| continue |
| parent_u = self._get_parent_coord(u) |
| parent_v = self._get_parent_coord(v) |
| node_u = leaf_nodes_map[u] |
| node_v = leaf_nodes_map[v] |
| if parent_u == parent_v: |
| node_u.intra_neighbors.add(v) |
| node_v.intra_neighbors.add(u) |
| else: |
| node_u.inter_neighbors.add(v) |
| node_v.inter_neighbors.add(u) |
|
|
| current_layer_nodes = list(leaf_nodes_map.values()) |
|
|
| for d in range(self.max_depth, 0, -1): |
| parent_layer_map = {} |
| for child in current_layer_nodes: |
| parent_coord = self._get_parent_coord(child.coord) |
| if parent_coord not in parent_layer_map: |
| parent_layer_map[parent_coord] = OctreeNode(*parent_coord, d - 1) |
| parent_node = parent_layer_map[parent_coord] |
| idx = self._get_local_index(child.coord) |
| parent_node.children[idx] = child |
|
|
| for child in current_layer_nodes: |
| parent_coord = self._get_parent_coord(child.coord) |
| parent_node = parent_layer_map[parent_coord] |
| child_idx = self._get_local_index(child.coord) |
|
|
| for sibling_coord in child.intra_neighbors: |
| sibling_idx = self._get_local_index(sibling_coord) |
| parent_node.intra_matrix[child_idx, sibling_idx] = 1 |
|
|
| for neighbor_coord in child.inter_neighbors: |
| neighbor_parent_coord = self._get_parent_coord(neighbor_coord) |
| if neighbor_parent_coord not in parent_node.inter_matrices: |
| parent_node.inter_matrices[neighbor_parent_coord] = np.zeros((8, 8), dtype=np.int8) |
| neighbor_child_idx = self._get_local_index(neighbor_coord) |
| parent_node.inter_matrices[neighbor_parent_coord][child_idx, neighbor_child_idx] = 1 |
|
|
| grandparent_p = self._get_parent_coord(parent_coord) |
| grandparent_n = self._get_parent_coord(neighbor_parent_coord) |
| if grandparent_p == grandparent_n: |
| parent_node.intra_neighbors.add(neighbor_parent_coord) |
| else: |
| parent_node.inter_neighbors.add(neighbor_parent_coord) |
|
|
| current_layer_nodes = list(parent_layer_map.values()) |
| if current_layer_nodes: |
| self.root = current_layer_nodes[0] |
|
|
|
|
| def reconstruct_next_level_optimized(output_tokens, resolution, max_children=NUM_OCTANTS): |
| """octant 绝对原位重建：M 已是 8x8、行列=octant；用 F mask 掉非 active 行列。""" |
| indices = np.arange(8) |
| offsets = np.stack([indices & 1, (indices >> 1) & 1, (indices >> 2) & 1], axis=1).astype(np.int32) |
|
|
| tri_blocks = [] |
| edge_keys = set() |
| tri_edge_keys = set() |
| active_keys = set() |
| key_to_coord = {} |
|
|
| for f_token in output_tokens: |
| v0_c, v1_c, v2_c, F0, F1, F2, M00, M11, M22, M01, M02, M12 = f_token |
| fa = [np.asarray(F0, dtype=bool), np.asarray(F1, dtype=bool), np.asarray(F2, dtype=bool)] |
|
|
| def intra(F, M): |
| return np.asarray(M, dtype=bool) & np.outer(F, F) |
|
|
| def inter(Fa, Fb, M): |
| return np.asarray(M, dtype=bool) & np.outer(Fa, Fb) |
|
|
| in0 = intra(fa[0], M00); in1 = intra(fa[1], M11); in2 = intra(fa[2], M22) |
| m01 = inter(fa[0], fa[1], M01); m02 = inter(fa[0], fa[2], M02); m12 = inter(fa[1], fa[2], M12) |
|
|
| A = np.zeros((24, 24), dtype=bool) |
| A[0:8, 0:8] = in0 | in0.T |
| A[8:16, 8:16] = in1 | in1.T |
| A[16:24, 16:24] = in2 | in2.T |
| A[0:8, 8:16] = m01; A[8:16, 0:8] = m01.T |
| A[0:8, 16:24] = m02; A[16:24, 0:8] = m02.T |
| A[8:16, 16:24] = m12; A[16:24, 8:16] = m12.T |
|
|
| |
| v_parents_arr = np.stack([np.array(v0_c[:3], dtype=np.int32), |
| np.array(v1_c[:3], dtype=np.int32), |
| np.array(v2_c[:3], dtype=np.int32)]) |
| node_coords = v_parents_arr[np.arange(24) // 8] * 2 + offsets[np.arange(24) % 8] |
| keys = [tuple(int(x) for x in node_coords[i]) for i in range(24)] |
| for i in range(24): |
| if keys[i] not in key_to_coord: |
| key_to_coord[keys[i]] = node_coords[i] |
|
|
| |
| valid_mask = A[_IDX_I, _IDX_J] & A[_IDX_J, _IDX_K] & A[_IDX_I, _IDX_K] |
| tris = _TRI_INDICES[valid_mask] |
| if len(tris) > 0: |
| p_idx = tris // 8 |
| c_idx = tris % 8 |
| tri_blocks.append(v_parents_arr[p_idx] * 2 + offsets[c_idx]) |
| for i, j, k in tris: |
| ka, kb, kc = sorted((keys[int(i)], keys[int(j)], keys[int(k)])) |
| tri_edge_keys.add((ka, kb)); tri_edge_keys.add((kb, kc)); tri_edge_keys.add((ka, kc)) |
|
|
| |
| ei, ej = np.where(np.triu(A, 1)) |
| for i, j in zip(ei.tolist(), ej.tolist()): |
| ki, kj = keys[i], keys[j] |
| edge_keys.add((ki, kj) if ki <= kj else (kj, ki)) |
|
|
| |
| for i in np.where(np.concatenate(fa))[0].tolist(): |
| active_keys.add(keys[i]) |
|
|
| all_blocks = list(tri_blocks) |
|
|
| |
| for a, b in (edge_keys - tri_edge_keys): |
| all_blocks.append(np.stack([key_to_coord[a], key_to_coord[b], key_to_coord[b]])[None]) |
|
|
| |
| edge_nodes = set() |
| for a, b in edge_keys: |
| edge_nodes.add(a); edge_nodes.add(b) |
| for a in (active_keys - edge_nodes): |
| c = key_to_coord[a] |
| all_blocks.append(np.stack([c, c, c])[None]) |
|
|
| if not all_blocks: |
| return np.array([]), np.array([]) |
|
|
| flat_v = np.concatenate(all_blocks, axis=0).reshape(-1, 3) |
| V, inv = np.unique(flat_v, axis=0, return_inverse=True) |
| F_raw = inv.reshape(-1, 3) |
| F_sorted = np.sort(F_raw, axis=1) |
| _, unique_face_indices = np.unique(F_sorted, axis=0, return_index=True) |
| F = F_raw[unique_face_indices] |
| V = undiscretize(V, num_discrete=resolution) |
| return V, F |
|
|
|
|
| def reconstruct_base_mesh(output_tokens, resolution): |
| all_face_coords = [] |
| for f_token in output_tokens: |
| v0_c, v1_c, v2_c = f_token |
| c0 = np.array(v0_c[:3], dtype=np.int32) |
| c1 = np.array(v1_c[:3], dtype=np.int32) |
| c2 = np.array(v2_c[:3], dtype=np.int32) |
| all_face_coords.append(np.stack([c0, c1, c2])) |
|
|
| if not all_face_coords: |
| return np.array([]), np.array([]) |
|
|
| flat_v = np.concatenate(all_face_coords, axis=0).reshape(-1, 3) |
| V, inv = np.unique(flat_v, axis=0, return_inverse=True) |
| F_raw = inv.reshape(-1, 3) |
| F_sorted = np.sort(F_raw, axis=1) |
| _, unique_face_indices = np.unique(F_sorted, axis=0, return_index=True) |
| F = F_raw[unique_face_indices] |
| V = undiscretize(V, num_discrete=resolution) |
| return V, F |
|
|
|
|
| class MeshCodec: |
| def __init__(self, args): |
| self.max_depth = args.max_depth |
| self.max_resolution = args.max_resolution |
| self.use_hourglass = args.use_hourglass |
| self.max_children = NUM_OCTANTS |
|
|
| self.tokenizer = DynamicMeshTokenizer(resolution=self.max_resolution) |
|
|
| def encode_mesh(self, mesh: trimesh.Trimesh, target_depth: int, corrupt=None) -> tuple: |
| if isinstance(mesh, trimesh.Scene): |
| mesh = mesh.dump(concatenate=True) |
| V = np.asarray(mesh.vertices, dtype=np.float64) |
|
|
| res = 2 ** (target_depth + 1) |
| v_restored = undiscretize(discretize(V, num_discrete=res), num_discrete=res) |
| quantized_mesh = trimesh.Trimesh(vertices=v_restored, faces=mesh.faces, process=False) |
|
|
| |
| raw_seq, seq = _fast_fix3_level_tokens(V, mesh.faces, target_depth, self.tokenizer, corrupt=corrupt) |
| return raw_seq, seq, quantized_mesh |
|
|
|
|
| def decode_tokens(self, token_sequence: list, target_depth: int): |
| if os.environ.get("FIX3_FAST_DECODE", "1") != "1": |
| return super().decode_tokens(token_sequence, target_depth) |
| if token_sequence is None or len(token_sequence) == 0: |
| return None |
| return _fast_decode_tokens(token_sequence, target_depth, self.tokenizer) |
|
|
|
|
| def decode_raw_tokens(self, raw_token_sequence: list, target_depth: int) -> trimesh.Trimesh: |
| if not raw_token_sequence: |
| return None |
| resolution = 2 ** (target_depth) |
| raw_t = self.tokenizer.decode_raw(raw_token_sequence) |
| v, f = reconstruct_base_mesh(raw_t, resolution) |
| if len(v) == 0 or len(f) == 0: |
| return None |
| return trimesh.Trimesh(vertices=v, faces=f, process=False) |
|
|
| def get_model_config(self) -> dict: |
| t = self.tokenizer |
| return { |
| "total_vocab_size": t.vocab_size, |
| "coord_vocab_size": t.resolution, |
| "f_vocab_size": t.f_vocab, |
| "row_vocab_size": t.row_vocab, |
| "coord_offset": t.coord_offset, |
| "f_offset": t.f_offset, |
| "intra_row_offset": t.intra_row_offset, |
| "inter_row_offset": t.inter_row_offset, |
| "pad_token_id": t.pad_id, |
| "num_octants": NUM_OCTANTS, |
| "subcodes_per_position": NUM_OCTANTS, |
| "tokens_per_face": DynamicMeshTokenizer.TOKENS_PER_FACE, |
| "coord_per_face": 9, |
| "refine_positions_per_face": DynamicMeshTokenizer.REFINE_POS_PER_FACE, |
| "input_num_categories": t.resolution + 1, |
| |
| "output_num_categories": 1 + t.f_vocab + 2 * t.row_vocab, |
| } |
|
|
|
|
| |
| |
| |
| if __name__ == "__main__": |
| |
| |
| |
| |
| |
| |
| import time |
|
|
| |
| INPUT_DIRECTORY = "/root/junkai_lin/mul_res_diffusion/demo/complex_test_data/files" |
| OUTPUT_DIRECTORY = "./output/fix3_octant_factorized2" |
| os.makedirs(OUTPUT_DIRECTORY, exist_ok=True) |
|
|
| @dataclass |
| class MeshCodecArgs: |
| max_depth: int = 10 |
| max_resolution: int = 1024 |
| use_hourglass: bool = True |
| max_children: int = 8 |
|
|
| |
| args = MeshCodecArgs() |
| codec = MeshCodec(args=args) |
| print(f"[*] Initialized fix3 Codec. Total Vocab Size: {codec.tokenizer.vocab_size} " |
| f"(tokens/face={DynamicMeshTokenizer.TOKENS_PER_FACE})") |
|
|
| |
| input_dir_path = Path(INPUT_DIRECTORY) |
| mesh_files = sorted(input_dir_path.rglob("*.obj")) + sorted(input_dir_path.rglob("*.glb")) |
| if not mesh_files: |
| print(f"[!] No .obj or .glb files found in {INPUT_DIRECTORY}") |
|
|
| DEPTHS = [0, 1, 2, 3, 4, 5, 6] |
|
|
| |
| for idx, file_path in enumerate(mesh_files, 1): |
| obj_name = file_path.stem |
| print(f"\n{'=' * 50}") |
| print(f"[{idx}/{len(mesh_files)}] Processing: {file_path.name}") |
| print(f"{'=' * 50}") |
| file_start_time = time.perf_counter() |
|
|
| try: |
| raw_mesh = trimesh.load(str(file_path), force="mesh") |
| normalize(raw_mesh) |
| except Exception as e: |
| print(f"[!] load/normalize failed: {e}") |
| continue |
|
|
| |
| gt_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_gt.obj") |
| raw_mesh.export(gt_path) |
| print(f"[*] Saved GT -> {gt_path}") |
|
|
| |
| for depth in DEPTHS: |
| print(f"\n--- Depth {depth} | res(fine)={2 ** (depth + 1)} | octant 原位 ---") |
|
|
| |
| t0 = time.perf_counter() |
| raw_tokens, tokens, quantized_mesh = codec.encode_mesh(raw_mesh, target_depth=depth) |
| encode_time = time.perf_counter() - t0 |
| nface = len(tokens) // DynamicMeshTokenizer.TOKENS_PER_FACE |
| print(f"[Timer] Encoding: {encode_time:.4f}s | flat_tokens={len(tokens)} " |
| f"(nface={nface}, 模型逻辑位置={18 * nface})") |
|
|
| |
| same = (tokens == codec.tokenizer.encode(codec.tokenizer.decode(tokens))) |
| print(f"[Check] token 往返一致: {same}") |
|
|
| |
| if len(tokens) > 0: |
| t1 = time.perf_counter() |
| recon_mesh = codec.decode_tokens(tokens, target_depth=depth) |
| decode_fine_time = time.perf_counter() - t1 |
| if recon_mesh is not None: |
| out_path = os.path.join( |
| OUTPUT_DIRECTORY, f"{obj_name}_recon_d{depth}_res{2 ** (depth + 1)}_fix3.obj") |
| recon_mesh.export(out_path) |
| print(f"[Timer] Decoding (Fine): {decode_fine_time:.4f}s " |
| f"-> V={len(recon_mesh.vertices)} F={len(recon_mesh.faces)}") |
| else: |
| print(f"[Timer] Decoding (Fine): {decode_fine_time:.4f}s -> empty") |
|
|
| |
| if len(raw_tokens) > 0: |
| t2 = time.perf_counter() |
| raw_recon_mesh = codec.decode_raw_tokens(raw_tokens, target_depth=depth) |
| decode_coarse_time = time.perf_counter() - t2 |
| if raw_recon_mesh is not None: |
| raw_out_path = os.path.join( |
| OUTPUT_DIRECTORY, f"{obj_name}_raw_recon_d{depth}_res{2 ** depth}.obj") |
| raw_recon_mesh.export(raw_out_path) |
| print(f"[Timer] Decoding (Coarse): {decode_coarse_time:.4f}s " |
| f"-> V={len(raw_recon_mesh.vertices)} F={len(raw_recon_mesh.faces)}") |
|
|
| |
| if quantized_mesh is not None: |
| quant_out_path = os.path.join( |
| OUTPUT_DIRECTORY, f"{obj_name}_dis_d{depth + 1}_res{2 ** (depth + 1)}.obj") |
| quantized_mesh.export(quant_out_path) |
|
|
| file_total_time = time.perf_counter() - file_start_time |
| print(f"\n[*] Finished {file_path.name} in {file_total_time:.2f}s") |
|
|
|
|
| |
| |
| |
|
|
| _OCT_W = np.array([1, 2, 4], dtype=np.int64) |
|
|
| def _corrupt_fid_nodes(fid, nodes, und, nN, target_depth, corrupt): |
| """2a 暴露偏差训练腐蚀(仅训练): 模拟级联实测的两类粗输入误差。 |
| (a) 伪面: 在已存在的 GT 节点间加不该有的面(过连接)。下游打包按 fid 查表 -> |
| 非边对的 inter 行自动全 0(empty),教模型"看 PC 把伪面 refine 成空"。 |
| (b) 节点坐标抖动 ±jitter_max(位移): 只改 coord_part 的输入坐标,目标 F/intra/inter |
| 仍按真节点身份查表不变,教模型"输入坐标偏了也 refine 对"。 |
| 目标无需重新导出 —— 全部沿用 fid 索引,与干净路径同一套向量化打包。返回 (fid_aug, nodes_coord)。 |
| """ |
| sp_rate = float(corrupt.get("spurious_rate", 0.0)) |
| jp = float(corrupt.get("jitter_prob", 0.0)) |
|
|
| if sp_rate > 0.0 and nN >= 3 and und.shape[0] > 0: |
| radius = int(corrupt.get("spatial_radius", 2)) |
| n_add = int(sp_rate * fid.shape[0]) |
| if n_add > 0: |
| coord2idx = {(int(c[0]), int(c[1]), int(c[2])): i for i, c in enumerate(nodes)} |
| existing = set(map(tuple, np.sort(fid, axis=1).tolist())) |
| ne = und.shape[0] |
| spurious, tries, maxtries = [], 0, n_add * 6 |
| while len(spurious) < n_add and tries < maxtries: |
| tries += 1 |
| ei = int(np.random.randint(ne)) |
| ai, bi = int(und[ei, 0]), int(und[ei, 1]) |
| base = nodes[ai] if np.random.random() < 0.5 else nodes[bi] |
| off = np.random.randint(-radius, radius + 1, size=3) |
| if not off.any(): |
| continue |
| ck = (int(base[0] + off[0]), int(base[1] + off[1]), int(base[2] + off[2])) |
| ci = coord2idx.get(ck) |
| if ci is None or ci == ai or ci == bi: |
| continue |
| tri = tuple(sorted((ai, bi, ci))) |
| if tri in existing: |
| continue |
| existing.add(tri) |
| spurious.append(tri) |
| if spurious: |
| fid = np.concatenate([fid, np.asarray(spurious, dtype=fid.dtype)], axis=0) |
|
|
| nodes_coord = nodes |
| if jp > 0.0: |
| jmax = int(corrupt.get("jitter_max", 1)) |
| res_c = 2 ** target_depth |
| nodes_coord = nodes.copy() |
| mask = np.random.random(nN) < jp |
| k = int(mask.sum()) |
| if k > 0: |
| noise = np.random.randint(-jmax, jmax + 1, size=(k, 3)) |
| nodes_coord[mask] = np.clip(nodes_coord[mask] + noise, 0, res_c - 1) |
| return fid, nodes_coord |
|
|
| def _fast_fix3_level_tokens(vertices, faces, target_depth, tok, corrupt=None): |
| """等价替换 MatrixOctree(max_depth=td+1).build + encode_level_tokens(td) + tokenizer.encode/encode_raw。 |
| 返回 (raw_token_sequence, token_sequence)，与原版逐位一致。""" |
| res = 2 ** (target_depth + 1) |
| Q = discretize(np.asarray(vertices, dtype=np.float64), num_discrete=res).astype(np.int64) |
| faces = np.asarray(faces, dtype=np.int64).reshape(-1, 3) |
|
|
| |
| P_all = Q >> 1 |
| i_all = (Q & 1) @ _OCT_W |
|
|
| |
| if faces.size: |
| e = np.concatenate([faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]], axis=0) |
| u, v = Q[e[:, 0]], Q[e[:, 1]] |
| keep = np.any(u != v, axis=1) |
| u, v = u[keep], v[keep] |
| else: |
| u = v = np.zeros((0, 3), dtype=np.int64) |
| pu, pv = u >> 1, v >> 1 |
| iu = (u & 1) @ _OCT_W |
| iv = (v & 1) @ _OCT_W |
|
|
| |
| nodes, inv = np.unique(np.concatenate([P_all, pu, pv], axis=0), axis=0, return_inverse=True) |
| nV, M = P_all.shape[0], pu.shape[0] |
| occ_idx = inv[:nV] |
| a, b = inv[nV:nV + M], inv[nV + M:] |
| nN = nodes.shape[0] |
| intra = a == b |
|
|
| |
| Fbits = np.zeros(nN, dtype=np.int64) |
| np.bitwise_or.at(Fbits, occ_idx, np.int64(1) << i_all) |
|
|
| |
| intra_rows = np.zeros((nN, 8), dtype=np.int64) |
| ai, iui, ivi = a[intra], iu[intra], iv[intra] |
| np.bitwise_or.at(intra_rows, (ai, iui), np.int64(1) << ivi) |
| np.bitwise_or.at(intra_rows, (ai, ivi), np.int64(1) << iui) |
|
|
| |
| im = ~intra |
| ea, eb, eiu, eiv = a[im], b[im], iu[im], iv[im] |
| dkey = np.concatenate([ea * nN + eb, eb * nN + ea]) |
| drow = np.concatenate([eiu, eiv]) |
| dcol = np.concatenate([eiv, eiu]) |
| ukeys, kinv = np.unique(dkey, return_inverse=True) |
| inter_rows = np.zeros((max(len(ukeys), 1), 8), dtype=np.int64) |
| if dkey.size: |
| np.bitwise_or.at(inter_rows, (kinv, drow), np.int64(1) << dcol) |
|
|
| |
| if ea.size: |
| und = np.unique(np.stack([np.minimum(ea, eb), np.maximum(ea, eb)], axis=1), axis=0) |
| else: |
| und = np.zeros((0, 2), dtype=np.int64) |
| deg = np.zeros(nN, dtype=np.int64) |
| if und.size: |
| np.add.at(deg, und[:, 0], 1) |
| np.add.at(deg, und[:, 1], 1) |
| order = np.lexsort((und[:, 1], und[:, 0])) |
| fu, fv = und[order, 0], und[order, 1] |
| indptr = np.zeros(nN + 1, dtype=np.int64) |
| np.add.at(indptr, fu + 1, 1) |
| indptr = np.cumsum(indptr) |
| tu, tv, tw = _triangles_csr(indptr, fv.astype(np.int64)) |
| else: |
| tu = tv = tw = np.zeros(0, dtype=np.int64) |
|
|
| |
| tris = np.stack([tu, tv, tw], axis=1) if tu.size else np.zeros((0, 3), dtype=np.int64) |
| covered = set() |
| for x, y, z in tris: |
| x, y, z = int(x), int(y), int(z) |
| covered.add((x, y)); covered.add((y, z)); covered.add((x, z)) |
| extra = [] |
| for x, y in und: |
| x, y = int(x), int(y) |
| if (x, y) not in covered: |
| extra.append((x, y, y)) |
| for x in np.flatnonzero(deg == 0): |
| x = int(x) |
| extra.append((x, x, x)) |
| fid = np.concatenate([tris, np.asarray(extra, dtype=np.int64).reshape(-1, 3)], axis=0) |
| if fid.shape[0] == 0: |
| return [], [] |
|
|
| |
| c0, c1, c2 = nodes[fid[:, 0]], nodes[fid[:, 1]], nodes[fid[:, 2]] |
| order = np.lexsort((c2[:, 0], c2[:, 1], c2[:, 2], |
| c1[:, 0], c1[:, 1], c1[:, 2], |
| c0[:, 0], c0[:, 1], c0[:, 2])) |
| fid = fid[order] |
| nodes_coord = nodes |
| if corrupt is not None: |
| fid, nodes_coord = _corrupt_fid_nodes(fid, nodes, und, nN, target_depth, corrupt) |
| c0, c1, c2 = nodes_coord[fid[:, 0]], nodes_coord[fid[:, 1]], nodes_coord[fid[:, 2]] |
|
|
| |
| |
| |
| co, fo = tok.coord_offset, tok.f_offset |
| io, eo = tok.intra_row_offset, tok.inter_row_offset |
| nF = fid.shape[0] |
|
|
| coord_part = np.concatenate([c0, c1, c2], axis=1) + co |
| F_part = np.stack([Fbits[fid[:, 0]], Fbits[fid[:, 1]], Fbits[fid[:, 2]]], axis=1) + fo |
| intra_part = np.concatenate([intra_rows[fid[:, 0]], |
| intra_rows[fid[:, 1]], |
| intra_rows[fid[:, 2]]], axis=1) + io |
|
|
| inter_blocks = [] |
| for x, y in ((0, 1), (0, 2), (1, 2)): |
| A_, B_ = fid[:, x], fid[:, y] |
| rows = np.zeros((nF, 8), dtype=np.int64) |
| same = A_ == B_ |
| rows[same] = intra_rows[A_[same]] |
| ns = np.flatnonzero(~same) |
| if ns.size: |
| dk = A_[ns] * nN + B_[ns] |
| pos = np.searchsorted(ukeys, dk) |
| pos_c = np.minimum(pos, max(len(ukeys) - 1, 0)) |
| found = (ukeys[pos_c] == dk) if len(ukeys) else np.zeros(ns.size, dtype=bool) |
| rows[ns[found]] = inter_rows[pos_c[found]] |
| inter_blocks.append(rows) |
| inter_part = np.concatenate(inter_blocks, axis=1) + eo |
|
|
| token = np.concatenate([coord_part, F_part, intra_part, inter_part], axis=1) |
| token_sequence = token.ravel().tolist() |
| raw_token_sequence = coord_part.reshape(-1, 1).tolist() |
| return raw_token_sequence, token_sequence |
|
|
| _BIT_SHIFTS = np.arange(NUM_OCTANTS, dtype=np.int64) |
|
|
| _NODE_PARENT = np.arange(24) // NUM_OCTANTS |
|
|
| _NODE_OCTANT = np.arange(24) % NUM_OCTANTS |
|
|
| _OCT_OFFSETS = np.stack([np.arange(8) & 1, (np.arange(8) >> 1) & 1, (np.arange(8) >> 2) & 1], |
| axis=1).astype(np.int32) |
|
|
| _INTER_ROW_SLOT = np.array([0, 0, 1]) |
|
|
| _INTER_COL_SLOT = np.array([1, 2, 2]) |
|
|
| def _unpack_rows(values): |
| """int 行子码 -> bool 位。values (...,) -> (..., 8), 与 _int_to_row 同序(bit i = (v>>i)&1)。""" |
| return ((values[..., None].astype(np.int64) >> _BIT_SHIFTS) & 1).astype(bool) |
|
|
| def _pack_keys(coords, resolution): |
| """(...,3) int 坐标 -> (...) int64 唯一 key。坐标范围 [0, resolution)。""" |
| c = coords.astype(np.int64) |
| return (c[..., 0] * resolution + c[..., 1]) * resolution + c[..., 2] |
|
|
| def _fast_reconstruct_next_level(token_array, resolution, tokenizer): |
| """token_array: (nface, TOKENS_PER_FACE) int。返回 (V, F), 语义同 reconstruct_next_level_optimized。""" |
| from modules.octree_codec import _TRI_INDICES |
|
|
| nface = token_array.shape[0] |
| parents = (token_array[:, :9] - tokenizer.coord_offset).reshape(nface, 3, 3).astype(np.int32) |
| occupancy = _unpack_rows(token_array[:, 9:12] - tokenizer.f_offset) |
|
|
| def rows_to_mat(lo, hi, offset): |
| return _unpack_rows(token_array[:, lo:hi] - offset) |
|
|
| intra = [rows_to_mat(12 + 8 * i, 20 + 8 * i, tokenizer.intra_row_offset) for i in range(3)] |
| inter = [rows_to_mat(36 + 8 * i, 44 + 8 * i, tokenizer.inter_row_offset) for i in range(3)] |
|
|
| |
| adjacency = np.zeros((nface, 24, 24), dtype=bool) |
| for slot in range(3): |
| f = occupancy[:, slot] |
| block = intra[slot] & (f[:, :, None] & f[:, None, :]) |
| block = block | np.swapaxes(block, 1, 2) |
| adjacency[:, slot * 8:(slot + 1) * 8, slot * 8:(slot + 1) * 8] = block |
| for b in range(3): |
| ra, ca = _INTER_ROW_SLOT[b], _INTER_COL_SLOT[b] |
| block = inter[b] & (occupancy[:, ra][:, :, None] & occupancy[:, ca][:, None, :]) |
| adjacency[:, ra * 8:(ra + 1) * 8, ca * 8:(ca + 1) * 8] = block |
| adjacency[:, ca * 8:(ca + 1) * 8, ra * 8:(ra + 1) * 8] = np.swapaxes(block, 1, 2) |
|
|
| |
| node_coords = parents[:, _NODE_PARENT] * 2 + _OCT_OFFSETS[_NODE_OCTANT] |
| node_keys = _pack_keys(node_coords, resolution) |
|
|
| |
| i_idx, j_idx, k_idx = _TRI_INDICES[:, 0], _TRI_INDICES[:, 1], _TRI_INDICES[:, 2] |
| valid = (adjacency[:, i_idx, j_idx] & adjacency[:, j_idx, k_idx] & adjacency[:, i_idx, k_idx]) |
| face_of_tri, tri_of_tri = np.nonzero(valid) |
| tri_nodes = _TRI_INDICES[tri_of_tri] |
| tri_coords = node_coords[face_of_tri[:, None], tri_nodes] |
| tri_keys = node_keys[face_of_tri[:, None], tri_nodes] |
|
|
| |
| upper = np.triu(np.ones((24, 24), dtype=bool), 1) |
| face_of_edge, node_i, node_j = np.nonzero(adjacency & upper) |
| key_i = node_keys[face_of_edge, node_i] |
| key_j = node_keys[face_of_edge, node_j] |
| edge_pairs = np.unique(np.stack([np.minimum(key_i, key_j), np.maximum(key_i, key_j)], axis=1), axis=0) |
|
|
| |
| if len(tri_keys): |
| tri_sorted = np.sort(tri_keys, axis=1) |
| covered = np.concatenate([tri_sorted[:, [0, 1]], tri_sorted[:, [1, 2]], tri_sorted[:, [0, 2]]], axis=0) |
| covered = np.unique(covered, axis=0) |
| else: |
| covered = np.zeros((0, 2), dtype=np.int64) |
|
|
| blocks = [tri_coords] if len(tri_coords) else [] |
|
|
| |
| if len(edge_pairs): |
| edge_flat = edge_pairs[:, 0] * (resolution ** 3) + edge_pairs[:, 1] |
| covered_flat = (covered[:, 0] * (resolution ** 3) + covered[:, 1]) if len(covered) \ |
| else np.zeros(0, dtype=np.int64) |
| open_edges = edge_pairs[~np.isin(edge_flat, covered_flat)] |
| if len(open_edges): |
| key_to_coord = _build_key_lookup(node_keys, node_coords) |
| ca = key_to_coord[open_edges[:, 0]] |
| cb = key_to_coord[open_edges[:, 1]] |
| blocks.append(np.stack([ca, cb, cb], axis=1)) |
|
|
| |
| active_keys = np.unique(node_keys[occupancy.reshape(nface, 24)]) |
| if len(active_keys): |
| edge_nodes = np.unique(edge_pairs.reshape(-1)) if len(edge_pairs) else np.zeros(0, dtype=np.int64) |
| isolated = active_keys[~np.isin(active_keys, edge_nodes)] |
| if len(isolated): |
| key_to_coord = _build_key_lookup(node_keys, node_coords) |
| c = key_to_coord[isolated] |
| blocks.append(np.stack([c, c, c], axis=1)) |
|
|
| if not blocks: |
| return np.array([]), np.array([]) |
|
|
| flat_v = np.concatenate(blocks, axis=0).reshape(-1, 3) |
| V, inverse = np.unique(flat_v, axis=0, return_inverse=True) |
| faces_raw = inverse.reshape(-1, 3) |
| _, unique_face_indices = np.unique(np.sort(faces_raw, axis=1), axis=0, return_index=True) |
| F = faces_raw[unique_face_indices] |
| return undiscretize(V, num_discrete=resolution), F |
|
|
| def _build_key_lookup(node_keys, node_coords): |
| """key -> 坐标 的稠密查表(用 np.unique 的第一次出现), 供退化面/孤立点还原坐标。""" |
| flat_keys = node_keys.reshape(-1) |
| flat_coords = node_coords.reshape(-1, 3) |
| uniq, first = np.unique(flat_keys, return_index=True) |
| lookup = np.zeros((int(uniq.max()) + 1, 3), dtype=np.int32) |
| lookup[uniq] = flat_coords[first] |
| return lookup |
|
|
| def _fast_decode_tokens(token_sequence, target_depth, tokenizer): |
| """decode_tokens 的向量化实现。返回 trimesh.Trimesh 或 None。""" |
| arr = np.asarray(token_sequence).reshape(-1) |
| arr = arr[arr != tokenizer.pad_id] |
| per_face = tokenizer.TOKENS_PER_FACE |
| nface = arr.shape[0] // per_face |
| if nface == 0: |
| return None |
| token_array = arr[:nface * per_face].reshape(nface, per_face).astype(np.int64) |
| V, F = _fast_reconstruct_next_level(token_array, 2 ** (target_depth + 1), tokenizer) |
| if len(V) == 0 or len(F) == 0: |
| return None |
| return trimesh.Trimesh(vertices=V, faces=F, process=False) |
|
|
|
|
| FastMeshCodec = MeshCodec |
|
|
