"""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 # 预计算 24 个节点(3 槽位 x 8 octant)中选 3 个的所有组合，用于向量化找三角形 _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] # ========================================== # Tokenizer：octant 原位 + 按行子码 # ========================================== 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 # octant 原位下固定为 8；保留字段仅为兼容 config self.pad_id = 0 self.coord_offset = 1 self.f_offset = self.coord_offset + self.resolution self.row_vocab = 1 << NUM_OCTANTS # 2^8 = 256 self.f_vocab = 1 << NUM_OCTANTS # F 也是 8bit 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 # ---- bit <-> int ---- 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)] # ------------------------------------ # 细分网格 (Fine Mesh) —— 扁平 1D int 布局，每 face 60 个 token： # 9 coord + 3 F + 6 矩阵 x 8 行子码 # 顺序: [coord(9)] [F0 F1 F2] [M00 M11 M22 各 8 行] [M01 M02 M12 各 8 行] # F: 每个 1 个 token；intra/inter: 每个矩阵按行拆成 8 个行子码(每行 8bit) # 所有真实 token 都带 offset(>=1)，0 仅作 trailing pad，所以 F 不需要复制避坑。 # 与现有 dataset/collate(1D int, pad_id=0) 兼容；模型侧整理成 9 个 refine 位置/face # (F 位置 1 子码、M 位置 8 行子码 sum 成一个 embedding token)，逻辑序列 18/face 不变。 # ------------------------------------ TOKENS_PER_FACE = 9 + 3 + 6 * NUM_OCTANTS # 60 REFINE_POS_PER_FACE = 9 # 模型逻辑: 3 F + 3 intra + 3 inter 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) # F: 1 个 token for M in (M00, M11, M22): seq.extend(self._mat_to_rows(M, self.intra_row_offset)) # 各 8 行 for M in (M01, M02, M12): seq.extend(self._mat_to_rows(M, self.inter_row_offset)) # 各 8 行 return seq def decode(self, seq): faces_data = [] B = self.TOKENS_PER_FACE # 60 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 # ------------------------------------ # 粗糙网格 (Raw Mesh)：仅顶点坐标，每 face 9 个单子码位置 # ------------------------------------ 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 # ========================================== # Builder：直接构建 octant 原位 8x8 矩阵 # ========================================== class OctreeFaceTokenBuilder: def __init__(self, tree, max_children=NUM_OCTANTS): self.tree = tree self.max_children = NUM_OCTANTS # octant 原位固定 8 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() # 1. 3-clique (三角形) 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]]))) # 2. 未被三角形覆盖的独立边 -> 退化面片 (A, B, B) 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])) # 3. 孤立顶点 -> (A, A, A) 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) # F：每个节点的 active octant 掩码 (8bit) 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) # 8x8 octant 原位、对称 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 = [] # 逐 face 向量化收集的三角形 (N,3,3)，最后统一去重 edge_keys = set() # 全局所有边: canonical (key_a, key_b) tri_edge_keys = set() # 被三角形覆盖的边(全局) active_keys = set() # 全局 active 节点 key key_to_coord = {} # key(int 三元组) -> np.int32 坐标(3,) 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 # 24 个节点(slot p, octant c)的世界坐标 = v_parents[p]*2 + offset[c] 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] # (24,3) 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] # 三角形(逐 face 向量化)；同时把三角形的边记入全局 covered 集合 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)) # 全局收集 active 节点 for i in np.where(np.concatenate(fa))[0].tolist(): active_keys.add(keys[i]) all_blocks = list(tri_blocks) # 未被任何三角形覆盖的边 -> 退化面 (A, B, B) [全局判断，与 raw 的 get_level_info 一致] 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]) # 孤立的 active 节点(全局无任何边) -> 点 (A, A, A) 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 # octant 原位固定 8(忽略 config 的 max_children) 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) # corrupt(仅训练时由 dataset 传入)= 2a 暴露偏差腐蚀; None 时与原版逐位一致(推理不受影响) 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, # intra/inter 行子码各 256 "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, # 每个 refine 位置 8 个子码槽(F 复制；intra/inter 为 8 行) "tokens_per_face": DynamicMeshTokenizer.TOKENS_PER_FACE, # 60 扁平 token/face (9+3+6*8) "coord_per_face": 9, "refine_positions_per_face": DynamicMeshTokenizer.REFINE_POS_PER_FACE, # 9 (3F+3intra+3inter) "input_num_categories": t.resolution + 1, # 输出端：F / intra_row / inter_row 三类，每类 256；coord 为输入端条件 "output_num_categories": 1 + t.f_vocab + 2 * t.row_vocab, } # ========================================== # 自测：encode -> decode -> reconstruct 往返 # ========================================== if __name__ == "__main__": # ========================================== # 编解码 + 重构导出流程 (参照 fix2 的 __main__) # 每个 mesh: 导出 GT；对每个 depth 做 # encode -> decode_tokens(细分/Fine) / decode_raw_tokens(粗糙/Coarse) / 量化网格 # 并打印计时、token 数(60/face)与 token 级往返一致性。 # ========================================== import time # 1. 输入/输出文件夹 (按需修改) 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 # fix3 octant 原位固定 8 (此字段实际被忽略) # 2. 初始化 Codec 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})") # 3. 收集所有 .obj / .glb 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] # 4. 逐文件处理 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 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 原位 ---") # Encoding 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})") # token 级往返一致性 (encode->decode->re-encode) same = (tokens == codec.tokenizer.encode(codec.tokenizer.decode(tokens))) print(f"[Check] token 往返一致: {same}") # Decoding (Fine) 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") # Decoding (Coarse / Raw) 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)}") # 量化网格 (输入离散化后的 GT，用于对照重构质量上限) 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") # ========================================================================== # encode_mesh / decode_tokens 的 numba 加速实现细节(原 *_fast.py 并入) # ========================================================================== _OCT_W = np.array([1, 2, 4], dtype=np.int64) # octant 局部索引 = x + 2y + 4z 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) # ---- 节点(td) = 所有顶点叶子坐标的 parent；占据 F = 所有顶点的 octant ---- P_all = Q >> 1 i_all = (Q & 1) @ _OCT_W # ---- 叶子边表(每面 3 条，跳过零长边 u==v，与原版一致) ---- 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 # 节点表(np.unique 行字典序 = 坐标 (x,y,z) 升序 = 原版 face 内顶点 sorted() 同序) 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 # ---- F 行子码(8bit int)：所有顶点的 octant 占据 ---- Fbits = np.zeros(nN, dtype=np.int64) np.bitwise_or.at(Fbits, occ_idx, np.int64(1) << i_all) # ---- intra 行子码 [nN,8]：绝对原位、对称双向、无对角填充(与原版一致) ---- 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) # ---- inter 行子码：有向 (a,b) 分组(双向都存，与原版 build 一致) ---- 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) # ---- 无向节点邻接 -> CSR -> numba 三角形枚举 ---- 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) # ---- faces_ids：三角形 + 未覆盖边退化 (A,B,B) + 孤立点 (A,A,A)，集合与原版一致 ---- 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])) # 主键=c0.z(最后) ... 与原版 7 键一致+补齐平局 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]] # ---- 60/face 扁平 token 向量化打包 ---- # 注: F_part/intra_part/inter_part 下面仍用 fid(节点身份)查表 -> 腐蚀的伪面自动得正确目标 # (非边对 inter 全 0),坐标抖动只影响上面 coord_part,目标不变。 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 # (nF,9) 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 # (nF,24) 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]] # inter8(a,a)=intra8(a)，与原版一致 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]] # 未找到 -> 保持全 0(原版 empty 矩阵) inter_blocks.append(rows) inter_part = np.concatenate(inter_blocks, axis=1) + eo # (nF,24) token = np.concatenate([coord_part, F_part, intra_part, inter_part], axis=1) # (nF,60) token_sequence = token.ravel().tolist() raw_token_sequence = coord_part.reshape(-1, 1).tolist() # 9 coord/face, 各 [t] return raw_token_sequence, token_sequence _BIT_SHIFTS = np.arange(NUM_OCTANTS, dtype=np.int64) # _int_to_row: bit i = (v>>i)&1 _NODE_PARENT = np.arange(24) // NUM_OCTANTS # 24 个节点 -> 哪个父顶点 slot _NODE_OCTANT = np.arange(24) % NUM_OCTANTS # 24 个节点 -> 哪个 octant _OCT_OFFSETS = np.stack([np.arange(8) & 1, (np.arange(8) >> 1) & 1, (np.arange(8) >> 2) & 1], axis=1).astype(np.int32) # octant -> (dx,dy,dz) _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) # (nface,3,8) bool def rows_to_mat(lo, hi, offset): return _unpack_rows(token_array[:, lo:hi] - offset) # (nface,8,8) bool, [行,列] 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)] # ---- 24x24 邻接: intra 块按 F&F 掩码后对称化, inter 块按 Fa&Fb 掩码 ---- adjacency = np.zeros((nface, 24, 24), dtype=bool) for slot in range(3): f = occupancy[:, slot] # (nface,8) 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) # ---- 24 个节点的世界坐标与 key ---- node_coords = parents[:, _NODE_PARENT] * 2 + _OCT_OFFSETS[_NODE_OCTANT] # (nface,24,3) node_keys = _pack_keys(node_coords, resolution) # (nface,24) # ---- 三角形: 三条边都在邻接里 ---- 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) # 各 (ntri,) tri_nodes = _TRI_INDICES[tri_of_tri] # (ntri,3) 节点编号 tri_coords = node_coords[face_of_tri[:, None], tri_nodes] # (ntri,3,3) tri_keys = node_keys[face_of_tri[:, None], tri_nodes] # (ntri,3) # ---- 全部边(上三角) ---- 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 [] # ---- 未被任何三角形覆盖的边 -> 退化面 (A,B,B) ---- 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 节点(无任何边) -> 点 (A,A,A) ---- 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] # 与原版 clean_seq 一致 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 # 兼容别名: 加速实现已并入 MeshCodec