geMoldInsight/src/moldinsight/core/mesh_generator.py

# src/core/mesh_generator.py
import logging
import numpy as np
from typing import Dict, List, Optional, Any
import trimesh
from trimesh import sample as trimesh_sample
from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
from OCC.Core.TopExp import TopExp_Explorer
from OCC.Core.TopAbs import TopAbs_FACE
from OCC.Core.BRep import BRep_Tool
from OCC.Core.TopoDS import TopoDS_Shape
from OCC.Core.TopLoc import TopLoc_Location

logger = logging.getLogger(__name__)


class MeshGenerator:
    """网格生成器 - 从PythonOCC形状生成点云，支持多级LOD"""

    def __init__(self, quality: str = "medium"):
        self.quality_settings = {
            "low": 1.0,
            "medium": 0.3,
            "high": 0.1
        }
        self.quality = self.quality_settings.get(quality, 0.3)

    def generate_mesh_from_shape(self, shape: TopoDS_Shape, num_points: int = 20000) -> Dict:
        """从PythonOCC形状生成点云数据"""
        try:
            mesh = BRepMesh_IncrementalMesh(shape, self.quality, False, 0.5, True)
            mesh.Perform()
            logger.info(f"OCC网格生成完成， 网格状态: {mesh.IsDone()}")

            all_vertices = []
            all_faces = []
            vertex_offset = 0

            explorer = TopExp_Explorer(shape, TopAbs_FACE)
            face_count = 0

            while explorer.More():
                face = explorer.Current()
                face_count += 1

                location = TopLoc_Location()
                face_triangulation = BRep_Tool.Triangulation(face, location)

                if face_triangulation is None:
                    logger.warning(f"面 {face_count} 没有三角剖分数据")
                    explorer.Next()
                    continue

                trsf = location.Transformation()

                nb_nodes = face_triangulation.NbNodes()
                nb_triangles = face_triangulation.NbTriangles()

                logger.debug(f"面 {face_count}: {nb_nodes} 个顶点, {nb_triangles} 个三角形")

                face_vertices = []
                for i in range(1, nb_nodes + 1):
                    pnt = face_triangulation.Node(i)
                    transformed = pnt.Transformed(trsf)
                    face_vertices.append([
                        float(transformed.X()),
                        float(transformed.Y()),
                        float(transformed.Z()),
                    ])

                face_indices = []
                for i in range(1, nb_triangles + 1):
                    tri = face_triangulation.Triangle(i)
                    idx1 = tri.Value(1)
                    idx2 = tri.Value(2)
                    idx3 = tri.Value(3)
                    face_indices.append([
                        vertex_offset + idx1 - 1,
                        vertex_offset + idx2 - 1,
                        vertex_offset + idx3 - 1
                    ])

                all_vertices.extend(face_vertices)
                all_faces.extend(face_indices)
                vertex_offset += len(face_vertices)

                explorer.Next()

            if len(all_vertices) == 0:
                logger.warning("未提取到任何顶点，使用示例数据")
                return self._create_sample_pointcloud()

            vertices = np.array(all_vertices, dtype=np.float32)
            faces = np.array(all_faces, dtype=np.int32)

            logger.info(f"总共提取了 {len(vertices)} 个顶点, {len(faces)} 个三角形面, {face_count} 个面")

            tri_mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=True)

            actual_num_points = min(num_points, len(faces) * 2)
            logger.info(f"采样点数: {actual_num_points}")

            points, face_idx = trimesh.sample.sample_surface(tri_mesh, actual_num_points)

            normals = tri_mesh.face_normals[face_idx]

            logger.info(f"生成了 {len(points)} 个点云点")

            return {
                "vertices": vertices.tolist(),
                "faces": faces.tolist(),
                "points": points.tolist(),
                "normals": normals.tolist(),
                "point_count": int(len(points)),
                "vertex_count": int(len(vertices)),
                "face_count": int(len(faces))
            }

        except Exception as e:
            logger.error(f"网格生成失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return self._create_sample_pointcloud()

    def generate_multi_lod_mesh(self, shape: TopoDS_Shape) -> Dict:
        """生成多级LOD网格 - 一次OCC剖分，trimesh简化，避免重复计算

        返回结构:
        {
            "lods": {
                "0": { "vertices": [...], "faces": [...], "vertex_count": N, "face_count": N },
                "1": { ... 50%简化 ... },
                "2": { ... 80%简化 ... }
            },
            "points": [...], "normals": [...], "point_count": N,
            "vertex_count": N, "face_count": N
        }
        """
        try:
            full_mesh_result = self.generate_mesh_from_shape(shape, num_points=20000)

            vertices = np.array(full_mesh_result["vertices"], dtype=np.float32)
            faces = np.array(full_mesh_result["faces"], dtype=np.int32)

            if len(vertices) == 0 or len(faces) == 0:
                sample = self._create_sample_pointcloud()
                return self._wrap_sample_as_lod(sample)

            tri_mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=True)
            full_face_count = len(tri_mesh.faces)
            logger.info(f"全精度网格: {len(tri_mesh.vertices)} 顶点, {full_face_count} 面")

            lods = {
                "0": self._mesh_to_lod_entry(tri_mesh, "LOD0-全精度")
            }

            lod_ratios = {"1": 0.50, "2": 0.20}
            for lod_level, ratio in lod_ratios.items():
                if full_face_count < 300:
                    lods[lod_level] = lods["0"]
                    continue

                target_faces = max(int(full_face_count * ratio), 200)
                try:
                    simplified = tri_mesh.simplify_quadric_decimation(target_faces)
                    if simplified is None or len(simplified.faces) < 3:
                        simplified = self._fast_decimate(tri_mesh, target_faces)
                    lods[lod_level] = self._mesh_to_lod_entry(simplified, f"LOD{lod_level}-简化{int((1-ratio)*100)}%")
                    logger.info(f"LOD{lod_level}: {len(simplified.vertices)} 顶点, {len(simplified.faces)} 面 (目标{target_faces})")
                except Exception as dec_err:
                    logger.warning(f"LOD{lod_level} 简化失败，回退到全精度: {dec_err}")
                    lods[lod_level] = lods["0"]

            result = {
                "lods": lods,
                "points": full_mesh_result["points"],
                "normals": full_mesh_result["normals"],
                "point_count": full_mesh_result["point_count"],
                "vertex_count": full_mesh_result["vertex_count"],
                "face_count": full_mesh_result["face_count"],
            }
            return result

        except Exception as e:
            logger.error(f"多级LOD网格生成失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            sample = self._create_sample_pointcloud()
            return self._wrap_sample_as_lod(sample)

    def _mesh_to_lod_entry(self, mesh: trimesh.Trimesh, label: str) -> Dict:
        return {
            "vertices": mesh.vertices.tolist(),
            "faces": mesh.faces.tolist(),
            "vertex_count": int(len(mesh.vertices)),
            "face_count": int(len(mesh.faces)),
        }

    def _fast_decimate(self, mesh: trimesh.Trimesh, target_faces: int) -> trimesh.Trimesh:
        """快速回退降采样：按面索引均匀采样"""
        if target_faces >= len(mesh.faces):
            return mesh
        step = max(len(mesh.faces) // target_faces, 1)
        indices = np.arange(0, len(mesh.faces), step)[:target_faces]
        return mesh.submesh([np.array(indices)], only_watertight=False, append=True)

    def _wrap_sample_as_lod(self, sample: Dict) -> Dict:
        lods = {
            "0": {
                "vertices": sample["vertices"],
                "faces": sample["faces"],
                "vertex_count": sample["vertex_count"],
                "face_count": sample["face_count"],
            }
        }
        lods["1"] = lods["0"]
        lods["2"] = lods["0"]
        return {
            "lods": lods,
            "points": sample["points"],
            "normals": sample["normals"],
            "point_count": sample["point_count"],
            "vertex_count": sample["vertex_count"],
            "face_count": sample["face_count"],
        }

    def _create_sample_pointcloud(self) -> Dict:
        """创建示例点云（备用）"""
        mesh = trimesh.creation.box([100, 80, 50])
        points, _ = trimesh.sample.sample_surface(mesh, 5000)
        normals = mesh.face_normals[:len(points)]

        return {
            "vertices": mesh.vertices.tolist(),
            "faces": mesh.faces.tolist(),
            "points": points.tolist(),
            "normals": normals.tolist(),
            "point_count": len(points),
            "vertex_count": len(mesh.vertices),
            "face_count": len(mesh.faces)
        }