Exports a CadQuery Model or Assembly to a PNG Image

png_export_from_cadquery_using_vtk.py

from math import degrees
from vtkmodules.vtkRenderingCore import vtkGraphicsFactory
from vtkmodules.vtkCommonColor import vtkNamedColors
from vtkmodules.vtkRenderingCore import (
    vtkActor,
    vtkPolyDataMapper as vtkMapper,
    vtkRenderer
)
from vtkmodules.vtkRenderingCore import vtkRenderWindow, vtkRenderWindowInteractor, vtkWindowToImageFilter
from vtkmodules.vtkFiltersExtraction import vtkExtractCellsByType
from vtkmodules.vtkCommonDataModel import VTK_TRIANGLE, VTK_LINE, VTK_VERTEX
from vtkmodules.vtkIOImage import vtkPNGWriter
import cadquery as cq


def process_cq_object(cq_obj):
    cq_objs = []
    face_actors = []
    edge_actors = []

    # Extract the parts out of an assembly
    if type(cq_obj).__name__ == "Assembly":
        for subassy in cq_obj.traverse():
            for shape, name, loc, col in subassy[1]:
                color = col.toTuple() if col else (0.5, 0.5, 0.5, 1.0)
                trans, rot = loc.toTuple()

                # Lower level shapes need to be named and wrapped in a cq.Workplane object
                model = cq.Workplane(shape)
                cq_objs.append((model, color, trans, rot, name))
    else:
        cq_objs.append((cq_obj, (0.5, 0.1, 0.8, 1.0), (0, 0, 0), (0, 0, 0)))

    for obj in cq_objs:
        # Tesselate the CQ object into VTK data
        vtk_data = obj[0].val().toVtkPolyData(1e-3, 0.1)
        color = obj[1]
        translation = obj[2]
        rotation = obj[3]

        # Extract faces
        extr = vtkExtractCellsByType()
        extr.SetInputDataObject(vtk_data)

        extr.AddCellType(VTK_LINE)
        extr.AddCellType(VTK_VERTEX)
        extr.Update()
        data_edges = extr.GetOutput()

        # Extract edges
        extr = vtkExtractCellsByType()
        extr.SetInputDataObject(vtk_data)

        extr.AddCellType(VTK_TRIANGLE)
        extr.Update()
        data_faces = extr.GetOutput()

        # Remove normals from edges
        data_edges.GetPointData().RemoveArray("Normals")

        # Set up the face and edge mappers and actors
        face_mapper = vtkMapper()
        face_actor = vtkActor()
        face_actor.SetMapper(face_mapper)
        edge_mapper = vtkMapper()
        edge_actor = vtkActor()
        edge_actor.SetMapper(edge_mapper)

        # Update the faces
        face_mapper.SetInputDataObject(data_faces)
        face_actor.SetPosition(*translation)
        face_actor.SetOrientation(*map(degrees, rotation))
        face_actor.GetProperty().SetColor(*color[:3])
        face_actor.GetProperty().SetOpacity(color[3])

        # Update the edges
        edge_mapper.SetInputDataObject(data_edges)
        edge_actor.SetPosition(*translation)
        edge_actor.SetOrientation(*map(degrees, rotation))
        edge_actor.GetProperty().SetColor(1.0, 1.0, 1.0)
        edge_actor.GetProperty().SetLineWidth(1)

        # Handle all actors
        face_actors.append(face_actor)
        edge_actors.append(edge_actor)

    return (face_actors, edge_actors)


def export(cq_object, path, opt=None):
    """
    Stand-in for the exporters.export method that is part of CadQuery.
    """

    # Handle view options that were passed in
    if opt != None:
        width = opt["width"] if "width" in opt else 800
        height = opt["height"] if "height" in opt else 600
        camera_position = opt["camera_position"] if "camera_position" in opt else (0, 0, 0)
        view_up_direction = opt["view_up_direction"] if "view_up_direction" in opt else (0, 0, 1)
        focal_point = opt["focal_point"] if "focal_point" in opt else (0, 0, 0)
        parallel_projection = opt["parallel_projection"] if "parallel_projection" in opt else False
    else:
        width = 800
        height = 600
        camera_position = (0, 0, 0)
        view_up_direction = (0, 0, 1)
        focal_point = (0, 0, 0)
        parallel_projection = False

    colors = vtkNamedColors()

    # Setup offscreen rendering
    graphics_factory = vtkGraphicsFactory()
    graphics_factory.SetOffScreenOnlyMode(1)
    graphics_factory.SetUseMesaClasses(1)

    # Process the CadQuery object into faces and edges
    face_actors, edge_actors = process_cq_object(cq_object)

    # A renderer and render window
    renderer = vtkRenderer()
    renderWindow = vtkRenderWindow()
    renderWindow.SetSize(width, height)
    # renderWindow.SetFullScreen(True)
    renderWindow.SetOffScreenRendering(1)

    renderWindow.AddRenderer(renderer)

    # Add the actors to the scene
    for face_actor in face_actors:
        renderer.AddActor(face_actor)
    for edge_actor in edge_actors:
        renderer.AddActor(edge_actor)

    renderer.SetBackground(0.5, 0.5, 0.5)

    # Render the scene
    renderWindow.Render()

    # Set the camera as the user requested
    camera = renderer.GetActiveCamera()
    camera.SetPosition(camera_position[0], camera_position[1], camera_position[2])
    camera.SetViewUp(view_up_direction[0], view_up_direction[1], view_up_direction[2])
    camera.SetFocalPoint(focal_point[0], focal_point[1], focal_point[2])
    if parallel_projection:
        camera.ParallelProjectionOn()
    else:
        camera.ParallelProjectionOff()

    # Export a PNG of the scene
    windowToImageFilter = vtkWindowToImageFilter()
    windowToImageFilter.SetInput(renderWindow)
    windowToImageFilter.Update()

    writer = vtkPNGWriter()
    writer.SetFileName(path)
    writer.SetInputConnection(windowToImageFilter.GetOutputPort())
    writer.Write()

png_path = "sample.png"
svg_path = "sample.svg"

# Set up a CadQuery object to display
# cq_object = cq.Workplane().box(10, 10, 10).edges().fillet(2.5)
cq_object = cq.Assembly()
cq_object.add(cq.Workplane().box(10, 10, 10), color=cq.Color(1.0, 0, 0), name="box1")
cq_object.add(cq.Workplane().box(5, 5, 5), color=cq.Color(0, 1.0, 0), loc=cq.Location((10.0, 10.0, 10.0), (0, 1, 0), 45), name="box2")

# Do the png export
png_opts = {
    "width": 600,
    "height": 600,
    "camera_position": (37, 37, 37),
    "view_up_direction": (0, 1, 0),
    "parallel_projection": True,
}
export(cq_object, png_path, png_opts)