""" AI 分型面检测模块 - 基于 GNN 的分型面预测框架 架构设计： 1. ShapeGraphBuilder - 将 OCC 形状转换为图表示（面为节点，共享边为图边） 2. PartingSurfaceGNN - 图神经网络模型定义 3. AIPartingSurfaceDetectorV2 - 增强版分型面检测器（集成 GNN） 图构建策略： - 节点：每个 TopoDS_Face 作为一个节点 - 节点特征：法向量(3) + 面积(1) + 曲率(2) + 面类型(1) = 7维 - 边：共享 TopoDS_Edge 的面之间建立边 - 边特征：共享边长度(1) + 二面角(1) = 2维 GNN 模型： - 3层 GraphConv + 全局池化 + MLP 分类头 - 输出：每个面的分型面归属概率 + 分型方向 依赖： - PyTorch + PyTorch Geometric（可选，缺失时回退到几何方法） """ from typing import Dict, List, Any, Optional, Tuple import numpy as np from utils.logger import get_logger logger = get_logger(__name__) _TORCH_AVAILABLE = False _TORCH_GEOMETRIC_AVAILABLE = False try: import torch import torch.nn as nn import torch.nn.functional as F _TORCH_AVAILABLE = True try: from torch_geometric.nn import GCNConv, global_mean_pool from torch_geometric.data import Data _TORCH_GEOMETRIC_AVAILABLE = True except ImportError: logger.info("PyTorch Geometric 未安装，GNN 模型不可用") except ImportError: logger.info("PyTorch 未安装，AI 分型面检测将使用几何回退方法") class ShapeGraphBuilder: """将 OCC 形状转换为图表示""" def build_graph(self, shape: Any) -> Optional[Dict]: """ 从 OCC 形状构建图数据 Returns: { "node_features": np.ndarray (N, 7), "edge_index": np.ndarray (2, E), "edge_features": np.ndarray (E, 2), "face_map": List[TopoDS_Face], "num_nodes": int, "num_edges": int } """ try: from OCC.Core.TopExp import TopExp_Explorer from OCC.Core.TopAbs import TopAbs_FACE, TopAbs_EDGE from OCC.Core.BRepAdaptor import BRepAdaptor_Surface from OCC.Core.GProp import GProp_GProps from OCC.Core.BRepGProp import brepgprop from OCC.Core.Bnd import Bnd_Box from OCC.Core.BRepBndLib import brepbndlib_Add from OCC.Core.TopTools import TopTools_IndexedDataMapOfShapeListOfShape from OCC.Core.TopExp import topexp_MapShapesAndAncestors from OCC.Core.TopoDS import TopoDS_Face, TopoDS_Edge, topods faces = [] face_features = [] explorer = TopExp_Explorer(shape, TopAbs_FACE) while explorer.More(): face = topods.Face(explorer.Current()) features = self._extract_face_features(face) if features is not None: faces.append(face) face_features.append(features) explorer.Next() if not faces: logger.warning("未找到面，无法构建图") return None node_features = np.array(face_features, dtype=np.float32) edge_map = TopTools_IndexedDataMapOfShapeListOfShape() topexp_MapShapesAndAncestors(shape, TopAbs_EDGE, TopAbs_FACE, edge_map) edge_list = [] edge_features_list = [] for i in range(1, edge_map.Extent() + 1): edge = topods.Edge(edge_map.FindKey(i)) face_list = edge_map.FindFromIndex(i) connected_faces = [] it = face_list.begin() while it != face_list.end(): f = topods.Face(it.Value()) try: idx = faces.index(f) connected_faces.append(idx) except ValueError: pass it.next_ptr() if len(connected_faces) >= 2: edge_feat = self._extract_edge_features(edge, connected_faces, faces) for j in range(len(connected_faces)): for k in range(j + 1, len(connected_faces)): edge_list.append([connected_faces[j], connected_faces[k]]) edge_features_list.append(edge_feat) if not edge_list: logger.warning("未找到边连接，返回无图边的图") edge_index = np.zeros((2, 0), dtype=np.int64) edge_features_arr = np.zeros((0, 2), dtype=np.float32) else: edge_index = np.array(edge_list, dtype=np.int64).T rev_edges = np.array([[e[1], e[0]] for e in edge_list], dtype=np.int64).T edge_index = np.concatenate([edge_index, rev_edges], axis=1) edge_features_arr = np.array(edge_features_list, dtype=np.float32) edge_features_arr = np.concatenate([edge_features_arr, edge_features_arr], axis=0) return { "node_features": node_features, "edge_index": edge_index, "edge_features": edge_features_arr, "face_map": faces, "num_nodes": len(faces), "num_edges": edge_index.shape[1] } except Exception as e: logger.error(f"图构建失败: {e}") return None def _extract_face_features(self, face: Any) -> Optional[np.ndarray]: """ 提取面特征：[nx, ny, nz, area, u_curvature, v_curvature, face_type] """ try: from OCC.Core.BRepAdaptor import BRepAdaptor_Surface from OCC.Core.GProp import GProp_GProps from OCC.Core.BRepGProp import brepgprop surface = BRepAdaptor_Surface(face) u = (surface.FirstUParameter() + surface.LastUParameter()) / 2 v = (surface.FirstVParameter() + surface.LastVParameter()) / 2 if surface.GetType() == 0: normal = surface.Plane().Position().Direction() face_type = 0.0 u_curv = 0.0 v_curv = 0.0 elif surface.GetType() == 1: normal = surface.Cylinder().Position().Direction() face_type = 1.0 radius = surface.Cylinder().Radius() u_curv = 1.0 / radius if radius > 0.001 else 0.0 v_curv = 0.0 elif surface.GetType() == 2: normal = surface.Cone().Position().Direction() face_type = 2.0 u_curv = 0.0 v_curv = 0.0 elif surface.GetType() == 3: normal = surface.Sphere().Position().Direction() face_type = 3.0 radius = surface.Sphere().Radius() u_curv = 1.0 / radius if radius > 0.001 else 0.0 v_curv = 1.0 / radius if radius > 0.001 else 0.0 elif surface.GetType() == 4: normal = surface.Torus().Position().Direction() face_type = 4.0 u_curv = 0.0 v_curv = 0.0 else: from OCC.Core.BRepLProp import BRepLProp_SLProps props = BRepLProp_SLProps(surface, 2, 0.001) props.SetParameters(u, v) if props.IsNormalDefined(): normal = props.Normal() else: normal = gp_Dir(0, 0, 1) face_type = 5.0 u_curv = 0.0 v_curv = 0.0 face_props = GProp_GProps() brepgprop.SurfaceProperties(face, face_props) area = face_props.Mass() return np.array([ normal.X(), normal.Y(), normal.Z(), area, u_curv, v_curv, face_type ], dtype=np.float32) except Exception as e: logger.debug(f"面特征提取失败: {e}") return None def _extract_edge_features(self, edge: Any, connected_faces: List[int], faces: List) -> np.ndarray: """ 提取边特征：[edge_length, dihedral_angle] """ try: from OCC.Core.BRepAdaptor import BRepAdaptor_Curve from OCC.Core.GProp import GProp_GProps from OCC.Core.BRepGProp import brepgprop curve = BRepAdaptor_Curve(edge) first = curve.FirstParameter() last = curve.LastParameter() edge_len = abs(last - first) dihedral = 0.0 if len(connected_faces) >= 2: n1 = self._get_face_normal_fast(faces[connected_faces[0]]) n2 = self._get_face_normal_fast(faces[connected_faces[1]]) if n1 is not None and n2 is not None: dot = np.clip(np.dot(n1, n2), -1.0, 1.0) dihedral = np.arccos(dot) return np.array([edge_len, dihedral], dtype=np.float32) except Exception: return np.array([0.0, 0.0], dtype=np.float32) def _get_face_normal_fast(self, face: Any) -> Optional[np.ndarray]: """快速获取面法向量（numpy数组）""" try: from OCC.Core.BRepAdaptor import BRepAdaptor_Surface surface = BRepAdaptor_Surface(face) if surface.GetType() == 0: n = surface.Plane().Position().Direction() return np.array([n.X(), n.Y(), n.Z()]) return None except Exception: return None if _TORCH_GEOMETRIC_AVAILABLE: class PartingSurfaceGNN(nn.Module): """ 分型面检测 GNN 模型 架构： - 3层 GCNConv (hidden_dim=64) - 全局平均池化 - 3层 MLP 分类头 - 输出：每个面的分型面归属概率 (0-1) """ def __init__(self, input_dim: int = 7, hidden_dim: int = 64, num_layers: int = 3, dropout: float = 0.3): super().__init__() self.input_dim = input_dim self.hidden_dim = hidden_dim self.num_layers = num_layers self.input_proj = nn.Linear(input_dim, hidden_dim) self.convs = nn.ModuleList() self.bns = nn.ModuleList() for _ in range(num_layers): self.convs.append(GCNConv(hidden_dim, hidden_dim)) self.bns.append(nn.BatchNorm1d(hidden_dim)) self.dropout = dropout self.mlp = nn.Sequential( nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden_dim, hidden_dim // 2), nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden_dim // 2, 1), ) def forward(self, data: Data) -> torch.Tensor: x, edge_index = data.x, data.edge_index x = self.input_proj(x) x = F.relu(x) for conv, bn in zip(self.convs, self.bns): x = conv(x, edge_index) x = bn(x) x = F.relu(x) x = F.dropout(x, p=self.dropout, training=self.training) out = self.mlp(x) return torch.sigmoid(out).squeeze(-1) class PartingDirectionHead(nn.Module): """ 分型方向预测头 基于全局池化的面特征，预测分型方向向量 """ def __init__(self, hidden_dim: int = 64): super().__init__() self.direction_mlp = nn.Sequential( nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 3), ) def forward(self, node_embeddings: torch.Tensor, batch: torch.Tensor) -> torch.Tensor: pooled = global_mean_pool(node_embeddings, batch) direction = self.direction_mlp(pooled) direction = F.normalize(direction, p=2, dim=-1) return direction class AIPartingSurfaceDetectorV2: """ 增强版 AI 分型面检测器 支持： 1. GNN 模型推理（需要 PyTorch + PyG） 2. 几何方法回退（无需任何 AI 依赖） 3. 模型训练数据收集 """ def __init__(self, model_path: Optional[str] = None, use_gnn: bool = True, device: str = "cpu"): self.model = None self.direction_head = None self.graph_builder = ShapeGraphBuilder() self.device = device self.use_gnn = use_gnn and _TORCH_GEOMETRIC_AVAILABLE if model_path and self.use_gnn: self._load_model(model_path) def _load_model(self, model_path: str): """加载训练好的 GNN 模型""" if not _TORCH_GEOMETRIC_AVAILABLE: logger.warning("PyTorch Geometric 不可用，无法加载 GNN 模型") return try: checkpoint = torch.load(model_path, map_location=self.device) self.model = PartingSurfaceGNN( input_dim=checkpoint.get("input_dim", 7), hidden_dim=checkpoint.get("hidden_dim", 64), ) self.model.load_state_dict(checkpoint["model_state_dict"]) self.model.to(self.device) self.model.eval() if "direction_head_state_dict" in checkpoint: self.direction_head = PartingDirectionHead( hidden_dim=checkpoint.get("hidden_dim", 64) ) self.direction_head.load_state_dict(checkpoint["direction_head_state_dict"]) self.direction_head.to(self.device) self.direction_head.eval() logger.info(f"GNN 模型加载成功: {model_path}") except Exception as e: logger.error(f"GNN 模型加载失败: {e}") self.model = None def detect(self, product_shape: Any, analysis: Dict) -> Optional[Dict]: """ 检测最优分型面 Args: product_shape: OpenCASCADE 形状对象 analysis: 几何分析结果 Returns: { "origin": [x, y, z], "normal": [nx, ny, nz], "confidence": float, "parting_line": [...], "method": "gnn" | "geometric" } """ if self.use_gnn and self.model is not None: result = self._detect_with_gnn(product_shape, analysis) if result is not None: return result return self._detect_with_geometry(product_shape, analysis) def _detect_with_gnn(self, shape: Any, analysis: Dict) -> Optional[Dict]: """使用 GNN 模型检测分型面""" if not _TORCH_GEOMETRIC_AVAILABLE: return None try: graph_data = self.graph_builder.build_graph(shape) if graph_data is None: return None node_features = torch.tensor( graph_data["node_features"], dtype=torch.float32 ).to(self.device) edge_index = torch.tensor( graph_data["edge_index"], dtype=torch.long ).to(self.device) data = Data(x=node_features, edge_index=edge_index) with torch.no_grad(): face_probs = self.model(data) if self.direction_head is not None: batch = torch.zeros( data.num_nodes, dtype=torch.long, device=self.device ) direction = self.direction_head(data.x, batch) normal = direction.cpu().numpy().tolist() else: normal = [0, 0, 1] parting_face_mask = face_probs.cpu().numpy() > 0.5 confidence = float(face_probs.mean().cpu().numpy()) bbox = analysis.get("bounding_box", {}) center = bbox.get("center", [0, 0, 0]) return { "origin": center, "normal": normal, "confidence": confidence, "method": "gnn", "face_probabilities": face_probs.cpu().numpy().tolist(), "parting_face_count": int(parting_face_mask.sum()), } except Exception as e: logger.warning(f"GNN 检测失败，回退到几何方法: {e}") return None def _detect_with_geometry(self, shape: Any, analysis: Dict) -> Dict: """几何方法回退：基于法向量统计的分型面检测""" try: graph_data = self.graph_builder.build_graph(shape) if graph_data is not None: node_features = graph_data["node_features"] normals = node_features[:, :3] areas = node_features[:, 3] total_area = areas.sum() if total_area > 0: weights = areas / total_area weighted_normal = np.sum(normals * weights[:, np.newaxis], axis=0) else: weighted_normal = np.mean(normals, axis=0) length = np.linalg.norm(weighted_normal) if length > 0.001: weighted_normal /= length else: weighted_normal = np.array([0, 0, 1]) dot_products = np.abs(np.dot(normals, weighted_normal)) confidence = float(np.mean(dot_products)) bbox = analysis.get("bounding_box", {}) center = bbox.get("center", [0, 0, 0]) return { "origin": center, "normal": weighted_normal.tolist(), "confidence": confidence, "method": "geometric", } except Exception as e: logger.warning(f"几何方法检测失败: {e}") bbox = analysis.get("bounding_box", {}) center = bbox.get("center", [0, 0, 0]) return { "origin": center, "normal": [0, 0, 1], "confidence": 0.5, "method": "fallback", } def collect_training_sample(self, shape: Any, analysis: Dict, ground_truth_normal: List[float], ground_truth_origin: List[float]) -> Optional[Dict]: """ 收集训练样本 Args: shape: OCC 形状 analysis: 几何分析 ground_truth_normal: 人工标注的分型方向 ground_truth_origin: 人工标注的分型面原点 Returns: 可序列化的训练样本 """ graph_data = self.graph_builder.build_graph(shape) if graph_data is None: return None return { "node_features": graph_data["node_features"].tolist(), "edge_index": graph_data["edge_index"].tolist(), "edge_features": graph_data["edge_features"].tolist(), "label_normal": ground_truth_normal, "label_origin": ground_truth_origin, "bounding_box": analysis.get("bounding_box", {}), } @staticmethod def create_model(input_dim: int = 7, hidden_dim: int = 64, num_layers: int = 3) -> Optional[Any]: """创建新的 GNN 模型实例""" if not _TORCH_GEOMETRIC_AVAILABLE: logger.warning("PyTorch Geometric 不可用，无法创建模型") return None return PartingSurfaceGNN( input_dim=input_dim, hidden_dim=hidden_dim, num_layers=num_layers, )