Spaces:
Sleeping
Sleeping
File size: 5,559 Bytes
182efca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """STEP file loading with XDE support and Shape Healing."""
from __future__ import annotations
import logging
from dataclasses import dataclass, field
from pathlib import Path
from OCP.IFSelect import IFSelect_RetDone
from OCP.STEPCAFControl import STEPCAFControl_Reader
from OCP.STEPControl import STEPControl_Reader
from OCP.ShapeFix import ShapeFix_Shape, ShapeFix_ShapeTolerance
from OCP.TCollection import TCollection_ExtendedString
from OCP.TDocStd import TDocStd_Document
from OCP.TopoDS import TopoDS_Shape
from OCP.XCAFApp import XCAFApp_Application
from OCP.XCAFDoc import XCAFDoc_ShapeTool
from OCP.TDF import TDF_LabelSequence
logger = logging.getLogger(__name__)
@dataclass
class StepFileInfo:
"""Metadata about a loaded STEP file."""
path: Path
shape: TopoDS_Shape
doc: TDocStd_Document | None = None
shape_tool: XCAFDoc_ShapeTool | None = None
reader: STEPControl_Reader | None = None
protocol: str = ""
unit: str = "mm"
num_roots: int = 0
extra: dict = field(default_factory=dict)
def _heal_shape(shape: TopoDS_Shape, tolerance: float = 0.01) -> TopoDS_Shape:
"""Apply shape healing to fix geometry issues."""
fixer = ShapeFix_Shape(shape)
fixer.SetPrecision(tolerance)
fixer.SetMaxTolerance(tolerance * 10)
fixer.Perform()
tol_fixer = ShapeFix_ShapeTolerance()
tol_fixer.SetTolerance(fixer.Shape(), tolerance)
logger.info("Shape healing completed")
return fixer.Shape()
def _detect_protocol(reader: STEPControl_Reader) -> str:
"""Detect AP protocol from STEP file."""
try:
ws = reader.WS()
model = ws.Model()
if model is not None:
header = str(model.Header()) if hasattr(model, "Header") else ""
for ap in ("AP214", "AP203", "AP242"):
if ap.lower() in header.lower() or ap in header:
return ap
except Exception:
pass
return "unknown"
def _load_xde(file_path: Path) -> tuple[TDocStd_Document, XCAFDoc_ShapeTool, str]:
"""Load STEP into XDE document and return (doc, shape_tool, protocol)."""
app = XCAFApp_Application.GetApplication_s()
doc = TDocStd_Document(TCollection_ExtendedString("MDTV-XCAF"))
app.InitDocument(doc)
xde_reader = STEPCAFControl_Reader()
xde_reader.SetNameMode(True)
xde_reader.SetColorMode(True)
xde_reader.SetLayerMode(True)
status = xde_reader.ReadFile(str(file_path))
if status != IFSelect_RetDone:
raise RuntimeError(f"Failed to read STEP file: {file_path} (status={status})")
protocol = _detect_protocol(xde_reader.Reader())
if not xde_reader.Transfer(doc):
raise RuntimeError(f"Failed to transfer STEP data: {file_path}")
shape_tool = XCAFDoc_ShapeTool.Set_s(doc.Main())
return doc, shape_tool, protocol
def _load_basic(file_path: Path) -> tuple[TopoDS_Shape, str, int, STEPControl_Reader]:
"""Fallback: load with basic STEPControl_Reader."""
reader = STEPControl_Reader()
status = reader.ReadFile(str(file_path))
if status != IFSelect_RetDone:
raise RuntimeError(f"Failed to read STEP file: {file_path}")
protocol = _detect_protocol(reader)
num_roots = reader.NbRootsForTransfer()
reader.TransferRoots()
shape = reader.OneShape()
return shape, protocol, num_roots, reader
def load_step(file_path: str | Path, heal: bool = True) -> StepFileInfo:
"""Load a STEP file, trying XDE reader first then falling back to basic.
Args:
file_path: Path to the .stp/.step file.
heal: Whether to apply shape healing.
Returns:
StepFileInfo with shape, XDE document (if available), and metadata.
"""
file_path = Path(file_path)
if not file_path.exists():
raise FileNotFoundError(f"STEP file not found: {file_path}")
logger.info("Loading STEP file: %s (%.1f MB)", file_path, file_path.stat().st_size / 1e6)
doc = None
shape_tool = None
shape = None
basic_reader = None
protocol = "unknown"
num_roots = 0
# Try XDE reader for assembly structure
try:
doc, shape_tool, protocol = _load_xde(file_path)
labels = TDF_LabelSequence()
shape_tool.GetFreeShapes(labels)
num_roots = labels.Length()
logger.info("XDE reader: %d free shape(s), protocol=%s", num_roots, protocol)
if num_roots > 0:
if num_roots == 1:
shape = shape_tool.GetShape(labels.Value(1))
else:
from OCP.BRep import BRep_Builder
from OCP.TopoDS import TopoDS_Compound
builder = BRep_Builder()
compound = TopoDS_Compound()
builder.MakeCompound(compound)
for i in range(1, num_roots + 1):
builder.Add(compound, shape_tool.GetShape(labels.Value(i)))
shape = compound
except Exception as e:
logger.warning("XDE reader failed: %s", e)
# Fallback to basic reader if XDE produced no shapes
if shape is None or shape.IsNull():
logger.info("Falling back to basic STEPControl_Reader")
shape, protocol, num_roots, basic_reader = _load_basic(file_path)
if shape.IsNull():
raise RuntimeError(f"Empty shape in STEP file: {file_path}")
if heal:
shape = _heal_shape(shape)
return StepFileInfo(
path=file_path,
shape=shape,
doc=doc,
shape_tool=shape_tool,
reader=basic_reader,
protocol=protocol,
num_roots=num_roots,
)
|