547 lines
19 KiB
Python
547 lines
19 KiB
Python
"""
|
||
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,
|
||
)
|