Files
yet-another-cad-viewer/yacv_server/tessellate.py
2024-02-01 21:03:20 +01:00

152 lines
4.5 KiB
Python

import concurrent
import copyreg
from concurrent.futures import ProcessPoolExecutor, Executor
from dataclasses import dataclass
from enum import Enum
from typing import Optional, Tuple, Callable
import OCP
from OCP.BRepAdaptor import BRepAdaptor_Curve
from OCP.GCPnts import GCPnts_TangentialDeflection
from OCP.TopoDS import TopoDS_Face, TopoDS_Edge, TopoDS_Shape
from build123d import Face, Vector, Shape
from partcad.wrappers import cq_serialize
class UVMode(Enum):
"""UV mode for tesselation"""
TRIPLANAR = 0
"""Triplanar UV mapping"""
FACE = 1
"""Use UV coordinates from each face"""
@dataclass
class TessellationUpdate:
"""Tessellation update"""
# Progress
root: TopoDS_Shape
"""The root shape that is being tessellated"""
progress: float
"""Progress in percent"""
# Current shape
shape: TopoDS_Shape
"""Shape that was tessellated"""
vertices: list[Vector]
"""List of vertices"""
indices: Optional[list[Tuple[int, int, int]]]
"""List of indices (only for faces)"""
progress_callback_t = Callable[[TessellationUpdate], None]
def _inflate_vec(*values: float):
pnt = OCP.gp.gp_Vec(values[0], values[1], values[2])
return pnt
def _reduce_vec(pnt: OCP.gp.gp_Vec):
return _inflate_vec, (pnt.X(), pnt.Y(), pnt.Z())
def tessellate(
ocp_shape: TopoDS_Shape,
progress_callback: progress_callback_t = None,
tolerance: float = 0.1,
angular_tolerance: float = 0.1,
uv: Optional[UVMode] = None,
executor: Executor = ProcessPoolExecutor(), # Set to ThreadPoolExecutor if pickling fails...
):
"""Tessellate a whole shape into a list of triangle vertices and a list of triangle indices.
It uses multiprocessing to speed up the process, and publishes progress updates to the callback.
"""
shape = Shape(ocp_shape)
_register_pickle_if_needed()
with executor:
futures = []
# Submit tessellation tasks
for face in shape.faces():
futures.append(executor.submit(_tessellate_element, face.wrapped, tolerance, angular_tolerance, uv))
for edge in shape.edges():
futures.append(executor.submit(_tessellate_element, edge.wrapped, tolerance, angular_tolerance, uv))
# Collect results as they come in
for i, future in enumerate(concurrent.futures.as_completed(futures)):
tessellation, shape = future.result()
is_face = isinstance(shape, TopoDS_Face)
update = TessellationUpdate(
root=ocp_shape,
progress=(i + 1) / len(futures),
shape=shape,
vertices=tessellation[0] if is_face else tessellation,
indices=tessellation[1] if is_face else None,
)
progress_callback(update)
_pickle_registered = False
def _register_pickle_if_needed():
global _pickle_registered
if _pickle_registered:
return
cq_serialize.register()
copyreg.pickle(OCP.gp.gp_Vec, _reduce_vec)
# Define the function that will tessellate each element in parallel
def _tessellate_element(element: TopoDS_Shape, tolerance, angular_tolerance, uv):
if isinstance(element, TopoDS_Face):
return _tessellate_face(element, tolerance, angular_tolerance, uv), element
elif isinstance(element, TopoDS_Edge):
return _tessellate_edge(element, angular_tolerance, angular_tolerance), element
TriMesh = Tuple[list[Vector], list[Tuple[int, int, int]]]
def _tessellate_face(
ocp_face: TopoDS_Face,
tolerance: float = 0.1,
angular_tolerance: float = 0.1,
uv: Optional[UVMode] = None,
) -> TriMesh:
"""Tessellate a face into a list of triangle vertices and a list of triangle indices"""
face = Face(ocp_face)
tri_mesh = face.tessellate(tolerance, angular_tolerance)
# TODO: UV mapping
return tri_mesh
def _tessellate_edge(
ocp_edge: TopoDS_Edge,
angular_deflection: float = 0.1,
curvature_deflection: float = 0.1,
) -> list[Vector]:
"""Tessellate a wire or edge into a list of ordered vertices"""
curve = BRepAdaptor_Curve(ocp_edge)
discretizer = GCPnts_TangentialDeflection(curve, angular_deflection, curvature_deflection)
assert discretizer.NbPoints() > 1, "Edge is too small??"
# TODO: get transformation??
# add vertices
vertices: list[Vector] = [
Vector(v.X(), v.Y(), v.Z())
for v in (
discretizer.Value(i) # .Transformed(transformation)
for i in range(1, discretizer.NbPoints() + 1)
)
]
return vertices