mul / utils /postprocessing.py
Upload folder using huggingface_hub
90ca518 verified 17 days ago
42.1 kB
from typing import *
import numpy as np
import torch
import utils3d
import nvdiffrast.torch as dr
from tqdm import tqdm
import trimesh
import trimesh.visual
import pyvista as pv
from pymeshfix import _meshfix
import igraph
PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53]
import itertools
import os as _os, sys as _sys, contextlib as _contextlib

MESHFIX_QUIET = _os.environ.get("MESHFIX_QUIET", "1") == "1"   # 屏蔽 PyMeshFix/TMesh C++ 层日志(Loading..%, WARNING-, % done)


@_contextlib.contextmanager
def _suppress_c_output(enabled=True):
    """临时把 fd1/fd2 重定向到 /dev/null: TMesh 的日志是 C 层 printf, Python 的 sys.stdout 拦不住。"""
    if not enabled:
        yield; return
    _sys.stdout.flush(); _sys.stderr.flush()
    saved_out, saved_err = _os.dup(1), _os.dup(2)
    devnull = _os.open(_os.devnull, _os.O_WRONLY)
    try:
        _os.dup2(devnull, 1); _os.dup2(devnull, 2)
        yield
    finally:
        _os.dup2(saved_out, 1); _os.dup2(saved_err, 2)
        _os.close(saved_out); _os.close(saved_err); _os.close(devnull)

def radical_inverse(base, n):
    val = 0
    inv_base = 1.0 / base
    inv_base_n = inv_base
    while n > 0:
        digit = n % base
        val += digit * inv_base_n
        n //= base
        inv_base_n *= inv_base
    return val

def halton_sequence(dim, n):
    return [radical_inverse(PRIMES[dim], n) for dim in range(dim)]

def hammersley_sequence(dim, n, num_samples):
    return [n / num_samples] + halton_sequence(dim - 1, n)

def sphere_hammersley_sequence(n, num_samples, offset=(0, 0), remap=False):
    u, v = hammersley_sequence(2, n, num_samples)
    u += offset[0] / num_samples
    v += offset[1]
    if remap:
        u = 2 * u if u < 0.25 else 2 / 3 * u + 1 / 3
    theta = np.arccos(1 - 2 * u) - np.pi / 2
    phi = v * 2 * np.pi
    return [phi, theta]

@torch.no_grad()
# def _fill_holes(
#     verts,
#     faces,
#     max_hole_size=0.04,
#     max_hole_nbe=32,
#     resolution=128,
#     num_views=500,
#     debug=False,
#     verbose=False,
#     only_fill_holes=False,
# ):
#     """
#     Rasterize a mesh from multiple views and remove invisible faces.
#     Also includes postprocessing to:
#         1. Remove connected components that are have low visibility.
#         2. Mincut to remove faces at the inner side of the mesh connected to the outer side with a small hole.

#     Args:
#         verts (torch.Tensor): Vertices of the mesh. Shape (V, 3).
#         faces (torch.Tensor): Faces of the mesh. Shape (F, 3).
#         max_hole_size (float): Maximum area of a hole to fill.
#         resolution (int): Resolution of the rasterization.
#         num_views (int): Number of views to rasterize the mesh.
#         verbose (bool): Whether to print progress.
#     """
#     if only_fill_holes:
#         mesh = _meshfix.PyTMesh()
#         mesh.load_array(verts.cpu().numpy(), faces.cpu().numpy())
#         mesh.fill_small_boundaries(nbe=max_hole_nbe, refine=True)
#         verts, faces = mesh.return_arrays()
#         verts = torch.tensor(verts, device='cuda', dtype=torch.float32)
#         faces = torch.tensor(faces, device='cuda', dtype=torch.int32)
#         if verbose:
#             tqdm.write('Skip mincut and run meshfix hole filling only')
#         return verts, faces
#     # Construct cameras
#     yaws = []
#     pitchs = []
#     for i in range(num_views):
#         y, p = sphere_hammersley_sequence(i, num_views)
#         yaws.append(y)
#         pitchs.append(p)
#     yaws = torch.tensor(yaws).cuda()
#     pitchs = torch.tensor(pitchs).cuda()
#     radius = 2.0
#     fov = torch.deg2rad(torch.tensor(40)).cuda()
#     projection = utils3d.torch.perspective_from_fov_xy(fov, fov, 1, 3)
#     views = []
#     for (yaw, pitch) in zip(yaws, pitchs):
#         orig = torch.tensor([
#             torch.sin(yaw) * torch.cos(pitch),
#             torch.cos(yaw) * torch.cos(pitch),
#             torch.sin(pitch),
#         ]).cuda().float() * radius
#         view = utils3d.torch.view_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
#         views.append(view)
#     views = torch.stack(views, dim=0)

#     # Rasterize
#     visblity = torch.zeros(faces.shape[0], dtype=torch.int32, device=verts.device)
#     rastctx = utils3d.torch.RastContext(backend='cuda')
#     for i in tqdm(range(views.shape[0]), total=views.shape[0], disable=not verbose, desc='Rasterizing'):
#         view = views[i]
#         buffers = utils3d.torch.rasterize_triangle_faces(
#             rastctx, verts[None], faces, resolution, resolution, view=view, projection=projection
#         )
#         face_id = buffers['face_id'][0][buffers['mask'][0] > 0.95] - 1
#         face_id = torch.unique(face_id).long()
#         visblity[face_id] += 1
#     visblity = visblity.float() / num_views
    
    

#     # Mincut
#     ## construct outer faces
#     edges, face2edge, edge_degrees = utils3d.torch.compute_edges(faces)
#     boundary_edge_indices = torch.nonzero(edge_degrees == 1).reshape(-1)
#     connected_components = utils3d.torch.compute_connected_components(faces, edges, face2edge)
#     outer_face_indices = torch.zeros(faces.shape[0], dtype=torch.bool, device=faces.device)
#     for i in range(len(connected_components)):
#         outer_face_indices[connected_components[i]] = visblity[connected_components[i]] > min(max(visblity[connected_components[i]].quantile(0.75).item(), 0.25), 0.5)
#     outer_face_indices = outer_face_indices.nonzero().reshape(-1)
    
#     ## construct inner faces
#     inner_face_indices = torch.nonzero(visblity == 0).reshape(-1)
#     if verbose:
#         tqdm.write(f'Found {inner_face_indices.shape[0]} invisible faces')
#     if inner_face_indices.shape[0] == 0:
#         return verts, faces
    
#     ## Construct dual graph (faces as nodes, edges as edges)
#     dual_edges, dual_edge2edge = utils3d.torch.compute_dual_graph(face2edge)
#     dual_edge2edge = edges[dual_edge2edge]
#     dual_edges_weights = torch.norm(verts[dual_edge2edge[:, 0]] - verts[dual_edge2edge[:, 1]], dim=1)
#     if verbose:
#         tqdm.write(f'Dual graph: {dual_edges.shape[0]} edges')

#     ## solve mincut problem
#     ### construct main graph
#     g = igraph.Graph()
#     g.add_vertices(faces.shape[0])
#     g.add_edges(dual_edges.cpu().numpy())
#     g.es['weight'] = dual_edges_weights.cpu().numpy()
    
#     ### source and target
#     g.add_vertex('s')
#     g.add_vertex('t')
    
#     ### connect invisible faces to source
#     g.add_edges([(f, 's') for f in inner_face_indices], attributes={'weight': torch.ones(inner_face_indices.shape[0], dtype=torch.float32).cpu().numpy()})
    
#     ### connect outer faces to target
#     g.add_edges([(f, 't') for f in outer_face_indices], attributes={'weight': torch.ones(outer_face_indices.shape[0], dtype=torch.float32).cpu().numpy()})
                
#     ### solve mincut
#     cut = g.mincut('s', 't', (np.array(g.es['weight']) * 1000).tolist())
#     remove_face_indices = torch.tensor([v for v in cut.partition[0] if v < faces.shape[0]], dtype=torch.long, device=faces.device)
#     if verbose:
#         tqdm.write(f'Mincut solved, start checking the cut')
    
#     ### check if the cut is valid with each connected component
#     to_remove_cc = utils3d.torch.compute_connected_components(faces[remove_face_indices])
#     if debug:
#         tqdm.write(f'Number of connected components of the cut: {len(to_remove_cc)}')
#     valid_remove_cc = []
#     cutting_edges = []
#     for cc in to_remove_cc:
#         #### check if the connected component has low visibility
#         visblity_median = visblity[remove_face_indices[cc]].median()
#         if debug:
#             tqdm.write(f'visblity_median: {visblity_median}')
#         if visblity_median > 0.25:
#             continue
        
#         #### check if the cuting loop is small enough
#         cc_edge_indices, cc_edges_degree = torch.unique(face2edge[remove_face_indices[cc]], return_counts=True)
#         cc_boundary_edge_indices = cc_edge_indices[cc_edges_degree == 1]
#         cc_new_boundary_edge_indices = cc_boundary_edge_indices[~torch.isin(cc_boundary_edge_indices, boundary_edge_indices)]
#         if len(cc_new_boundary_edge_indices) > 0:
#             cc_new_boundary_edge_cc = utils3d.torch.compute_edge_connected_components(edges[cc_new_boundary_edge_indices])
#             cc_new_boundary_edges_cc_center = [verts[edges[cc_new_boundary_edge_indices[edge_cc]]].mean(dim=1).mean(dim=0) for edge_cc in cc_new_boundary_edge_cc]
#             cc_new_boundary_edges_cc_area = []
#             for i, edge_cc in enumerate(cc_new_boundary_edge_cc):
#                 _e1 = verts[edges[cc_new_boundary_edge_indices[edge_cc]][:, 0]] - cc_new_boundary_edges_cc_center[i]
#                 _e2 = verts[edges[cc_new_boundary_edge_indices[edge_cc]][:, 1]] - cc_new_boundary_edges_cc_center[i]
#                 cc_new_boundary_edges_cc_area.append(torch.norm(torch.cross(_e1, _e2, dim=-1), dim=1).sum() * 0.5)
#             if debug:
#                 cutting_edges.append(cc_new_boundary_edge_indices)
#                 tqdm.write(f'Area of the cutting loop: {cc_new_boundary_edges_cc_area}')
#             if any([l > max_hole_size for l in cc_new_boundary_edges_cc_area]):
#                 continue
            
#         valid_remove_cc.append(cc)
        
#     if debug:
#         face_v = verts[faces].mean(dim=1).cpu().numpy()
#         vis_dual_edges = dual_edges.cpu().numpy()
#         vis_colors = np.zeros((faces.shape[0], 3), dtype=np.uint8)
#         vis_colors[inner_face_indices.cpu().numpy()] = [0, 0, 255]
#         vis_colors[outer_face_indices.cpu().numpy()] = [0, 255, 0]
#         vis_colors[remove_face_indices.cpu().numpy()] = [255, 0, 255]
#         if len(valid_remove_cc) > 0:
#             vis_colors[remove_face_indices[torch.cat(valid_remove_cc)].cpu().numpy()] = [255, 0, 0]
#         utils3d.io.write_ply('dbg_dual.ply', face_v, edges=vis_dual_edges, vertex_colors=vis_colors)
        
#         vis_verts = verts.cpu().numpy()
#         vis_edges = edges[torch.cat(cutting_edges)].cpu().numpy()
#         utils3d.io.write_ply('dbg_cut.ply', vis_verts, edges=vis_edges)
        
    
#     if len(valid_remove_cc) > 0:
#         remove_face_indices = remove_face_indices[torch.cat(valid_remove_cc)]
#         mask = torch.ones(faces.shape[0], dtype=torch.bool, device=faces.device)
#         mask[remove_face_indices] = 0
#         faces = faces[mask]
#         faces, verts = utils3d.torch.remove_unreferenced_vertices(faces, verts)
#         if verbose:
#             tqdm.write(f'Removed {(~mask).sum()} faces by mincut')
#     else:
#         if verbose:
#             tqdm.write(f'Removed 0 faces by mincut')
            
#     mesh = _meshfix.PyTMesh()
#     mesh.load_array(verts.cpu().numpy(), faces.cpu().numpy())
#     mesh.fill_small_boundaries(nbe=max_hole_nbe, refine=True)
#     verts, faces = mesh.return_arrays()
#     verts, faces = torch.tensor(verts, device='cuda', dtype=torch.float32), torch.tensor(faces, device='cuda', dtype=torch.int32)

#     return verts, faces

def _detect_degenerate_faces(
    verts: torch.Tensor,
    faces: torch.Tensor,
    area_eps: float = 1e-12,
) -> torch.Tensor:
    """
    检测退化三角面。
    判定条件：
        1. 一个面中存在重复顶点索引；
        2. 三角形二倍面积小于等于 area_eps；
        3. 三角形顶点坐标中存在 NaN 或 Inf。
    Args:
        verts:
            [V, 3] 顶点。
        faces:
            [F, 3] 三角面。
        area_eps:
            二倍面积阈值。
    Returns:
        [F] bool mask，True 表示退化面。
    """
    if faces.numel() == 0:
        return torch.zeros(
            faces.shape[0],
            dtype=torch.bool,
            device=faces.device,
        )

    faces_long = faces.long()

    repeated_vertex = (
        (faces_long[:, 0] == faces_long[:, 1])
        | (faces_long[:, 1] == faces_long[:, 2])
        | (faces_long[:, 2] == faces_long[:, 0])
    )

    triangles = verts[faces_long]

    invalid_position = ~torch.isfinite(triangles).all(
        dim=-1
    ).all(dim=-1)

    edge_1 = triangles[:, 1] - triangles[:, 0]
    edge_2 = triangles[:, 2] - triangles[:, 0]

    double_area = torch.linalg.vector_norm(
        torch.cross(edge_1, edge_2, dim=-1),
        dim=-1,
    )

    zero_area = double_area <= area_eps

    return repeated_vertex | zero_area | invalid_position


def _find_matching_vertex(
    position: np.ndarray,
    vertices_list: list,
    vertex_buckets: dict,
    tolerance: float,
    unavailable_indices: set,
):
    """
    在已有顶点中查找与 position 接近的顶点。
    会搜索当前量化单元及其相邻的 26 个单元，避免浮点量化边界问题。
    """
    base_key = np.round(
        position / tolerance
    ).astype(np.int64)

    best_index = None
    best_distance = np.inf

    for offset in itertools.product(
        (-1, 0, 1),
        repeat=3,
    ):
        key = tuple(
            base_key + np.asarray(offset, dtype=np.int64)
        )

        for candidate in vertex_buckets.get(key, []):
            if candidate in unavailable_indices:
                continue

            distance = np.linalg.norm(
                vertices_list[candidate] - position
            )

            if (
                distance <= tolerance * 2.0
                and distance < best_distance
            ):
                best_distance = distance
                best_index = candidate

    return best_index


def _restore_degenerate_faces(
    vertices: np.ndarray,
    faces: np.ndarray,
    original_vertices: np.ndarray,
    preserved_faces: np.ndarray,
    vertex_match_eps: float = 1e-7,
):
    """
    将预先保存的退化面恢复到处理后的 mesh 中。
    PyMeshFix 可能重排或删除顶点，因此不能直接使用原 face index。
    本函数根据顶点坐标重新建立映射。
    同一个原始顶点 ID 会映射到同一个新顶点 ID；
    不同但坐标相同的原始顶点会尽量保持为不同顶点。
    Args:
        vertices:
            处理后的 [V, 3] 顶点。
        faces:
            处理后的 [F, 3] 面。
        original_vertices:
            处理前的原始顶点。
        preserved_faces:
            需要恢复的原始退化面索引。
        vertex_match_eps:
            顶点坐标匹配的相对容差。
    Returns:
        output_vertices, output_faces
    """
    vertices = np.asarray(vertices)
    faces = np.asarray(faces)
    original_vertices = np.asarray(original_vertices)
    preserved_faces = np.asarray(preserved_faces)

    if preserved_faces.size == 0:
        return vertices, faces

    if vertices.ndim != 2:
        vertices = vertices.reshape(-1, 3)

    if faces.size == 0:
        faces = np.empty(
            (0, 3),
            dtype=preserved_faces.dtype,
        )
    else:
        faces = faces.reshape(-1, 3)

    if len(vertices) > 0:
        bbox_diagonal = np.linalg.norm(
            vertices.max(axis=0)
            - vertices.min(axis=0)
        )
    elif len(original_vertices) > 0:
        bbox_diagonal = np.linalg.norm(
            original_vertices.max(axis=0)
            - original_vertices.min(axis=0)
        )
    else:
        bbox_diagonal = 1.0

    tolerance = max(
        float(vertex_match_eps)
        * max(float(bbox_diagonal), 1.0),
        1e-12,
    )

    vertices_list = [
        vertex.copy()
        for vertex in vertices
    ]

    vertex_buckets = {}

    for vertex_index, vertex in enumerate(vertices_list):
        key = tuple(
            np.round(
                vertex / tolerance
            ).astype(np.int64)
        )

        vertex_buckets.setdefault(
            key,
            [],
        ).append(vertex_index)

    original_vertex_ids = np.unique(
        preserved_faces.reshape(-1)
    )

    old_to_new = {}

    # 防止两个不同的原始顶点 ID 因为坐标相同而被错误合并。
    claimed_existing_indices = set()

    for old_vertex_id in original_vertex_ids:
        old_vertex_id = int(old_vertex_id)
        position = original_vertices[old_vertex_id]

        matched_index = _find_matching_vertex(
            position=position,
            vertices_list=vertices_list,
            vertex_buckets=vertex_buckets,
            tolerance=tolerance,
            unavailable_indices=claimed_existing_indices,
        )

        if matched_index is None:
            matched_index = len(vertices_list)
            vertices_list.append(position.copy())

            key = tuple(
                np.round(
                    position / tolerance
                ).astype(np.int64)
            )

            vertex_buckets.setdefault(
                key,
                [],
            ).append(matched_index)
        else:
            claimed_existing_indices.add(matched_index)

        old_to_new[old_vertex_id] = matched_index

    restored_faces = np.empty_like(
        preserved_faces,
        dtype=np.int64,
    )

    for column in range(3):
        restored_faces[:, column] = np.fromiter(
            (
                old_to_new[int(vertex_id)]
                for vertex_id
                in preserved_faces[:, column]
            ),
            dtype=np.int64,
            count=preserved_faces.shape[0],
        )

    output_vertices = np.asarray(
        vertices_list,
        dtype=vertices.dtype,
    )

    restored_faces = restored_faces.astype(
        faces.dtype,
        copy=False,
    )

    output_faces = np.concatenate(
        [faces, restored_faces],
        axis=0,
    )

    return output_vertices, output_faces


def _run_meshfix_hole_filling(
    verts: torch.Tensor,
    faces: torch.Tensor,
    max_hole_nbe: int,
):
    """
    调用 PyMeshFix 进行小边界补洞。
    注意：PyMeshFix 可能自行清理退化面。
    需要保留的退化面应在调用本函数前暂时移除，
    并在本函数执行结束后恢复。
    """
    if faces.numel() == 0:
        return verts, faces

    device = verts.device
    verts_dtype = verts.dtype
    faces_dtype = faces.dtype

    with _suppress_c_output(MESHFIX_QUIET):
        mesh = _meshfix.PyTMesh()

        mesh.load_array(
            verts.detach().cpu().numpy(),
            faces.detach().cpu().numpy(),
        )

        mesh.fill_small_boundaries(
            nbe=max_hole_nbe,
            refine=True,
        )

        output_verts, output_faces = mesh.return_arrays()

    output_verts = torch.as_tensor(
        output_verts,
        device=device,
        dtype=verts_dtype,
    )

    output_faces = torch.as_tensor(
        output_faces,
        device=device,
        dtype=faces_dtype,
    )

    return output_verts, output_faces


def _fill_holes(
    verts,
    faces,
    max_hole_size=0.04,
    max_hole_nbe=32,
    resolution=128,
    num_views=500,
    debug=False,
    verbose=False,
    only_fill_holes=False,
    keep_degenerate_faces=False,
    degenerate_area_eps=1e-12,
    vertex_match_eps=1e-7,
):
    """
    Rasterize a mesh from multiple views and remove invisible faces.
    后处理包括：
        1. 删除低可见性的连通区域；
        2. 使用 mincut 删除通过小孔与外侧连接的内部面；
        3. 使用 PyMeshFix 填补小边界洞。
    Args:
        verts:
            torch.Tensor，[V, 3] 顶点坐标。
        faces:
            torch.Tensor，[F, 3] 三角面。
        max_hole_size:
            mincut 新产生切割环允许的最大面积。
        max_hole_nbe:
            PyMeshFix 填补边界时允许的最大边数量。
        resolution:
            光栅化分辨率。
        num_views:
            多视角光栅化视角数量。
        debug:
            是否输出 debug 文件。
        verbose:
            是否输出进度信息。
        only_fill_holes:
            True 时跳过光栅化和 mincut，只运行 MeshFix 补洞。
        keep_degenerate_faces:
            False：
                保持原始行为。退化面可能在 mincut 或 PyMeshFix 中消失。
            True：
                处理前保存所有退化面，并暂时将它们从处理网格中移除；
                补洞结束后再将其恢复。
        degenerate_area_eps:
            判断零面积面的二倍面积阈值。
        vertex_match_eps:
            恢复退化面时的顶点坐标匹配容差。
    Returns:
        verts:
            处理后的顶点。
        faces:
            处理后的三角面。
    注意：
        keep_degenerate_faces=True 时，退化面会被追加到输出 faces 末尾，
        所以输出 face 顺序不保证与输入完全一致。
    """
    if verts is None or faces is None:
        return verts, faces

    if faces.numel() == 0:
        return verts, faces

    original_device = verts.device
    original_verts_dtype = verts.dtype
    original_faces_dtype = faces.dtype

    original_vertices_numpy = None
    preserved_degenerate_faces_numpy = None

    # ============================================================
    # 保存退化面，并将其从后续处理网格中暂时移除
    # ============================================================

    if keep_degenerate_faces:
        degenerate_mask = _detect_degenerate_faces(
            verts=verts,
            faces=faces,
            area_eps=degenerate_area_eps,
        )

        num_degenerate = int(
            degenerate_mask.sum().item()
        )

        if verbose:
            tqdm.write(
                f'Found {num_degenerate} degenerate input faces'
            )

        if num_degenerate > 0:
            original_vertices_numpy = (
                verts.detach().cpu().numpy().copy()
            )

            preserved_degenerate_faces_numpy = (
                faces[degenerate_mask]
                .detach()
                .cpu()
                .numpy()
                .copy()
            )

            faces = faces[~degenerate_mask]

    def finalize_output(
        output_verts: torch.Tensor,
        output_faces: torch.Tensor,
    ):
        """
        在返回前恢复之前保存的退化面。
        """
        if (
            not keep_degenerate_faces
            or preserved_degenerate_faces_numpy is None
            or preserved_degenerate_faces_numpy.size == 0
        ):
            return output_verts, output_faces

        output_vertices_numpy = (
            output_verts.detach().cpu().numpy()
        )

        output_faces_numpy = (
            output_faces.detach().cpu().numpy()
        )

        (
            output_vertices_numpy,
            output_faces_numpy,
        ) = _restore_degenerate_faces(
            vertices=output_vertices_numpy,
            faces=output_faces_numpy,
            original_vertices=original_vertices_numpy,
            preserved_faces=preserved_degenerate_faces_numpy,
            vertex_match_eps=vertex_match_eps,
        )

        output_verts = torch.as_tensor(
            output_vertices_numpy,
            device=original_device,
            dtype=original_verts_dtype,
        )

        output_faces = torch.as_tensor(
            output_faces_numpy,
            device=original_device,
            dtype=original_faces_dtype,
        )

        if verbose:
            tqdm.write(
                f'Restored '
                f'{len(preserved_degenerate_faces_numpy)} '
                f'degenerate faces'
            )

        return output_verts, output_faces

    # 如果所有输入面都是退化面，没有可补洞的有效表面
    if faces.numel() == 0:
        if keep_degenerate_faces:
            return finalize_output(
                verts[:0],
                faces.reshape(0, 3),
            )

        return verts, faces

    # ============================================================
    # 只运行 PyMeshFix 补洞
    # ============================================================

    if only_fill_holes:
        verts, faces = _run_meshfix_hole_filling(
            verts=verts,
            faces=faces,
            max_hole_nbe=max_hole_nbe,
        )

        if verbose:
            tqdm.write(
                'Skip mincut and run MeshFix hole filling only'
            )

        return finalize_output(verts, faces)

    device = verts.device

    # ============================================================
    # Construct cameras
    # ============================================================

    yaws = []
    pitches = []

    for view_index in range(num_views):
        yaw, pitch = sphere_hammersley_sequence(
            view_index,
            num_views,
        )

        yaws.append(yaw)
        pitches.append(pitch)

    yaws = torch.as_tensor(
        yaws,
        device=device,
        dtype=torch.float32,
    )

    pitches = torch.as_tensor(
        pitches,
        device=device,
        dtype=torch.float32,
    )

    radius = 2.0

    fov = torch.deg2rad(
        torch.tensor(
            40.0,
            device=device,
            dtype=torch.float32,
        )
    )

    projection = utils3d.torch.perspective_from_fov_xy(
        fov,
        fov,
        1,
        3,
    )

    target = torch.zeros(
        3,
        device=device,
        dtype=torch.float32,
    )

    up = torch.tensor(
        [0.0, 0.0, 1.0],
        device=device,
        dtype=torch.float32,
    )

    views = []

    for yaw, pitch in zip(yaws, pitches):
        origin = torch.stack(
            [
                torch.sin(yaw) * torch.cos(pitch),
                torch.cos(yaw) * torch.cos(pitch),
                torch.sin(pitch),
            ],
            dim=0,
        ).float() * radius

        view = utils3d.torch.view_look_at(
            origin,
            target,
            up,
        )

        views.append(view)

    views = torch.stack(
        views,
        dim=0,
    )

    # ============================================================
    # Rasterize
    # ============================================================

    visibility = torch.zeros(
        faces.shape[0],
        dtype=torch.int32,
        device=device,
    )

    rastctx = utils3d.torch.RastContext(
        backend='cuda'
    )

    for view_index in tqdm(
        range(views.shape[0]),
        total=views.shape[0],
        disable=not verbose,
        desc='Rasterizing',
    ):
        view = views[view_index]

        buffers = utils3d.torch.rasterize_triangle_faces(
            rastctx,
            verts[None],
            faces,
            resolution,
            resolution,
            view=view,
            projection=projection,
        )

        valid_pixels = buffers['mask'][0] > 0.95

        face_id = (
            buffers['face_id'][0][valid_pixels] - 1
        )

        if face_id.numel() == 0:
            continue

        face_id = torch.unique(
            face_id
        ).long()

        # 避免异常背景 ID
        face_id = face_id[
            (face_id >= 0)
            & (face_id < faces.shape[0])
        ]

        if face_id.numel() > 0:
            visibility[face_id] += 1

    visibility = (
        visibility.float()
        / float(num_views)
    )

    # ============================================================
    # Mincut: construct outer faces
    # ============================================================

    (
        edges,
        face2edge,
        edge_degrees,
    ) = utils3d.torch.compute_edges(faces)

    boundary_edge_indices = torch.nonzero(
        edge_degrees == 1,
        as_tuple=False,
    ).reshape(-1)

    connected_components = (
        utils3d.torch.compute_connected_components(
            faces,
            edges,
            face2edge,
        )
    )

    outer_face_mask = torch.zeros(
        faces.shape[0],
        dtype=torch.bool,
        device=faces.device,
    )

    for component in connected_components:
        component_visibility = visibility[component]

        if component_visibility.numel() == 0:
            continue

        visibility_threshold = min(
            max(
                component_visibility
                .quantile(0.75)
                .item(),
                0.25,
            ),
            0.5,
        )

        outer_face_mask[component] = (
            component_visibility
            > visibility_threshold
        )

    outer_face_indices = torch.nonzero(
        outer_face_mask,
        as_tuple=False,
    ).reshape(-1)

    # ============================================================
    # Construct inner faces
    # ============================================================

    inner_face_indices = torch.nonzero(
        visibility == 0,
        as_tuple=False,
    ).reshape(-1)

    if verbose:
        tqdm.write(
            f'Found {inner_face_indices.shape[0]} '
            f'invisible faces'
        )

    # 保持原函数行为：
    # 若不存在不可见面，则不执行 mincut 和最终 MeshFix。
    if inner_face_indices.shape[0] == 0:
        return finalize_output(verts, faces)

    # ============================================================
    # Construct dual graph
    # ============================================================

    (
        dual_edges,
        dual_edge2edge,
    ) = utils3d.torch.compute_dual_graph(
        face2edge
    )

    dual_edge_vertices = edges[
        dual_edge2edge
    ]

    dual_edges_weights = torch.linalg.vector_norm(
        verts[dual_edge_vertices[:, 0]]
        - verts[dual_edge_vertices[:, 1]],
        dim=1,
    )

    if verbose:
        tqdm.write(
            f'Dual graph: {dual_edges.shape[0]} edges'
        )

    # ============================================================
    # Solve mincut
    # ============================================================

    graph = igraph.Graph()

    graph.add_vertices(
        faces.shape[0]
    )

    if dual_edges.numel() > 0:
        graph.add_edges(
            dual_edges.detach().cpu().numpy()
        )

        graph.es['weight'] = (
            dual_edges_weights
            .detach()
            .cpu()
            .numpy()
            .tolist()
        )
    else:
        graph.es['weight'] = []

    graph.add_vertex('s')
    graph.add_vertex('t')

    # Connect invisible faces to source
    if inner_face_indices.numel() > 0:
        graph.add_edges(
            [
                (int(face_index), 's')
                for face_index
                in inner_face_indices.tolist()
            ],
            attributes={
                'weight': np.ones(
                    inner_face_indices.shape[0],
                    dtype=np.float32,
                )
            },
        )

    # Connect outer faces to target
    if outer_face_indices.numel() > 0:
        graph.add_edges(
            [
                (int(face_index), 't')
                for face_index
                in outer_face_indices.tolist()
            ],
            attributes={
                'weight': np.ones(
                    outer_face_indices.shape[0],
                    dtype=np.float32,
                )
            },
        )

    mincut_weights = (
        np.asarray(
            graph.es['weight'],
            dtype=np.float64,
        )
        * 1000.0
    ).tolist()

    cut = graph.mincut(
        's',
        't',
        mincut_weights,
    )

    remove_face_indices = torch.tensor(
        [
            vertex_id
            for vertex_id in cut.partition[0]
            if vertex_id < faces.shape[0]
        ],
        dtype=torch.long,
        device=faces.device,
    )

    if verbose:
        tqdm.write(
            'Mincut solved, start checking the cut'
        )

    # ============================================================
    # Validate each connected component selected by mincut
    # ============================================================

    if remove_face_indices.numel() > 0:
        to_remove_components = (
            utils3d.torch.compute_connected_components(
                faces[remove_face_indices]
            )
        )
    else:
        to_remove_components = []

    if debug:
        tqdm.write(
            'Number of connected components of the cut: '
            f'{len(to_remove_components)}'
        )

    valid_remove_components = []
    cutting_edges = []

    for component in to_remove_components:
        selected_face_indices = remove_face_indices[
            component
        ]

        component_visibility_median = visibility[
            selected_face_indices
        ].median()

        if debug:
            tqdm.write(
                'visibility_median: '
                f'{component_visibility_median.item()}'
            )

        # 可见性过高，不删除
        if component_visibility_median > 0.25:
            continue

        # 检查切割形成的新边界环是否足够小
        (
            component_edge_indices,
            component_edge_counts,
        ) = torch.unique(
            face2edge[selected_face_indices],
            return_counts=True,
        )

        component_boundary_edge_indices = (
            component_edge_indices[
                component_edge_counts == 1
            ]
        )

        component_new_boundary_edge_indices = (
            component_boundary_edge_indices[
                ~torch.isin(
                    component_boundary_edge_indices,
                    boundary_edge_indices,
                )
            ]
        )

        if component_new_boundary_edge_indices.numel() > 0:
            new_boundary_components = (
                utils3d.torch
                .compute_edge_connected_components(
                    edges[
                        component_new_boundary_edge_indices
                    ]
                )
            )

            component_areas = []

            for edge_component in new_boundary_components:
                edge_indices = (
                    component_new_boundary_edge_indices[
                        edge_component
                    ]
                )

                edge_vertex_indices = edges[
                    edge_indices
                ]

                edge_center = verts[
                    edge_vertex_indices
                ].mean(dim=1).mean(dim=0)

                edge_vector_1 = (
                    verts[edge_vertex_indices[:, 0]]
                    - edge_center
                )

                edge_vector_2 = (
                    verts[edge_vertex_indices[:, 1]]
                    - edge_center
                )

                loop_area = (
                    torch.linalg.vector_norm(
                        torch.cross(
                            edge_vector_1,
                            edge_vector_2,
                            dim=-1,
                        ),
                        dim=1,
                    ).sum()
                    * 0.5
                )

                component_areas.append(loop_area)

            if debug:
                cutting_edges.append(
                    component_new_boundary_edge_indices
                )

                tqdm.write(
                    'Area of the cutting loop: '
                    f'{[float(area.item()) for area in component_areas]}'
                )

            if any(
                float(area.item()) > max_hole_size
                for area in component_areas
            ):
                continue

        valid_remove_components.append(
            component
        )

    # ============================================================
    # Debug outputs
    # ============================================================

    if debug:
        face_centers = (
            verts[faces]
            .mean(dim=1)
            .detach()
            .cpu()
            .numpy()
        )

        visualization_dual_edges = (
            dual_edges.detach().cpu().numpy()
        )

        visualization_colors = np.zeros(
            (faces.shape[0], 3),
            dtype=np.uint8,
        )

        visualization_colors[
            inner_face_indices.detach().cpu().numpy()
        ] = [0, 0, 255]

        visualization_colors[
            outer_face_indices.detach().cpu().numpy()
        ] = [0, 255, 0]

        visualization_colors[
            remove_face_indices.detach().cpu().numpy()
        ] = [255, 0, 255]

        if len(valid_remove_components) > 0:
            valid_component_indices = torch.cat(
                valid_remove_components
            )

            visualization_colors[
                remove_face_indices[
                    valid_component_indices
                ].detach().cpu().numpy()
            ] = [255, 0, 0]

        utils3d.io.write_ply(
            'dbg_dual.ply',
            face_centers,
            edges=visualization_dual_edges,
            vertex_colors=visualization_colors,
        )

        if len(cutting_edges) > 0:
            debug_cutting_edge_indices = torch.cat(
                cutting_edges
            )

            visualization_vertices = (
                verts.detach().cpu().numpy()
            )

            visualization_edges = (
                edges[debug_cutting_edge_indices]
                .detach()
                .cpu()
                .numpy()
            )

            utils3d.io.write_ply(
                'dbg_cut.ply',
                visualization_vertices,
                edges=visualization_edges,
            )

    # ============================================================
    # Remove validated mincut components
    # ============================================================

    if len(valid_remove_components) > 0:
        valid_component_indices = torch.cat(
            valid_remove_components
        )

        remove_face_indices = remove_face_indices[
            valid_component_indices
        ]

        if remove_face_indices.numel() > 0:
            keep_face_mask = torch.ones(
                faces.shape[0],
                dtype=torch.bool,
                device=faces.device,
            )

            keep_face_mask[
                remove_face_indices
            ] = False

            faces = faces[
                keep_face_mask
            ]

            # 该函数只删除不再被任何面引用的顶点，
            # 不会主动删除仍存在的面。
            (
                faces,
                verts,
            ) = utils3d.torch.remove_unreferenced_vertices(
                faces,
                verts,
            )

            if verbose:
                tqdm.write(
                    f'Removed '
                    f'{int((~keep_face_mask).sum().item())} '
                    f'faces by mincut'
                )
        elif verbose:
            tqdm.write(
                'Removed 0 faces by mincut'
            )
    elif verbose:
        tqdm.write(
            'Removed 0 faces by mincut'
        )

    # ============================================================
    # MeshFix hole filling
    # ============================================================

    verts, faces = _run_meshfix_hole_filling(
        verts=verts,
        faces=faces,
        max_hole_nbe=max_hole_nbe,
    )

    # 恢复提前保存的退化面
    return finalize_output(
        verts,
        faces,
    )

if __name__ == '__main__':
    import trimesh
    input_mesh_path = '/root/junkai_lin/mul_res_diffusion/pretrained_model/gobjaverse_fix3_octant_factorized_1_to_1024_ce_altpc_decoded_deg_w1_hg246/single_layer_ep28_res32to64_s20/000-002_21538150cde042bd87a0b0a75935c5a1_facenum_8572/pred_consensus.obj'
    output_mesh_path = '/root/junkai_lin/mul_res_diffusion/pretrained_model/gobjaverse_fix3_octant_factorized_1_to_1024_ce_altpc_decoded_deg_w1_hg246/single_layer_ep28_res32to64_s20/000-002_21538150cde042bd87a0b0a75935c5a1_facenum_8572/pred_consensus_fix2.obj'
    mesh = trimesh.load(input_mesh_path, force='mesh')
    verts = torch.tensor(mesh.vertices, device='cuda', dtype=torch.float32)
    faces = torch.tensor(mesh.faces, device='cuda', dtype=torch.int32)
    new_verts, new_faces = _fill_holes(verts, faces, debug=True, verbose=True, only_fill_holes=True, keep_degenerate_faces=True)
    new_mesh = trimesh.Trimesh(vertices=new_verts.cpu().numpy(), faces=new_faces.cpu().numpy())
    new_mesh.export(output_mesh_path)
    # from pathlib import Path

    # import torch
    # import trimesh

    # input_dir = Path("/root/junkai_lin/experiment/render/result_control/t-flow-ab")
    # output_dir = Path("/root/junkai_lin/experiment/render/result_control/t-flow-ab_fix_holes")
    # output_dir.mkdir(parents=True, exist_ok=True)

    # obj_paths = sorted(input_dir.glob("*.obj"))

    # print(f"Found {len(obj_paths)} obj files in {input_dir}")

    # for obj_path in obj_paths:
    #     print(f"\nProcessing: {obj_path.name}")

    #     try:
    #         mesh = trimesh.load(str(obj_path), force="mesh")

    #         verts = torch.tensor(mesh.vertices, device="cuda", dtype=torch.float32)
    #         faces = torch.tensor(mesh.faces, device="cuda", dtype=torch.int32)

    #         new_verts, new_faces = _fill_holes(
    #             verts,
    #             faces,
    #             debug=True,
    #             verbose=True,
    #             only_fill_holes=True,
    #             max_hole_nbe=64,
    #             resolution=256,
    #             max_hole_size=0.12,
    #             num_views=1000
    #         )

    #         new_mesh = trimesh.Trimesh(
    #             vertices=new_verts.detach().cpu().numpy(),
    #             faces=new_faces.detach().cpu().numpy(),
    #             process=False,
    #         )

    #         out_path = output_dir / f"{obj_path.stem}_fix2.obj"
    #         new_mesh.export(str(out_path))

    #         print(f"Saved: {out_path}")

    #     except Exception as e:
    #         print(f"Failed: {obj_path.name}")
    #         print(f"Error: {e}")