Imprinting and matching

cadquery/assembly.py

@@ -543,6 +543,28 @@ def _flatten(self, parents=[]):
 
         return rv
 
+    def __iter__(
+        self,
+        loc: Optional[Location] = None,
+        name: Optional[str] = None,
+        color: Optional[Color] = None,
+    ) -> Iterator[Tuple[Shape, str, Location, Optional[Color]]]:
+        """
+        Assembly iterator yielding shapes, names, locations and colors.
+        """
+
+        name = f"{name}/{self.name}" if name else self.name
+        loc = loc * self.loc if loc else self.loc
+        color = self.color if self.color else color
+
+        if self.obj:
+            yield self.obj if isinstance(self.obj, Shape) else Compound.makeCompound(
+                s for s in self.obj.vals() if isinstance(s, Shape)
+            ), name, loc, color
+
+        for ch in self.children:
+            yield from ch.__iter__(loc, name, color)
+
     def toCompound(self) -> Compound:
         """
         Returns a Compound made from this Assembly (including all children) with the

cadquery/occ_impl/assembly.py

@@ -1,4 +1,15 @@
-from typing import Union, Iterable, Tuple, Dict, overload, Optional, Any, List, cast
+from typing import (
+    Union,
+    Iterable,
+    Iterator,
+    Tuple,
+    Dict,
+    overload,
+    Optional,
+    Any,
+    List,
+    cast,
+)
 from typing_extensions import Protocol
 from math import degrees
 
@@ -12,7 +23,7 @@
 from OCP.Quantity import Quantity_ColorRGBA
 from OCP.BRepAlgoAPI import BRepAlgoAPI_Fuse
 from OCP.TopTools import TopTools_ListOfShape
-from OCP.BOPAlgo import BOPAlgo_GlueEnum
+from OCP.BOPAlgo import BOPAlgo_GlueEnum, BOPAlgo_MakeConnected
 from OCP.TopoDS import TopoDS_Shape
 
 from vtkmodules.vtkRenderingCore import (
@@ -21,8 +32,11 @@
     vtkRenderer,
 )
 
+from vtkmodules.vtkFiltersExtraction import vtkExtractCellsByType
+from vtkmodules.vtkCommonDataModel import VTK_TRIANGLE, VTK_LINE, VTK_VERTEX
+
 from .geom import Location
-from .shapes import Shape, Compound
+from .shapes import Shape, Solid, Compound
 from .exporters.vtk import toString
 from ..cq import Workplane
 
@@ -131,6 +145,14 @@ def children(self) -> Iterable["AssemblyProtocol"]:
     def traverse(self) -> Iterable[Tuple[str, "AssemblyProtocol"]]:
         ...
 
+    def __iter__(
+        self,
+        loc: Optional[Location] = None,
+        name: Optional[str] = None,
+        color: Optional[Color] = None,
+    ) -> Iterator[Tuple[Shape, str, Location, Optional[Color]]]:
+        ...
+
 
 def setName(l: TDF_Label, name: str, tool):
 
@@ -227,75 +249,93 @@ def _toCAF(el, ancestor, color) -> TDF_Label:
 
 def toVTK(
     assy: AssemblyProtocol,
-    renderer: vtkRenderer = vtkRenderer(),
-    loc: Location = Location(),
     color: Tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0),
     tolerance: float = 1e-3,
     angularTolerance: float = 0.1,
 ) -> vtkRenderer:
 
-    loc = loc * assy.loc
-    trans, rot = loc.toTuple()
+    renderer = vtkRenderer()
+
+    for shape, _, loc, col_ in assy:
+
+        col = col_.toTuple() if col_ else color
+        trans, rot = loc.toTuple()
+
+        data = shape.toVtkPolyData(tolerance, angularTolerance)
+
+        # extract faces
+        extr = vtkExtractCellsByType()
+        extr.SetInputDataObject(data)
+
+        extr.AddCellType(VTK_LINE)
+        extr.AddCellType(VTK_VERTEX)
+        extr.Update()
+        data_edges = extr.GetOutput()
+
+        # extract edges
+        extr = vtkExtractCellsByType()
+        extr.SetInputDataObject(data)
 
-    if assy.color:
-        color = assy.color.toTuple()
+        extr.AddCellType(VTK_TRIANGLE)
+        extr.Update()
+        data_faces = extr.GetOutput()
 
-    if assy.shapes:
-        data = Compound.makeCompound(assy.shapes).toVtkPolyData(
-            tolerance, angularTolerance
-        )
+        # remove normals from edges
+        data_edges.GetPointData().RemoveArray("Normals")
 
+        # add both to the renderer
         mapper = vtkMapper()
-        mapper.SetInputData(data)
+        mapper.AddInputDataObject(data_faces)
 
         actor = vtkActor()
         actor.SetMapper(mapper)
         actor.SetPosition(*trans)
         actor.SetOrientation(*map(degrees, rot))
-        actor.GetProperty().SetColor(*color[:3])
-        actor.GetProperty().SetOpacity(color[3])
+        actor.GetProperty().SetColor(*col[:3])
+        actor.GetProperty().SetOpacity(col[3])
 
         renderer.AddActor(actor)
 
-    for child in assy.children:
-        renderer = toVTK(child, renderer, loc, color, tolerance, angularTolerance)
+        mapper = vtkMapper()
+        mapper.AddInputDataObject(data_edges)
+
+        actor = vtkActor()
+        actor.SetMapper(mapper)
+        actor.SetPosition(*trans)
+        actor.SetOrientation(*map(degrees, rot))
+        actor.GetProperty().SetColor(0, 0, 0)
+        actor.GetProperty().SetLineWidth(2)
+
+        renderer.AddActor(actor)
 
     return renderer
 
 
 def toJSON(
     assy: AssemblyProtocol,
-    loc: Location = Location(),
     color: Tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0),
     tolerance: float = 1e-3,
 ) -> List[Dict[str, Any]]:
     """
     Export an object to a structure suitable for converting to VTK.js JSON.
     """
 
-    loc = loc * assy.loc
-    trans, rot = loc.toTuple()
-
-    if assy.color:
-        color = assy.color.toTuple()
-
     rv = []
 
-    if assy.shapes:
+    for shape, _, loc, col_ in assy:
+
         val: Any = {}
 
         data = toString(Compound.makeCompound(assy.shapes), tolerance)
+        trans, rot = loc.toTuple()
 
         val["shape"] = data
-        val["color"] = color
+        val["color"] = col_.toTuple() if col_ else color
         val["position"] = trans
         val["orientation"] = rot
 
         rv.append(val)
 
-    for child in assy.children:
-        rv.extend(toJSON(child, loc, color, tolerance))
-
     return rv
 
 
@@ -331,19 +371,9 @@ def toFusedCAF(
     shapes: List[Shape] = []
     colors = []
 
-    def extract_shapes(assy, parent_loc=None, parent_color=None):
-
-        loc = parent_loc * assy.loc if parent_loc else assy.loc
-        color = assy.color if assy.color else parent_color
-
-        for shape in assy.shapes:
-            shapes.append(shape.moved(loc).copy())
-            colors.append(color)
-
-        for ch in assy.children:
-            extract_shapes(ch, loc, color)
-
-    extract_shapes(assy)
+    for shape, _, loc, color in assy:
+        shapes.append(shape.moved(loc).copy())
+        colors.append(color)
 
     # Initialize with a dummy value for mypy
     top_level_shape = cast(TopoDS_Shape, None)
@@ -411,3 +441,37 @@ def extract_shapes(assy, parent_loc=None, parent_color=None):
                     color_tool.SetColor(cur_lbl, color.wrapped, XCAFDoc_ColorGen)
 
     return top_level_lbl, doc
+
+
+def imprint(assy: AssemblyProtocol) -> Tuple[Shape, Dict[Shape, Tuple[str, ...]]]:
+    """
+    Imprint all the solids and construct a dictionary mapping imprinted solids to names from the input assy.
+    """
+
+    # make the id map
+    id_map = {}
+
+    for obj, name, loc, _ in assy:
+        for s in obj.moved(loc).Solids():
+            id_map[s] = name
+
+    # connect topologically
+    bldr = BOPAlgo_MakeConnected()
+    bldr.SetRunParallel(True)
+    bldr.SetUseOBB(True)
+
+    for obj in id_map:
+        bldr.AddArgument(obj.wrapped)
+
+    bldr.Perform()
+    res = Shape(bldr.Shape())
+
+    # make the connected solid -> id map
+    origins: Dict[Shape, Tuple[str, ...]] = {}
+
+    for s in res.Solids():
+        ids = tuple(id_map[Solid(el)] for el in bldr.GetOrigins(s.wrapped))
+        # if GetOrigins yields nothing, solid was not modified
+        origins[s] = ids if ids else (id_map[s],)
+
+    return res, origins

cadquery/occ_impl/exporters/assembly.py

@@ -140,10 +140,9 @@ def _vtkRenderWindow(
     Convert an assembly to a vtkRenderWindow. Used by vtk based exporters.
     """
 
-    renderer = vtkRenderer()
+    renderer = toVTK(assy, tolerance=tolerance, angularTolerance=angularTolerance)
     renderWindow = vtkRenderWindow()
     renderWindow.AddRenderer(renderer)
-    toVTK(assy, renderer, tolerance=tolerance, angularTolerance=angularTolerance)
 
     renderer.ResetCamera()
     renderer.SetBackground(1, 1, 1)

tests/test_assembly.py

@@ -1475,3 +1475,42 @@ def test_point_constraint(simple_assy2):
     t2 = assy.children[1].loc.wrapped.Transformation().TranslationPart()
 
     assert t2.Modulus() == pytest.approx(1)
+
+
+@pytest.fixture
+def touching_assy():
+
+    b1 = cq.Workplane().box(1, 1, 1)
+    b2 = cq.Workplane(origin=(1, 0, 0)).box(1, 1, 1)
+
+    return cq.Assembly().add(b1).add(b2)
+
+
+@pytest.fixture
+def disjoint_assy():
+
+    b1 = cq.Workplane().box(1, 1, 1)
+    b2 = cq.Workplane(origin=(2, 0, 0)).box(1, 1, 1)
+
+    return cq.Assembly().add(b1).add(b2)
+
+
+def test_imprinting(touching_assy, disjoint_assy):
+
+    # normal usecase
+    r, o = cq.occ_impl.assembly.imprint(touching_assy)
+
+    assert len(r.Solids()) == 2
+    assert len(r.Faces()) == 11
+
+    for s in r.Solids():
+        assert s in o
+
+    # edge usecase
+    r, o = cq.occ_impl.assembly.imprint(disjoint_assy)
+
+    assert len(r.Solids()) == 2
+    assert len(r.Faces()) == 12
+
+    for s in r.Solids():
+        assert s in o