mul / modules /octree_codec_acc.py
Upload folder using huggingface_hub
90ca518 verified 16 days ago
36.5 kB
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")