Implement exporting to Gltf (#802)

* Implement GLTF export

* Extend the save method

* Implement tests

* Update tests/test_assembly.py

Co-authored-by: Marcus Boyd <[email protected]>

* Fix gltf export

* Improve tests

* refactor tests with pytest parameterize

Co-authored-by: Marcus Boyd <[email protected]>

Author: adam-urbanczyk

cadquery/assembly.py

@@ -12,7 +12,13 @@
     ConstraintMarker,
     Constraint as ConstraintPOD,
 )
-from .occ_impl.exporters.assembly import exportAssembly, exportCAF
+from .occ_impl.exporters.assembly import (
+    exportAssembly,
+    exportCAF,
+    exportVTKJS,
+    exportVRML,
+    exportGLTF,
+)
 
 from .selectors import _expression_grammar as _selector_grammar
 from OCP.BRepTools import BRepTools
@@ -22,7 +28,7 @@
 # type definitions
 AssemblyObjects = Union[Shape, Workplane, None]
 ConstraintKinds = Literal["Plane", "Point", "Axis", "PointInPlane"]
-ExportLiterals = Literal["STEP", "XML"]
+ExportLiterals = Literal["STEP", "XML", "GLTF", "VTKJS", "VRML"]
 
 PATH_DELIM = "/"
 
@@ -462,18 +468,24 @@ def solve(self) -> "Assembly":
         return self
 
     def save(
-        self, path: str, exportType: Optional[ExportLiterals] = None
+        self,
+        path: str,
+        exportType: Optional[ExportLiterals] = None,
+        tolerance: float = 0.1,
+        angularTolerance: float = 0.1,
     ) -> "Assembly":
         """
         save as STEP or OCCT native XML file
 
         :param path: filepath
         :param exportType: export format (default: None, results in format being inferred form the path)
+        :param tolerance: the deflection tolerance, in model units. Only used for GLTF. Default 0.1.
+        :param angularTolerance: the angular tolerance, in radians. Only used for GLTF. Default 0.1.
         """
 
         if exportType is None:
             t = path.split(".")[-1].upper()
-            if t in ("STEP", "XML"):
+            if t in ("STEP", "XML", "VRML", "VTKJS", "GLTF"):
                 exportType = cast(ExportLiterals, t)
             else:
                 raise ValueError("Unknown extension, specify export type explicitly")
@@ -482,6 +494,12 @@ def save(
             exportAssembly(self, path)
         elif exportType == "XML":
             exportCAF(self, path)
+        elif exportType == "VRML":
+            exportVRML(self, path)
+        elif exportType == "GLTF":
+            exportGLTF(self, path, True, tolerance, angularTolerance)
+        elif exportType == "VTKJS":
+            exportVTKJS(self, path)
         else:
             raise ValueError(f"Unknown format: {exportType}")
 

cadquery/occ_impl/exporters/assembly.py

@@ -2,6 +2,7 @@
 
 from tempfile import TemporaryDirectory
 from shutil import make_archive
+from itertools import chain
 
 from vtkmodules.vtkIOExport import vtkJSONSceneExporter, vtkVRMLExporter
 from vtkmodules.vtkRenderingCore import vtkRenderer, vtkRenderWindow
@@ -17,6 +18,9 @@
 )
 from OCP.TCollection import TCollection_ExtendedString, TCollection_AsciiString
 from OCP.PCDM import PCDM_StoreStatus
+from OCP.RWGltf import RWGltf_CafWriter
+from OCP.TColStd import TColStd_IndexedDataMapOfStringString
+from OCP.Message import Message_ProgressRange
 
 from ..assembly import AssemblyProtocol, toCAF, toVTK
 
@@ -116,3 +120,27 @@ def exportVRML(assy: AssemblyProtocol, path: str):
     exporter.SetFileName(path)
     exporter.SetRenderWindow(_vtkRenderWindow(assy))
     exporter.Write()
+
+
+def exportGLTF(
+    assy: AssemblyProtocol,
+    path: str,
+    binary: bool = True,
+    tolerance: float = 0.1,
+    angularTolerance: float = 0.1,
+):
+    """
+    Export an assembly to a gltf file.
+    """
+
+    # mesh all the shapes
+    for _, el in assy.traverse():
+        for s in el.shapes:
+            s.mesh(tolerance, angularTolerance)
+
+    _, doc = toCAF(assy, True)
+
+    writer = RWGltf_CafWriter(TCollection_AsciiString(path), binary)
+    return writer.Perform(
+        doc, TColStd_IndexedDataMapOfStringString(), Message_ProgressRange()
+    )

tests/test_assembly.py

@@ -50,6 +50,22 @@ def nested_assy():
     return assy
 
 
+@pytest.fixture
+def nested_assy_sphere():
+
+    b1 = cq.Workplane().box(1, 1, 1).faces("<Z").tag("top_face").end()
+    b2 = cq.Workplane().box(1, 1, 1).faces("<Z").tag("bottom_face").end()
+    b3 = cq.Workplane().pushPoints([(-2, 0), (2, 0)]).tag("pts").sphere(1).tag("boxes")
+
+    assy = cq.Assembly(b1, loc=cq.Location(cq.Vector(0, 0, 0)), name="TOP")
+    assy2 = cq.Assembly(b2, loc=cq.Location(cq.Vector(0, 4, 0)), name="SECOND")
+    assy2.add(b3, loc=cq.Location(cq.Vector(0, 4, 0)), name="BOTTOM")
+
+    assy.add(assy2, color=cq.Color("green"))
+
+    return assy
+
+
 @pytest.fixture
 def empty_top_assy():
 
@@ -168,28 +184,43 @@ def test_toJSON(simple_assy, nested_assy, empty_top_assy):
     assert len(r3) == 1
 
 
-def test_save(simple_assy, nested_assy):
+@pytest.mark.parametrize(
+    "extension, args",
+    [
+        ("step", ()),
+        ("xml", ()),
+        ("stp", ("STEP",)),
+        ("caf", ("XML",)),
+        ("wrl", ("VRML",)),
+    ],
+)
+def test_save(extension, args, nested_assy, nested_assy_sphere):
+
+    filename = "nested." + extension
+    nested_assy.save(filename, *args)
+    assert os.path.exists(filename)
+
+
+def test_save_gltf(nested_assy_sphere):
+
+    nested_assy_sphere.save("nested.glb", "GLTF")
+    assert os.path.exists("nested.glb")
+    assert os.path.getsize("nested.glb") > 50 * 1024
 
-    simple_assy.save("simple.step")
-    assert os.path.exists("simple.step")
 
-    simple_assy.save("simple.xml")
-    assert os.path.exists("simple.xml")
+def test_save_vtkjs(nested_assy):
 
-    simple_assy.save("simple.step")
-    assert os.path.exists("simple.step")
+    nested_assy.save("nested", "VTKJS")
+    assert os.path.exists("nested.zip")
 
-    simple_assy.save("simple.stp", "STEP")
-    assert os.path.exists("simple.stp")
 
-    simple_assy.save("simple.caf", "XML")
-    assert os.path.exists("simple.caf")
+def test_save_raises(nested_assy):
 
     with pytest.raises(ValueError):
-        simple_assy.save("simple.dxf")
+        nested_assy.save("nested.dxf")
 
     with pytest.raises(ValueError):
-        simple_assy.save("simple.step", "DXF")
+        nested_assy.save("nested.step", "DXF")
 
 
 def test_constrain(simple_assy, nested_assy):