139 lines
4.3 KiB
Python
139 lines
4.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# @author Arnaud Morin <arnaud.gfpv@mailops.fr>
|
|
|
|
"""Headless 3D render helper -- shared by every part module and the assembly.
|
|
|
|
from render import render
|
|
render(shape_or_assembly, "build/foo.png", view="iso")
|
|
|
|
Views: iso, top, front, right, or an (elev, azim) tuple.
|
|
"""
|
|
|
|
import matplotlib
|
|
matplotlib.use("Agg")
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
|
|
|
|
import cadquery as cq
|
|
|
|
VIEWS = {
|
|
"iso": (26, -55),
|
|
"iso2": (20, 35),
|
|
"top": (89, -90),
|
|
"front": (0, -90),
|
|
"right": (0, 0),
|
|
}
|
|
|
|
DEFAULT_COLOR = (0.20, 0.20, 0.22)
|
|
|
|
|
|
def _iter_parts(assy, parent_loc=None):
|
|
"""Yield (located_shape, rgba|None) for every solid in an Assembly."""
|
|
loc = assy.loc if parent_loc is None else parent_loc * assy.loc
|
|
if assy.obj is not None:
|
|
shape = assy.obj.val() if isinstance(assy.obj, cq.Workplane) else assy.obj
|
|
col = assy.color.toTuple() if assy.color is not None else None
|
|
yield shape.moved(loc), col
|
|
for child in assy.children:
|
|
yield from _iter_parts(child, loc)
|
|
|
|
|
|
def _shapes(obj):
|
|
if isinstance(obj, cq.Assembly):
|
|
return list(_iter_parts(obj))
|
|
if isinstance(obj, cq.Workplane):
|
|
return [(obj.val(), None)]
|
|
return [(obj, None)]
|
|
|
|
|
|
def render(obj, path, view="iso", tol=0.08, figsize=(9, 9), dpi=110,
|
|
edges=True, title=None):
|
|
parts = _shapes(obj)
|
|
|
|
fig = plt.figure(figsize=figsize)
|
|
ax = fig.add_subplot(111, projection="3d")
|
|
|
|
lo = np.array([+1e9, +1e9, +1e9])
|
|
hi = np.array([-1e9, -1e9, -1e9])
|
|
|
|
for shape, col in parts:
|
|
verts, tris = shape.tessellate(tol)
|
|
if not tris:
|
|
continue
|
|
v = np.array([[p.x, p.y, p.z] for p in verts])
|
|
lo = np.minimum(lo, v.min(axis=0))
|
|
hi = np.maximum(hi, v.max(axis=0))
|
|
face = col[:3] if col else DEFAULT_COLOR
|
|
pc = Poly3DCollection(
|
|
v[np.array(tris)],
|
|
facecolor=face,
|
|
edgecolor=(0, 0, 0, 0.25) if edges else "none",
|
|
linewidth=0.15 if edges else 0,
|
|
)
|
|
pc.set_alpha(col[3] if col else 1.0)
|
|
ax.add_collection3d(pc)
|
|
|
|
ctr = (lo + hi) / 2.0
|
|
span = float((hi - lo).max()) * 0.55 or 1.0
|
|
ax.set_xlim(ctr[0] - span, ctr[0] + span)
|
|
ax.set_ylim(ctr[1] - span, ctr[1] + span)
|
|
ax.set_zlim(ctr[2] - span, ctr[2] + span)
|
|
try:
|
|
ax.set_box_aspect((1, 1, 1))
|
|
except Exception:
|
|
pass
|
|
|
|
elev, azim = VIEWS.get(view, view) if isinstance(view, str) else view
|
|
ax.view_init(elev=elev, azim=azim)
|
|
ax.set_axis_off()
|
|
if title:
|
|
ax.set_title(title)
|
|
fig.savefig(path, dpi=dpi, bbox_inches="tight", facecolor="white")
|
|
plt.close(fig)
|
|
print(f"wrote {path}")
|
|
return path
|
|
|
|
|
|
def render_grid(obj, path, views=("iso", "top", "front", "right"), tol=0.08, dpi=100):
|
|
"""2x2 contact sheet -- the quickest way to eyeball a part or assembly."""
|
|
parts = _shapes(obj)
|
|
fig = plt.figure(figsize=(12, 12))
|
|
for i, view in enumerate(views, start=1):
|
|
ax = fig.add_subplot(2, 2, i, projection="3d")
|
|
lo = np.array([+1e9] * 3)
|
|
hi = np.array([-1e9] * 3)
|
|
for shape, col in parts:
|
|
verts, tris = shape.tessellate(tol)
|
|
if not tris:
|
|
continue
|
|
v = np.array([[p.x, p.y, p.z] for p in verts])
|
|
lo = np.minimum(lo, v.min(axis=0))
|
|
hi = np.maximum(hi, v.max(axis=0))
|
|
pc = Poly3DCollection(
|
|
v[np.array(tris)],
|
|
facecolor=col[:3] if col else DEFAULT_COLOR,
|
|
edgecolor=(0, 0, 0, 0.25),
|
|
linewidth=0.15,
|
|
)
|
|
pc.set_alpha(col[3] if col else 1.0)
|
|
ax.add_collection3d(pc)
|
|
ctr = (lo + hi) / 2.0
|
|
span = float((hi - lo).max()) * 0.55 or 1.0
|
|
ax.set_xlim(ctr[0] - span, ctr[0] + span)
|
|
ax.set_ylim(ctr[1] - span, ctr[1] + span)
|
|
ax.set_zlim(ctr[2] - span, ctr[2] + span)
|
|
try:
|
|
ax.set_box_aspect((1, 1, 1))
|
|
except Exception:
|
|
pass
|
|
elev, azim = VIEWS.get(view, view) if isinstance(view, str) else view
|
|
ax.view_init(elev=elev, azim=azim)
|
|
ax.set_axis_off()
|
|
ax.set_title(str(view))
|
|
fig.savefig(path, dpi=dpi, bbox_inches="tight", facecolor="white")
|
|
plt.close(fig)
|
|
print(f"wrote {path}")
|
|
return path
|