import numpy as np import trimesh from tqdm import tqdm import os import sys import glob import multiprocessing import random from collections import Counter 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 import itertools from dataclasses import dataclass # 预计算 24 个节点中选 3 个的所有可能组合 (共 2024 种) _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 # ========================================== class DynamicMeshTokenizer: def __init__(self, resolution, max_children): self.resolution = resolution self.max_children = max_children self.m_bits = max_children * max_children # 1. 特殊控制 Token self.pad_id = 0 # 2. 计算各特征在词表中的偏移量 (Offsets) self.coord_offset = 1 self.f_offset = self.coord_offset + self.resolution self.m_intra_offset = self.f_offset + 256 # 1x8 对应 256 种可能 self.m_vocab_size = 2 ** self.m_bits self.m_inter_offset = self.m_intra_offset + self.m_vocab_size # 最终 Transformer 的总词表大小 self.vocab_size = self.m_inter_offset + self.m_vocab_size def _list_to_int(self, bit_list): res = 0 for i, bit in enumerate(bit_list): if bit: res |= (1 << i) return res def _int_to_list(self, val, length): return [(val >> i) & 1 for i in range(length)] def _m_to_int(self, M): flat = [bit for row in M for bit in row] return self._list_to_int(flat) def _int_to_m(self, val): flat = self._int_to_list(val, self.m_bits) return [flat[i*self.max_children : (i+1)*self.max_children] for i in range(self.max_children)] # ------------------------------------ # 细分网格 (Fine Mesh) Tokenizer # ------------------------------------ def encode(self, faces_data, use_hourglass=True): """将细分几何特征列表编码为 1D Token 序列 (18 Tokens / Face)""" seq = [] for face in faces_data: v0, v1, v2, F0, F1, F2, M00, M11, M22, M01, M02, M12 = face seq.extend([v0[0] + self.coord_offset, v0[1] + self.coord_offset, v0[2] + self.coord_offset]) seq.extend([v1[0] + self.coord_offset, v1[1] + self.coord_offset, v1[2] + self.coord_offset]) seq.extend([v2[0] + self.coord_offset, v2[1] + self.coord_offset, v2[2] + self.coord_offset]) seq.extend([self._list_to_int(F0) + self.f_offset, self._list_to_int(F1) + self.f_offset, self._list_to_int(F2) + self.f_offset]) seq.extend([self._m_to_int(M00) + self.m_intra_offset, self._m_to_int(M11) + self.m_intra_offset, self._m_to_int(M22) + self.m_intra_offset]) seq.extend([self._m_to_int(M01) + self.m_inter_offset, self._m_to_int(M02) + self.m_inter_offset, self._m_to_int(M12) + self.m_inter_offset]) return seq def decode(self, seq): """将细分网格的 1D Token 序列解码回几何特征列表""" faces_data = [] clean_seq = [t for t in seq if t not in ([self.pad_id])] for i in range(0, len(clean_seq), 18): chunk = clean_seq[i:i+18] if len(chunk) < 18: break v0 = (chunk[0] - self.coord_offset, chunk[1] - self.coord_offset, chunk[2] - self.coord_offset) v1 = (chunk[3] - self.coord_offset, chunk[4] - self.coord_offset, chunk[5] - self.coord_offset) v2 = (chunk[6] - self.coord_offset, chunk[7] - self.coord_offset, chunk[8] - self.coord_offset) F0 = self._int_to_list(chunk[9] - self.f_offset, 8) F1 = self._int_to_list(chunk[10] - self.f_offset, 8) F2 = self._int_to_list(chunk[11] - self.f_offset, 8) M00 = self._int_to_m(chunk[12] - self.m_intra_offset) M11 = self._int_to_m(chunk[13] - self.m_intra_offset) M22 = self._int_to_m(chunk[14] - self.m_intra_offset) M01 = self._int_to_m(chunk[15] - self.m_inter_offset) M02 = self._int_to_m(chunk[16] - self.m_inter_offset) M12 = self._int_to_m(chunk[17] - self.m_inter_offset) faces_data.append((v0, v1, v2, F0, F1, F2, M00, M11, M22, M01, M02, M12)) return faces_data # ------------------------------------ # [新增] 基础/粗糙网格 (Raw Mesh) Tokenizer # ------------------------------------ def encode_raw(self, raw_faces_data, use_hourglass=True): """将粗糙层级特征（仅含顶点）编码为 1D Token 序列 (9 Tokens / Face)""" seq = [] for face in raw_faces_data: v0, v1, v2 = face seq.extend([v0[0] + self.coord_offset, v0[1] + self.coord_offset, v0[2] + self.coord_offset]) seq.extend([v1[0] + self.coord_offset, v1[1] + self.coord_offset, v1[2] + self.coord_offset]) seq.extend([v2[0] + self.coord_offset, v2[1] + self.coord_offset, v2[2] + self.coord_offset]) return seq def decode_raw(self, seq): """将粗糙层级的 1D Token 序列解码回几何特征列表 [(v0, v1, v2), ...]""" faces_data = [] clean_seq = [t for t in seq if t not in ([self.pad_id])] # 每个面片固定转化为 9 个 Token for i in range(0, len(clean_seq), 9): chunk = clean_seq[i:i+9] if len(chunk) < 9: break v0 = (chunk[0] - self.coord_offset, chunk[1] - self.coord_offset, chunk[2] - self.coord_offset) v1 = (chunk[3] - self.coord_offset, chunk[4] - self.coord_offset, chunk[5] - self.coord_offset) v2 = (chunk[6] - self.coord_offset, chunk[7] - self.coord_offset, chunk[8] - self.coord_offset) faces_data.append((v0, v1, v2)) return faces_data # --- 以下为原有的类声明 (OctreeFaceTokenBuilder, OctreeNode, MatrixOctree 等保持不变) --- class OctreeFaceTokenBuilder: def __init__(self, tree, max_children=4): self.tree = tree self.max_children = max_children 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. 提取未形成三角形的独立边 (线段结构，例如细长桌腿) for u in range(len(coords)): for v in adj[u]: edge = tuple(sorted([coords[u], coords[v]])) if edge not in covered_edges: # 将 1D 的线段转换为退化面片 (A, B, B) 以适配 Tokenizer 格式 degenerate_face = (edge[0], edge[1], edge[1]) raw_faces_set.add(degenerate_face) # 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: isolated_face = (coords[u], coords[u], coords[u]) raw_faces_set.add(isolated_face) raw_faces_at_depth = list(raw_faces_set) # 保持原始排序逻辑以确保输出的 Determinism 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 reduce_node_children_to_max(self, nodes_map): def bit_dist(a, b): return bin(a ^ b).count('1') # Pass 1: 基于原始图结构为所有节点一次性计算 merge plan，消除顺序依赖。 merge_plan = {} for coord, n in nodes_map.items(): active_indices = set() has_intra_row = n.intra_matrix.any(axis=1) has_intra_col = n.intra_matrix.any(axis=0) for i in range(8): if has_intra_row[i] or has_intra_col[i]: active_indices.add(i) for mat in n.inter_matrices.values(): has_inter_row = mat.any(axis=1) for i in range(8): if has_inter_row[i]: active_indices.add(i) for n_inter_coord in n.inter_neighbors: if n_inter_coord in nodes_map: neighbor = nodes_map[n_inter_coord] if coord in neighbor.inter_matrices: has_in_col = neighbor.inter_matrices[coord].any(axis=0) for i in range(8): if has_in_col[i]: active_indices.add(i) if len(active_indices) <= self.max_children: merge_plan[coord] = None continue child_degrees = [0] * 8 intra_row_counts = np.count_nonzero(n.intra_matrix, axis=1) intra_col_counts = np.count_nonzero(n.intra_matrix, axis=0) for i in range(8): child_degrees[i] += intra_row_counts[i] + intra_col_counts[i] out_inter_counts = np.zeros(8, dtype=int) for mat in n.inter_matrices.values(): out_inter_counts += np.count_nonzero(mat, axis=1) in_inter_counts = np.zeros(8, dtype=int) for n_inter_coord in n.inter_neighbors: if n_inter_coord in nodes_map: neighbor = nodes_map[n_inter_coord] if coord in neighbor.inter_matrices: in_inter_counts += np.count_nonzero(neighbor.inter_matrices[coord], axis=0) for i in range(8): child_degrees[i] += 2 * (out_inter_counts[i] + in_inter_counts[i]) active_list = sorted(list(active_indices)) scored = sorted( active_list, key=lambda idx: (-child_degrees[idx], idx) ) keep_nodes = sorted(scored[:self.max_children]) mapping = {idx: idx for idx in range(8)} for m_idx in active_list: if m_idx in keep_nodes: continue target_merge = min(keep_nodes, key=lambda x: (bit_dist(x, m_idx), x)) mapping[m_idx] = target_merge merge_plan[coord] = mapping # Pass 2: 根据全局 merge plan 同步重映射所有 intra/inter 矩阵。 for coord, n in nodes_map.items(): mapping = merge_plan.get(coord) if mapping is not None: old_intra = n.intra_matrix new_intra = np.zeros_like(old_intra) nz = np.argwhere(old_intra != 0) for i, j in nz: new_intra[mapping[int(i)], mapping[int(j)]] = 1 np.fill_diagonal(new_intra, 0) n.intra_matrix = new_intra for neigh_coord, mat in list(n.inter_matrices.items()): src_map = merge_plan.get(coord) dst_map = merge_plan.get(neigh_coord) if src_map is None and dst_map is None: continue new_mat = np.zeros_like(mat) nz = np.argwhere(mat != 0) for i, j in nz: ni = src_map[int(i)] if src_map is not None else int(i) nj = dst_map[int(j)] if dst_map is not None else int(j) new_mat[ni, nj] = 1 n.inter_matrices[neigh_coord] = new_mat def encode_level_tokens(self, depth): raw_faces, nodes_map = self.get_level_info(depth) self.reduce_node_children_to_max(nodes_map) node_F = {} node_active = {} empty_mat = np.zeros((8, 8), dtype=np.int8) for coord, n in nodes_map.items(): active_indices = set() has_intra_row = n.intra_matrix.any(axis=1) has_intra_col = n.intra_matrix.any(axis=0) for i in range(8): if has_intra_row[i] or has_intra_col[i]: active_indices.add(i) for mat in n.inter_matrices.values(): has_inter_row = mat.any(axis=1) for i in range(8): if has_inter_row[i]: active_indices.add(i) for n_inter_coord in n.inter_neighbors: if n_inter_coord in nodes_map: neighbor = nodes_map[n_inter_coord] if n.coord in neighbor.inter_matrices: has_in_col = neighbor.inter_matrices[n.coord].any(axis=0) for i in range(8): if has_in_col[i]: active_indices.add(i) F = [0] * 8 for idx in active_indices: F[idx] = 1 node_F[coord] = F node_active[coord] = sorted(list(active_indices)) memo_intra = {} memo_inter = {} def extract_NxN_intra(node_coord): if node_coord in memo_intra: return memo_intra[node_coord] active = node_active[node_coord][:self.max_children] n = nodes_map[node_coord] if len(active) > 0: M = np.zeros((self.max_children, self.max_children), dtype=np.int8) sub = n.intra_matrix[np.ix_(active, active)] M[:len(active), :len(active)] = sub np.fill_diagonal(M[:len(active), :len(active)], 1) res = M.tolist() else: res = np.zeros((self.max_children, self.max_children), dtype=np.int8).tolist() memo_intra[node_coord] = res return res def extract_NxN_inter(coord_A, coord_B): if coord_A == coord_B: return extract_NxN_intra(coord_A) key = (coord_A, coord_B) if key in memo_inter: return memo_inter[key] active_A = node_active[coord_A][:self.max_children] active_B = node_active[coord_B][:self.max_children] nA = nodes_map[coord_A] orig_M = nA.inter_matrices.get(coord_B, empty_mat) if len(active_A) > 0 and len(active_B) > 0: M = np.zeros((self.max_children, self.max_children), dtype=np.int8) M[:len(active_A), :len(active_B)] = orig_M[np.ix_(active_A, active_B)] res = M.tolist() else: res = np.zeros((self.max_children, self.max_children), dtype=np.int8).tolist() memo_inter[key] = res return res 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 = extract_NxN_intra(v0_c) M11 = extract_NxN_intra(v1_c) M22 = extract_NxN_intra(v2_c) M01 = extract_NxN_inter(v0_c, v1_c) M02 = extract_NxN_inter(v0_c, v2_c) M12 = extract_NxN_inter(v1_c, v2_c) f_encoded = ( v0_c, v1_c, v2_c, F0, F1, F2, M00, M11, M22, M01, M02, M12 ) output_tokens.append(f_encoded) 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=4): all_face_coords = [] indices = np.arange(8) offsets = np.stack([indices & 1, (indices >> 1) & 1, (indices >> 2) & 1], axis=1).astype(np.int32) def expand_intra(F, M_NxN): M8x8 = np.zeros((8, 8), dtype=bool) active = [i for i, val in enumerate(F) if val] for i_n, i_o in enumerate(active): for j_n, j_o in enumerate(active): if i_n < max_children and j_n < max_children and M_NxN[i_n][j_n]: M8x8[i_o][j_o] = True return M8x8 def expand_inter(FA, FB, M_NxN): M8x8 = np.zeros((8, 8), dtype=bool) activeA = [i for i, val in enumerate(FA) if val] activeB = [i for i, val in enumerate(FB) if val] for i_n, i_o in enumerate(activeA): for j_n, j_o in enumerate(activeB): if i_n < max_children and j_n < max_children and M_NxN[i_n][j_n]: M8x8[i_o][j_o] = True return M8x8 for f_token in output_tokens: v0_c, v1_c, v2_c, F0, F1, F2, M00, M11, M22, M01, M02, M12 = f_token v_parents = [np.array(v[:3], dtype=np.int32) for v in [v0_c, v1_c, v2_c]] in0 = expand_intra(F0, M00) in1 = expand_intra(F1, M11) in2 = expand_intra(F2, M22) m01 = expand_inter(F0, F1, M01) m02 = expand_inter(F0, F2, M02) m12 = expand_inter(F1, F2, M12) # in_mats = [in0, in1, in2] # inter_edges = { # (0, 1): m01, # (1, 2): m12, # (0, 2): m02 # } # def are_connected_opt(p_a, c_a, p_b, c_b): # # p_a 和 p_b 是 0, 1, 2 这三个顶点索引。 # # 只要索引相同，说明是查询同一个顶点内部；索引不同，无脑查 inter_edges # if p_a == p_b: # return in_mats[p_a][c_a, c_b] # if p_a < p_b: # return inter_edges[(p_a, p_b)][c_a, c_b] # else: # return inter_edges[(p_b, p_a)][c_b, c_a] # from itertools import combinations_with_replacement # for pi, pj, pk in combinations_with_replacement(range(3), 3): # for i in range(8): # start_j = i + 1 if pi == pj else 0 # for j in range(start_j, 8): # if not are_connected_opt(pi, i, pj, j): continue # start_k = j + 1 if pj == pk else 0 # for k in range(start_k, 8): # if are_connected_opt(pj, j, pk, k) and are_connected_opt(pk, k, pi, i): # c0 = v_parents[pi] * 2 + offsets[i] # c1 = v_parents[pj] * 2 + offsets[j] # c2 = v_parents[pk] * 2 + offsets[k] # all_face_coords.append(np.stack([c0, c1, c2])) # 1. 构建 24x24 全局邻接矩阵 A A = np.zeros((24, 24), dtype=bool) # 填充内部矩阵 (Intra)，使用 | .T 确保无向图的对称性 A[0:8, 0:8] = in0 | in0.T A[8:16, 8:16] = in1 | in1.T A[16:24, 16:24] = in2 | in2.T # 填充外部跨体素矩阵 (Inter) 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 # 2. 向量化寻找所有的 3-clique (三角形) # 只需要一行代码，判断 i-j, j-k, i-k 是否同时连通 valid_mask = A[_IDX_I, _IDX_J] & A[_IDX_J, _IDX_K] & A[_IDX_I, _IDX_K] valid_triangles = _TRI_INDICES[valid_mask] # 过滤出所有形成三角形的节点组合 # 3. 批量计算物理坐标 if len(valid_triangles) > 0: # 将 0~23 的索引还原为父节点索引 (p_idx: 0~2) 和子节点索引 (c_idx: 0~7) p_idx = valid_triangles // 8 c_idx = valid_triangles % 8 v_parents_arr = np.stack(v_parents) # shape: (3, 3) # 批量算出所有符合条件的三角面的世界坐标，shape: (N, 3, 3) coords = v_parents_arr[p_idx] * 2 + offsets[c_idx] all_face_coords.append(coords) 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) # mask_not_degenerate = (F_raw[:, 0] != F_raw[:, 1]) & \ # (F_raw[:, 1] != F_raw[:, 2]) & \ # (F_raw[:, 2] != F_raw[:, 0]) # F_valid = F_raw[mask_not_degenerate] F_valid = F_raw F_sorted = np.sort(F_valid, axis=1) _, unique_face_indices = np.unique(F_sorted, axis=0, return_index=True) F = F_valid[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) # mask_not_degenerate = (F_raw[:, 0] != F_raw[:, 1]) & \ # (F_raw[:, 1] != F_raw[:, 2]) & \ # (F_raw[:, 2] != F_raw[:, 0]) # F_valid = F_raw[mask_not_degenerate] F_valid = F_raw F_sorted = np.sort(F_valid, axis=1) _, unique_face_indices = np.unique(F_sorted, axis=0, return_index=True) F = F_valid[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 = args.max_children self.tokenizer = DynamicMeshTokenizer(resolution=self.max_resolution, max_children=self.max_children) def encode_mesh(self, mesh: trimesh.Trimesh, target_depth: int) -> tuple: mesh_copy = mesh.copy() if isinstance(mesh_copy, trimesh.Scene): mesh_copy = mesh_copy.dump(concatenate=True) v_restored = discretize(mesh_copy.vertices, num_discrete=2**(target_depth+1)) v_restored = undiscretize(v_restored, num_discrete=2**(target_depth+1)) # 4. 生成新的量化网格 quantized_mesh = trimesh.Trimesh(vertices=v_restored, faces=mesh.faces, process=False) # tree = MatrixOctree(max_depth=self.max_depth) tree = MatrixOctree(max_depth=target_depth+1) tree.build(mesh_copy.vertices, mesh_copy.faces) builder = OctreeFaceTokenBuilder(tree, max_children=self.max_children) raw_t, out_t = builder.encode_level_tokens(depth=target_depth) token_sequence = self.tokenizer.encode(out_t, self.use_hourglass) raw_token_sequence = self.tokenizer.encode_raw(raw_t, self.use_hourglass) return raw_token_sequence, token_sequence, quantized_mesh def decode_tokens(self, token_sequence: list, target_depth: int) -> trimesh.Trimesh: if not token_sequence: return None resolution = 2 ** (target_depth + 1) out_t = self.tokenizer.decode(token_sequence) v_next, f_next = reconstruct_next_level_optimized(out_t, resolution, max_children=self.max_children) if len(v_next) == 0 or len(f_next) == 0: return None detokenized_mesh = trimesh.Trimesh(vertices=v_next, faces=f_next, process=False) # trimesh.repair.fix_normals(detokenized_mesh) # clean_mesh(detokenized_mesh) return detokenized_mesh 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) # [修改] 先使用 tokenizer 从 1D 整数序列解码回 raw_t 结构 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 detokenized_mesh = trimesh.Trimesh(vertices=v, faces=f, process=False) # trimesh.repair.fix_normals(detokenized_mesh) # clean_mesh(detokenized_mesh) return detokenized_mesh def get_model_config(self) -> dict: return { "total_vocab_size": self.tokenizer.vocab_size, "coord_vocab_size": self.tokenizer.resolution, "f_vocab_size": 256, "m_vocab_size": self.tokenizer.m_vocab_size, "coord_offset": self.tokenizer.coord_offset, "f_offset": self.tokenizer.f_offset, "m_intra_offset": self.tokenizer.m_intra_offset, "m_inter_offset": self.tokenizer.m_inter_offset, "pad_token_id": self.tokenizer.pad_id, "max_children": self.max_children, "input_num_categories": self.tokenizer.resolution+1, "output_num_categories": 1 + 2*2**self.tokenizer.m_bits + 256, } # ========================================== # 测试重构与无损化流程 # ========================================== if __name__ == "__main__": # 1. 定义输入和输出文件夹 INPUT_DIRECTORY = "/root/junkai_lin/mul_res_diffusion/demo/complex_test_data/files" # 【修改这里】替换为你的模型文件夹路径 OUTPUT_DIRECTORY = "./output/complex_child8" os.makedirs(OUTPUT_DIRECTORY, exist_ok=True) import time @dataclass class MeshCodecArgs: max_depth: int = 10 max_resolution: int = 1024 use_hourglass: bool = True max_children: int = 8 # 2. 初始化 Codec (放在循环外，避免重复加载) args = MeshCodecArgs() codec = MeshCodec(args=args) print(f"[*] Initialized Codec. Total Vocab Size: {codec.tokenizer.vocab_size}") # 3. 获取所有 .obj 和 .glb 文件 input_dir_path = Path(INPUT_DIRECTORY) # rglob 会递归查找子文件夹，如果只需要当前层级，可以把 rglob 改成 glob mesh_files = list(input_dir_path.rglob("*.obj")) + list(input_dir_path.rglob("*.glb")) if not mesh_files: print(f"[!] No .obj or .glb files found in {INPUT_DIRECTORY}") # 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: # 加载并归一化 Mesh raw_mesh = trimesh.load(str(file_path), force="mesh") normalize(raw_mesh) # clean_mesh(raw_mesh) if raw_mesh: out_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_gt.obj") raw_mesh.export(out_path) print(f"[*] Successfully saved GT to: {out_path}") # 开始不同深度 (Depth) 的推理 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() # 记录单个文件处理开始 raw_mesh = trimesh.load(str(file_path), force="mesh") normalize(raw_mesh) if raw_mesh: out_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_gt.obj") raw_mesh.export(out_path) # 开始不同深度 (Depth) 的推理 for depth in [0,1,2,3,4,5,6]: print(f"\n--- Depth {depth} | N={args.max_children} ---") # --- 计时：Encoding --- t0 = time.perf_counter() raw_tokens, tokens, quantized_mesh = codec.encode_mesh(raw_mesh, target_depth=depth) encode_time = time.perf_counter() - t0 print(f"[Timer] Encoding: {encode_time:.4f}s") # --- 计时：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 print(f"[Timer] Decoding (Fine): {decode_fine_time:.4f}s") # 导出保存 out_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_recon_d{depth}_res{2**(depth+1)}_N{args.max_children}.obj") recon_mesh.export(out_path) # --- 计时：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 print(f"[Timer] Decoding (Coarse): {decode_coarse_time:.4f}s") 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) if quantized_mesh: 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")