Material and texture support, simpler tessellation and valid models

yacv_server/gltf.py

@@ -1,37 +1,33 @@
 import numpy as np
 from pygltflib import *
 
-from tessellate import TessellationUpdate
 
-
-def create_gltf_from_update(update: TessellationUpdate) -> GLTF2:
-    """Create a glTF object from a tessellation update."""
-    return create_gltf(
-        np.array(list(map(lambda v: [v.X, v.Y, v.Z], update.vertices))),
-        np.array(update.indices) if update.indices else None
-    )
-
-
-def create_gltf(vertices: np.ndarray, indices_in: Optional[np.ndarray]) -> GLTF2:
+def create_gltf(vertices: np.ndarray, indices: np.ndarray, tex_coord: np.ndarray, mode: int = TRIANGLES,
+                material: Optional[Material] = None, add_checkerboard_image: bool = False) -> GLTF2:
     """Create a glTF object from vertices and optionally indices.
 
     If indices are not set, vertices are interpreted as line_strip."""
 
     assert vertices.ndim == 2
     assert vertices.shape[1] == 3
-    vertices = vertices  # .astype(np.float16)
+    vertices = vertices.astype(np.float32)
     vertices_blob = vertices.tobytes()
     # print(vertices)
 
-    if indices_in is not None:
-        assert indices_in.ndim == 2
-        assert indices_in.shape[1] == 3
-        indices = indices_in  # .astype(np.uint8)
-    else:
-        indices = np.array(list(map(lambda i: [i, i + 1], range(len(vertices) - 1))), dtype=np.uint8)
+    indices = indices.astype(np.uint8)
     indices_blob = indices.flatten().tobytes()
     # print(indices)
 
+    tex_coord = tex_coord.astype(np.float32)
+    tex_coord_blob = tex_coord.tobytes()
+    # print(tex_coord)
+
+    image_blob = b''
+    if add_checkerboard_image:
+        image_blob = base64.decodebytes(
+            b'iVBORw0KGgoAAAANSUhEUgAAABAAAAAQAQMAAAAlPW0iAAAABlBMVEX'
+            b'////MzMw46qqDAAAAEElEQVQImWNg+M+AFeEQBgB+vw/xfUUZkgAAAABJRU5ErkJggg==')
+
     gltf = GLTF2(
         scene=0,
         scenes=[Scene(nodes=[0])],
@@ -40,51 +36,81 @@ def create_gltf(vertices: np.ndarray, indices_in: Optional[np.ndarray]) -> GLTF2
             Mesh(
                 primitives=[
                     Primitive(
-                        attributes=Attributes(POSITION=1),
+                        attributes=Attributes(POSITION=1, TEXCOORD_0=2) if len(tex_coord) > 0 else Attributes(
+                            POSITION=1),
                         indices=0,
-                        mode=TRIANGLES if indices_in is not None else LINE_STRIP
-                        # TODO: Also support POINTS mode
+                        mode=mode,
+                        material=0 if material is not None else None,
                     )
                 ]
             )
         ],
+        materials=[material] if material is not None else [],
         accessors=[
-            Accessor(
-                bufferView=0,
-                componentType=UNSIGNED_BYTE,
-                count=indices.size,
-                type=SCALAR,
-                max=[int(indices.max())],
-                min=[int(indices.min())],
-            ),
-            Accessor(
-                bufferView=1,
-                componentType=FLOAT,
-                count=len(vertices),
-                type=VEC3,
-                max=vertices.max(axis=0).tolist(),
-                min=vertices.min(axis=0).tolist(),
-            ),
-        ],
+                      Accessor(
+                          bufferView=0,
+                          componentType=UNSIGNED_BYTE,
+                          count=indices.size,
+                          type=SCALAR,
+                          max=[int(indices.max())],
+                          min=[int(indices.min())],
+                      ),
+                      Accessor(
+                          bufferView=1,
+                          componentType=FLOAT,
+                          count=len(vertices),
+                          type=VEC3,
+                          max=vertices.max(axis=0).tolist(),
+                          min=vertices.min(axis=0).tolist(),
+                      ),
+                  ] + ([
+                           Accessor(
+                               bufferView=2,
+                               componentType=FLOAT,
+                               count=len(tex_coord),
+                               type=VEC2,
+                               max=tex_coord.max(axis=0).tolist(),
+                               min=tex_coord.min(axis=0).tolist(),
+                           )] if len(tex_coord) > 0 else [])
+        ,
         bufferViews=[
-            BufferView(
-                buffer=0,
-                byteLength=len(indices_blob),
-                target=ELEMENT_ARRAY_BUFFER,
-            ),
-            BufferView(
-                buffer=0,
-                byteOffset=len(indices_blob),
-                byteLength=len(vertices_blob),
-                target=ARRAY_BUFFER,
-            ),
-        ],
+                        BufferView(
+                            buffer=0,
+                            byteLength=len(indices_blob),
+                            target=ELEMENT_ARRAY_BUFFER,
+                        ),
+                        BufferView(
+                            buffer=0,
+                            byteOffset=len(indices_blob),
+                            byteLength=len(vertices_blob),
+                            target=ARRAY_BUFFER,
+                        ),
+                    ] + (
+                        [
+                            BufferView(
+                                buffer=0,
+                                byteOffset=len(indices_blob) + len(vertices_blob),
+                                byteLength=len(tex_coord_blob),
+                                target=ARRAY_BUFFER,
+                            ),
+                        ] if len(tex_coord) > 0 else []) + (
+                        [
+                            BufferView(
+                                buffer=0,
+                                byteOffset=len(indices_blob) + len(
+                                    vertices_blob) + len(tex_coord_blob),
+                                byteLength=len(image_blob),
+                            ),
+                        ] if add_checkerboard_image else []),
         buffers=[
             Buffer(
-                byteLength=len(indices_blob) + len(vertices_blob)
+                byteLength=len(indices_blob) + len(vertices_blob) + len(tex_coord_blob) + len(image_blob),
             )
         ],
+        textures=[Texture(source=0)] if add_checkerboard_image else [],
+        images=[Image(bufferView=3, mimeType=IMAGEPNG, )] if add_checkerboard_image else [],
     )
 
-    gltf.set_binary_blob(indices_blob + vertices_blob)
+    gltf.set_binary_blob(indices_blob + vertices_blob + tex_coord_blob + image_blob)
+
     return gltf

yacv_server/logo/logo.py

@@ -1,7 +1,7 @@
 from build123d import *
+from tqdm import tqdm
 
-from gltf import create_gltf_from_update
-from tessellate import tessellate, TessellationUpdate
+from tessellate import tessellate, tessellate_count
 
 
 def logo() -> Compound:
@@ -14,14 +14,6 @@ def logo() -> Compound:
 if __name__ == "__main__":
     obj = logo()
 
-
-    def progress(update: TessellationUpdate):
-        gltf = create_gltf_from_update(update)
-        print(gltf)
-        if update.is_face:
-            gltf.save("logo_face.glb")
-        else:
-            gltf.save("logo_edge.glb")
-
-
-    tessellate(obj.wrapped, progress)
+    for update in tqdm(tessellate(obj.wrapped), total=tessellate_count(obj.wrapped)):
+        # print(update.gltf)
+        update.gltf.save(f'logo_{update.kind}.glb')  # Will overwrite the file for each update

yacv_server/requirements.txt

@@ -1,3 +0,0 @@
-build123d==0.3.0
-aiohttp==3.9.3
-partcad==0.3.84

yacv_server/tessellate.py

@@ -2,38 +2,45 @@ import concurrent
 import copyreg
 from concurrent.futures import ProcessPoolExecutor, Executor
 from dataclasses import dataclass
-from enum import Enum
-from typing import Optional, Tuple, Callable
+from typing import Tuple, Callable, Generator
 
 import OCP
+import numpy as np
+from OCP.BRep import BRep_Tool
 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 OCP.TopLoc import TopLoc_Location
+from OCP.TopoDS import TopoDS_Face, TopoDS_Edge, TopoDS_Shape, TopoDS_Vertex
+from build123d import Face, Vector, Shape, Vertex
 from partcad.wrappers import cq_serialize
+from pygltflib import LINE_STRIP, GLTF2, Material, PbrMetallicRoughness, TRIANGLES, POINTS, TextureInfo
+
+from gltf import create_gltf
 
 
 @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)"""
+    """(Sub)shape that was tessellated"""
+    gltf: GLTF2
+    """The valid GLTF containing only the current shape"""
 
     @property
-    def is_face(self):
-        return isinstance(self.shape, TopoDS_Face)
+    def kind(self) -> str:
+        """The kind of the shape"""
+        if isinstance(self.shape, TopoDS_Face):
+            return "face"
+        elif isinstance(self.shape, TopoDS_Edge):
+            return "edge"
+        elif isinstance(self.shape, TopoDS_Vertex):
+            return "vertex"
+        else:
+            raise ValueError(f"Unknown shape type: {self.shape}")
 
 
 progress_callback_t = Callable[[TessellationUpdate], None]
@@ -48,13 +55,18 @@ def _reduce_vec(pnt: OCP.gp.gp_Vec):
     return _inflate_vec, (pnt.X(), pnt.Y(), pnt.Z())
 
 
+def tessellate_count(ocp_shape: TopoDS_Shape) -> int:
+    """Count the number of elements that will be tessellated"""
+    shape = Shape(ocp_shape)
+    return len(shape.faces()) + len(shape.edges()) + len(shape.vertices())
+
+
 def tessellate(
         ocp_shape: TopoDS_Shape,
-        progress_callback: progress_callback_t = None,
         tolerance: float = 0.1,
         angular_tolerance: float = 0.1,
         executor: Executor = ProcessPoolExecutor(),  # Set to ThreadPoolExecutor if pickling fails...
-):
+) -> Generator[TessellationUpdate, None, None]:
     """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.
@@ -69,19 +81,17 @@ def tessellate(
             futures.append(executor.submit(_tessellate_element, face.wrapped, tolerance, angular_tolerance))
         for edge in shape.edges():
             futures.append(executor.submit(_tessellate_element, edge.wrapped, tolerance, angular_tolerance))
+        for vertex in shape.vertices():
+            futures.append(executor.submit(_tessellate_element, vertex.wrapped, tolerance, angular_tolerance))
 
         # 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,
+            sub_shape, gltf = future.result()
+            yield TessellationUpdate(
                 progress=(i + 1) / len(futures),
-                shape=shape,
-                vertices=tessellation[0] if is_face else tessellation,
-                indices=tessellation[1] if is_face else None,
+                shape=sub_shape,
+                gltf=gltf,
             )
-            progress_callback(update)
 
 
 _pickle_registered = False
@@ -96,11 +106,16 @@ def _register_pickle_if_needed():
 
 
 # Define the function that will tessellate each element in parallel
-def _tessellate_element(element: TopoDS_Shape, tolerance: float, angular_tolerance: float):
+def _tessellate_element(
+        element: TopoDS_Shape, tolerance: float, angular_tolerance: float) -> Tuple[TopoDS_Shape, GLTF2]:
     if isinstance(element, TopoDS_Face):
-        return _tessellate_face(element, tolerance, angular_tolerance), element
+        return element, _tessellate_face(element, tolerance, angular_tolerance)
     elif isinstance(element, TopoDS_Edge):
-        return _tessellate_edge(element, angular_tolerance, angular_tolerance), element
+        return element, _tessellate_edge(element, angular_tolerance, angular_tolerance)
+    elif isinstance(element, TopoDS_Vertex):
+        return element, _tessellate_vertex(element)
+    else:
+        raise ValueError(f"Unknown element type: {element}")
 
 
 TriMesh = Tuple[list[Vector], list[Tuple[int, int, int]]]
@@ -110,35 +125,66 @@ def _tessellate_face(
         ocp_face: TopoDS_Face,
         tolerance: float = 0.1,
         angular_tolerance: float = 0.1
-) -> TriMesh:
+) -> GLTF2:
     """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 of each face
+    # Get UV of each face from the parameters
+    loc = TopLoc_Location()
+    poly = BRep_Tool.Triangulation_s(face.wrapped, loc)
+    uv = [
+        [v.X(), v.Y()]
+        for v in (poly.UVNode(i) for i in range(1, poly.NbNodes() + 1))
+    ]
 
-    return tri_mesh
+    vertices = np.array(list(map(lambda v: [v.X, v.Y, v.Z], tri_mesh[0])))
+    indices = np.array(tri_mesh[1])
+    tex_coord = np.array(uv)
+    mode = TRIANGLES
+    material = Material(pbrMetallicRoughness=PbrMetallicRoughness(
+        baseColorFactor=[0.3, 1.0, 0.2, 1.0], roughnessFactor=0.1, baseColorTexture=TextureInfo(index=0)),
+        alphaCutoff=None)
+    return create_gltf(vertices, indices, tex_coord, mode, material, add_checkerboard_image=True)
 
 
 def _tessellate_edge(
         ocp_edge: TopoDS_Edge,
         angular_deflection: float = 0.1,
         curvature_deflection: float = 0.1,
-) -> list[Vector]:
+) -> GLTF2:
     """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??
+    # TODO: get and apply transformation??
 
     # add vertices
-    vertices: list[Vector] = [
-        Vector(v.X(), v.Y(), v.Z())
+    vertices: list[list[float]] = [
+        [v.X(), v.Y(), v.Z()]
         for v in (
             discretizer.Value(i)  # .Transformed(transformation)
             for i in range(1, discretizer.NbPoints() + 1)
         )
     ]
+    indices = np.array(list(map(lambda i: [i, i + 1], range(len(vertices) - 1))), dtype=np.uint8)
+    tex_coord = np.array([], dtype=np.float32)
+    mode = LINE_STRIP
+    material = Material(
+        pbrMetallicRoughness=PbrMetallicRoughness(baseColorFactor=[1.0, 1.0, 0.5, 1.0]),
+        alphaCutoff=None)
+    return create_gltf(np.array(vertices), indices, tex_coord, mode, material)
 
-    return vertices
+
+def _tessellate_vertex(ocp_vertex: TopoDS_Vertex) -> GLTF2:
+    """Tessellate a vertex into a list of triangle vertices and a list of triangle indices"""
+    c = Vertex(ocp_vertex).center()
+    vertices = np.array([[c.X, c.Y, c.Z]])
+    indices = np.array([0])
+    tex_coord = np.array([], dtype=np.float32)
+    mode = POINTS
+    material = Material(
+        pbrMetallicRoughness=PbrMetallicRoughness(baseColorFactor=[1.0, 0.5, 0.5, 1.0]),
+        alphaCutoff=None)
+    return create_gltf(vertices, indices, tex_coord, mode, material)