from . import (
Shape,
Workplane,
Assembly,
Sketch,
Compound,
Color,
Vector,
Location,
Face,
Edge,
)
from .occ_impl.assembly import _loc2vtk, toVTKAssy
from .occ_impl.nurbs import Curve, Surface
from typing import Union, List, Tuple, Iterable, cast, Optional
from OCP.TopoDS import TopoDS_Shape
from OCP.Geom import Geom_BSplineSurface
from vtkmodules.vtkInteractionWidgets import vtkOrientationMarkerWidget
from vtkmodules.vtkRenderingAnnotation import vtkAxesActor
from vtkmodules.vtkInteractionStyle import vtkInteractorStyleTrackballCamera
from vtkmodules.vtkRenderingCore import (
vtkMapper,
vtkRenderWindowInteractor,
vtkActor,
vtkProp3D,
vtkPolyDataMapper,
vtkAssembly,
vtkRenderWindow,
vtkWindowToImageFilter,
vtkRenderer,
)
from vtkmodules.vtkCommonCore import vtkPoints
from vtkmodules.vtkCommonDataModel import vtkCellArray, vtkPolyData
from vtkmodules.vtkCommonColor import vtkNamedColors
from vtkmodules.vtkIOImage import vtkPNGWriter
from .utils import instance_of
DEFAULT_COLOR = (1, 0.8, 0)
DEFAULT_EDGE_COLOR = (0, 0, 0)
DEFAULT_PT_SIZE = 7.5
DEFAULT_PT_COLOR = "darkviolet"
DEFAULT_CTRL_PT_COLOR = "crimson"
DEFAULT_CTRL_PT_SIZE = 7.5
SPECULAR = 0.3
SPECULAR_POWER = 100
SPECULAR_COLOR = vtkNamedColors().GetColor3d("White")
ShapeLike = Union[Shape, Workplane, Assembly, Sketch, TopoDS_Shape]
Showable = Union[
ShapeLike,
List[ShapeLike],
Vector,
List[Vector],
vtkProp3D,
List[vtkProp3D],
Location,
List[Location],
]
def _to_assy(
*objs: ShapeLike,
color: Tuple[float, float, float] = DEFAULT_COLOR,
alpha: float = 1,
) -> Assembly:
"""
Convert shapes to Assembly.
"""
assy = Assembly(color=Color(*color, alpha))
for obj in objs:
if isinstance(obj, (Shape, Workplane, Assembly)):
assy.add(obj)
elif isinstance(obj, Sketch):
assy.add(Compound.makeCompound(obj))
elif isinstance(obj, TopoDS_Shape):
assy.add(Shape(obj))
else:
raise ValueError(f"{obj} has unsupported type {type(obj)}")
return assy
def _split_showables(
objs,
) -> Tuple[List[ShapeLike], List[Vector], List[Location], List[vtkProp3D]]:
"""
Split into showables and others.
"""
rv_s: List[ShapeLike] = []
rv_v: List[Vector] = []
rv_l: List[Location] = []
rv_a: List[vtkProp3D] = []
for el in objs:
if instance_of(el, ShapeLike):
rv_s.append(el)
elif isinstance(el, Curve):
rv_s.append(el.edge())
elif isinstance(el, Surface):
rv_s.append(el.face())
elif isinstance(el, Vector):
rv_v.append(el)
elif isinstance(el, Location):
rv_l.append(el)
elif isinstance(el, vtkProp3D):
rv_a.append(el)
elif isinstance(el, list):
tmp1, tmp2, tmp3, tmp4 = _split_showables(el) # split recursively
rv_s.extend(tmp1)
rv_v.extend(tmp2)
rv_l.extend(tmp3)
rv_a.extend(tmp4)
return rv_s, rv_v, rv_l, rv_a
def _to_vtk_pts(
vecs: List[Vector], size: float = DEFAULT_PT_SIZE, color: str = DEFAULT_PT_COLOR
) -> vtkActor:
"""
Convert Vectors to vtkActor.
"""
rv = vtkActor()
mapper = vtkPolyDataMapper()
points = vtkPoints()
verts = vtkCellArray()
data = vtkPolyData()
data.SetPoints(points)
data.SetVerts(verts)
for v in vecs:
ix = points.InsertNextPoint(*v.toTuple())
verts.InsertNextCell(1)
verts.InsertCellPoint(ix)
mapper.SetInputData(data)
rv.SetMapper(mapper)
rv.GetProperty().SetColor(vtkNamedColors().GetColor3d(color))
rv.GetProperty().SetPointSize(size)
return rv
def _to_vtk_axs(locs: List[Location], scale: float = 0.1) -> List[vtkProp3D]:
"""
Convert Locations to vtkActor.
"""
rv = vtkAssembly()
for l in locs:
trans, rot = _loc2vtk(l)
ax = vtkAxesActor()
ax.SetAxisLabels(0)
ax.SetPosition(*trans)
ax.SetOrientation(*rot)
ax.SetScale(scale)
rv.AddPart(ax)
return [rv]
def _to_vtk_shapes(
obj: List[ShapeLike],
color: Tuple[float, float, float] = DEFAULT_COLOR,
edgecolor: Tuple[float, float, float] = DEFAULT_EDGE_COLOR,
edges: bool = True,
linewidth: float = 2,
alpha: float = 1,
tolerance: float = 1e-3,
) -> List[vtkProp3D]:
"""
Convert Shapes to vtkAssembly.
"""
return toVTKAssy(
_to_assy(*obj, color=color, alpha=alpha),
edgecolor=(*edgecolor, 1),
edges=edges,
linewidth=linewidth,
tolerance=tolerance,
)
[docs]
def ctrlPts(
s: Union[Face, Edge, Surface, Curve],
size: float = DEFAULT_CTRL_PT_SIZE,
color: str = DEFAULT_CTRL_PT_COLOR,
) -> vtkActor:
"""
Convert Edge, Face, Surface or Curve to a vtkActor representing control points.
"""
# handle Surface, Curve first
if isinstance(s, Surface):
return ctrlPts(s.face(), size, color)
elif isinstance(s, Curve):
return ctrlPts(s.edge(), size, color)
rv = vtkActor()
mapper = vtkPolyDataMapper()
points = vtkPoints()
cells = vtkCellArray()
data = vtkPolyData()
data.SetPoints(points)
data.SetVerts(cells)
data.SetLines(cells)
if isinstance(s, Face):
if isinstance(s._geomAdaptor(), Geom_BSplineSurface):
surf = cast(Geom_BSplineSurface, s._geomAdaptor())
else:
raise ValueError(
f"Only NURBS surfaces are supported, encountered {s._geomAdaptor()}"
)
Nu = surf.NbUPoles()
Nv = surf.NbVPoles()
u_periodic = surf.IsUPeriodic()
v_periodic = surf.IsVPeriodic()
# add points
for i in range(Nu):
for j in range(Nv):
pt = surf.Pole(i + 1, j + 1)
points.InsertNextPoint(pt.X(), pt.Y(), pt.Z())
# u edges
for j in range(Nv):
for i in range(Nu - 1):
cells.InsertNextCell(2, (Nv * i + j, Nv * (i + 1) + j))
if u_periodic:
cells.InsertNextCell(2, (Nv * (i + 1) + j, 0 + j))
# v edges
for i in range(Nu):
for j in range(Nv - 1):
cells.InsertNextCell(2, (Nv * i + j, Nv * i + j + 1))
if v_periodic:
cells.InsertNextCell(2, (Nv * i + j + 1, Nv * i + 0))
else:
if s.geomType() == "BSPLINE":
curve = s._geomAdaptor().BSpline()
else:
raise ValueError(
f"Only NURBS curves are supported, encountered {s.geomType()}"
)
for pt in curve.Poles():
points.InsertNextPoint(pt.X(), pt.Y(), pt.Z())
N = curve.NbPoles()
for i in range(N - 1):
cells.InsertNextCell(2, (i, i + 1))
if curve.IsPeriodic():
cells.InsertNextCell(2, (i + 1, 0))
mapper.SetInputData(data)
rv.SetMapper(mapper)
props = rv.GetProperty()
props.SetColor(vtkNamedColors().GetColor3d(color))
props.SetPointSize(size)
props.SetLineWidth(size / 3)
props.SetRenderPointsAsSpheres(True)
return rv
def _iterate_actors(
obj: Union[vtkProp3D, vtkActor, List[vtkProp3D]]
) -> Iterable[vtkActor]:
"""
Iterate over vtkActors, other props are ignored.
"""
if isinstance(obj, vtkActor):
yield obj
elif isinstance(obj, list):
for el in obj:
if isinstance(el, vtkActor):
yield el
[docs]
def style(
obj: Showable,
scale: float = 0.2,
alpha: float = 1,
tolerance: float = 1e-2,
edges: bool = True,
mesh: bool = False,
specular: bool = True,
markersize: float = 5,
linewidth: float = 2,
spheres: bool = False,
tubes: bool = False,
color: str = "gold",
edgecolor: str = "black",
meshcolor: str = "lightgrey",
vertexcolor: str = "cyan",
**kwargs,
) -> List[vtkProp3D]:
"""
Apply styling to CQ objects. To be used in conjunction with show.
"""
# styling functions
def _apply_style(actor):
props = actor.GetProperty()
props.SetEdgeColor(vtkNamedColors().GetColor3d(meshcolor))
props.SetVertexColor(vtkNamedColors().GetColor3d(vertexcolor))
props.SetPointSize(markersize)
props.SetLineWidth(linewidth)
props.SetRenderPointsAsSpheres(spheres)
props.SetRenderLinesAsTubes(tubes)
props.SetEdgeVisibility(mesh)
if specular:
props.SetSpecular(SPECULAR)
props.SetSpecularPower(SPECULAR_POWER)
props.SetSpecularColor(SPECULAR_COLOR)
def _apply_color(actor):
props = actor.GetProperty()
props.SetColor(vtkNamedColors().GetColor3d(color))
props.SetOpacity(alpha)
# split showables
shapes, vecs, locs, actors = _split_showables([obj,])
# convert to a prop
rv: Union[vtkActor, List[vtkProp3D]]
if shapes:
rv = _to_vtk_shapes(
shapes,
color=vtkNamedColors().GetColor3d(color),
edgecolor=vtkNamedColors().GetColor3d(edgecolor),
edges=edges,
linewidth=linewidth,
alpha=alpha,
tolerance=tolerance,
)
# apply style to every actor
for a in _iterate_actors(rv):
_apply_style(a)
elif vecs:
tmp = _to_vtk_pts(vecs)
_apply_style(tmp)
_apply_color(tmp)
rv = [tmp]
elif locs:
rv = _to_vtk_axs(locs, scale=scale)
else:
rv = []
for p in actors:
for a in _iterate_actors(p):
_apply_style(a)
_apply_color(a)
rv.append(a)
return rv
[docs]
def show(
*objs: Showable,
scale: float = 0.2,
alpha: float = 1,
tolerance: float = 1e-3,
edges: bool = False,
specular: bool = True,
title: str = "CQ viewer",
screenshot: Optional[str] = None,
interact: bool = True,
zoom: float = 1.0,
roll: float = -35,
elevation: float = -45,
azimuth: float = 0,
position: Optional[Tuple[float, float, float]] = None,
focus: Optional[Tuple[float, float, float]] = None,
viewup: Optional[Tuple[float, float, float]] = None,
clipping_range: Optional[Tuple[float, float]] = None,
width: Union[int, float] = 0.5,
height: Union[int, float] = 0.5,
trihedron: bool = True,
bgcolor: tuple[float, float, float] = (1, 1, 1),
gradient: bool = True,
xpos: Union[int, float] = 0,
ypos: Union[int, float] = 0,
fxaa: bool = False,
orthographic: bool = False,
):
"""
Show CQ objects using VTK. This functions optionally allows to make screenshots.
"""
# split objects
shapes, vecs, locs, props = _split_showables(objs)
# construct the assy
assy = _to_assy(*shapes, alpha=alpha)
# construct the points and locs
pts = _to_vtk_pts(vecs)
axs = _to_vtk_axs(locs, scale=scale)
# assy+renderer
renderer = vtkRenderer()
for act in toVTKAssy(assy, tolerance=tolerance):
renderer.AddActor(act)
# VTK window boilerplate
win = vtkRenderWindow()
# Render off-screen when not interacting
if not interact:
win.SetOffScreenRendering(1)
win.SetWindowName(title)
win.AddRenderer(renderer)
# get renderer and actor
for act in cast(Iterable[vtkActor], renderer.GetActors()):
propt = act.GetProperty()
if edges:
propt.EdgeVisibilityOn()
if specular:
propt.SetSpecular(SPECULAR)
propt.SetSpecularPower(SPECULAR_POWER)
propt.SetSpecularColor(SPECULAR_COLOR)
# rendering related settings
vtkMapper.SetResolveCoincidentTopologyToPolygonOffset()
vtkMapper.SetResolveCoincidentTopologyPolygonOffsetParameters(1, 0)
vtkMapper.SetResolveCoincidentTopologyLineOffsetParameters(-1, 0)
# create a VTK interactor
inter = vtkRenderWindowInteractor()
inter.SetInteractorStyle(vtkInteractorStyleTrackballCamera())
inter.SetRenderWindow(win)
# construct an axes indicator
axes = vtkAxesActor()
axes.SetDragable(0)
tp = axes.GetXAxisCaptionActor2D().GetCaptionTextProperty()
tp.SetColor(0, 0, 0)
axes.GetYAxisCaptionActor2D().GetCaptionTextProperty().ShallowCopy(tp)
axes.GetZAxisCaptionActor2D().GetCaptionTextProperty().ShallowCopy(tp)
# add to an orientation widget
if trihedron:
orient_widget = vtkOrientationMarkerWidget()
orient_widget.SetOrientationMarker(axes)
orient_widget.SetViewport(0.9, 0.0, 1.0, 0.2)
orient_widget.SetZoom(1.1)
orient_widget.SetInteractor(inter)
orient_widget.EnabledOn()
orient_widget.InteractiveOff()
# use gradient background
renderer.SetBackground(*bgcolor)
if gradient:
renderer.GradientBackgroundOn()
# use FXXAA
renderer.SetUseFXAA(fxaa)
# add pts and locs
renderer.AddActor(pts)
for ax in axs:
renderer.AddActor(ax)
# add other vtk actors
for p in props:
renderer.AddActor(p)
# set camera
camera = renderer.GetActiveCamera()
# set perspective or parallel (orthographic) projection
camera.SetParallelProjection(orthographic)
# Update camera position with user provided absolute positions
if viewup:
camera.SetViewUp(*viewup)
if focus:
camera.SetFocalPoint(*focus)
if position:
camera.SetPosition(*position)
if not (position or focus):
renderer.ResetCamera() # fit all if no explicit position provided
# Update camera position with user defined relative positions
camera.Roll(roll)
camera.Elevation(elevation)
camera.Azimuth(azimuth)
# Update camera view frustum
camera.Zoom(zoom)
if clipping_range:
camera.SetClippingRange(*clipping_range)
else:
renderer.ResetCameraClippingRange()
# initialize and set size
inter.Initialize()
w, h = win.GetScreenSize()
win.SetSize(
int(w * width) if isinstance(width, float) else width,
int(h * height) if isinstance(height, float) else height,
) # is height, width specified as float assume it is relative
# set position
win.SetPosition(
int(w * xpos) if isinstance(xpos, float) else xpos,
int(h * ypos) if isinstance(ypos, float) else ypos,
)
# show and return
win.Render()
# make a screenshot
if screenshot:
win2image = vtkWindowToImageFilter()
win2image.SetInput(win)
win2image.SetInputBufferTypeToRGB()
win2image.ReadFrontBufferOff()
win2image.Update()
writer = vtkPNGWriter()
writer.SetFileName(screenshot)
writer.SetInputConnection(win2image.GetOutputPort())
writer.Write()
# start interaction
if interact:
inter.Start()
# alias
show_object = show