Skip to content

Instantly share code, notes, and snippets.

@jmwright
Last active April 22, 2024 15:37
Show Gist options
  • Save jmwright/e9ce05b197ea64d54bab7ca4ee4469ee to your computer and use it in GitHub Desktop.
Save jmwright/e9ce05b197ea64d54bab7ca4ee4469ee to your computer and use it in GitHub Desktop.
Exports a CadQuery Model or Assembly to a PNG Image
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment