from __future__ import annotations
import math
import struct
import numpy as np
import numpy.typing as npt
from earcut.earcut import earcut
CoordinateType = npt.NDArray[np.float32]
LineType = list[CoordinateType]
PolygonType = list[LineType]
MultiPolygonsType = list[PolygonType]
PolygonAsTriangleType = list[npt.NDArray[np.float32]] # the array shape is 3, 3
[docs]
class TriangleSoup:
def __init__(self) -> None:
self.triangles: list[PolygonAsTriangleType] = []
@property
def vertices(self) -> npt.NDArray[np.float32]:
"""Extract the unique vertices that compose the triangle set."""
return np.unique(np.concatenate(self.triangles[0]), axis=0)
@property
def triangle_indices(self) -> npt.NDArray[np.uint8]:
"""Express the triangles as triplets of vertex indices."""
flattened_triangles = np.concatenate(self.triangles[0])
indices = np.full(flattened_triangles.shape[0], dtype=np.uint8, fill_value=-1)
for vertex_idx in range(self.vertices.shape[0]):
indices[
np.all(flattened_triangles == self.vertices[vertex_idx], axis=1)
] = vertex_idx
return indices.reshape(-1, 3)
[docs]
@staticmethod
def from_wkb_multipolygon(
wkb: bytes, associated_data: list[bytes] | None = None
) -> TriangleSoup:
"""
:param wkb: Well-Known Binary binary string describing a multipolygon
:param associated_data: array of multipolygons containing data attached to the wkb
parameter multipolygon. Must be the same size as wkb.
"""
multipolygons: list[MultiPolygonsType] = [parse(bytes(wkb))]
if associated_data is None:
associated_data = []
for additional_wkb in associated_data:
multipolygons.append(parse(bytes(additional_wkb)))
triangles_array: list[PolygonAsTriangleType] = [
[] for _ in range(len(multipolygons))
]
for i in range(len(multipolygons[0])):
polygon: PolygonType = multipolygons[0][i]
additional_polygons = [mp[i] for mp in multipolygons[1:]]
triangles = triangulate(polygon, additional_polygons)
for array, tri in zip(triangles_array, triangles):
array += tri
ts = TriangleSoup()
ts.triangles = triangles_array
return ts
[docs]
def get_position_array(self) -> bytes:
"""
Returns a binary array of vertex positions
"""
vertex_triangles = self.triangles[0]
vertex_array = vertex_attribute_to_array(vertex_triangles)
return b"".join([vertex.tobytes() for vertex in vertex_array])
[docs]
def get_data_array(self, index: int) -> bytes:
"""
Returns a binary array of vertex data
"""
vertex_triangles = self.triangles[1 + index]
vertex_array = vertex_attribute_to_array(vertex_triangles)
return b"".join([vertex.tobytes() for vertex in vertex_array])
[docs]
def get_data(self, index: int) -> npt.NDArray[np.float32]:
vertex_triangles = self.triangles[1 + index]
return np.array(vertex_attribute_to_array(vertex_triangles))
[docs]
def get_normal_array(self) -> bytes:
"""
Returns a binary array of vertex normals
"""
vertex_array = list(self.compute_normals())
return b"".join([vertex.tobytes() for vertex in vertex_array])
[docs]
def compute_normals(self) -> npt.NDArray[np.float32]:
"""Compute vertex normals and returns them as a numpy array."""
normals: list[CoordinateType] = []
for t in self.triangles[0]:
U = t[1] - t[0]
V = t[2] - t[0]
N = np.cross(U, V)
norm = np.linalg.norm(N)
if norm == 0:
normals.append(np.array([0, 0, 1], dtype=np.float32))
else:
normals.append(N / norm)
return np.array(face_attribute_to_array(normals))
[docs]
def get_bbox(self) -> list[npt.NDArray[np.float32]]:
"""
Returns the bbox in this format: [[minX, minY, minZ],[maxX, maxY, maxZ]]
"""
mins = np.array([np.min(t, 0) for t in self.triangles[0]])
maxs = np.array([np.max(t, 0) for t in self.triangles[0]])
return [np.min(mins, 0), np.max(maxs, 0)]
[docs]
def face_attribute_to_array(triangles: list[CoordinateType]) -> list[CoordinateType]:
array = []
for face in triangles:
array += [face, face, face]
return array
[docs]
def vertex_attribute_to_array(triangles: PolygonAsTriangleType) -> list[CoordinateType]:
array = []
for face in triangles:
for vertex in face:
array.append(vertex)
return array
[docs]
def parse(wkb: bytes) -> MultiPolygonsType:
multipolygon: MultiPolygonsType = []
byteorder = struct.unpack("b", wkb[0:1])
bo = "<" if byteorder[0] else ">"
geomtype = struct.unpack(bo + "I", wkb[1:5])[0]
has_z = (geomtype == 1006) or (geomtype == 1015)
# MultipolygonZ or polyhedralSurface
pnt_offset = 24 if has_z else 16
pnt_unpack = "ddd" if has_z else "dd"
geom_nb = struct.unpack(bo + "I", wkb[5:9])[0]
# print(struct.unpack('b', wkb[9:10])[0])
# print(struct.unpack('I', wkb[10:14])[0]) # 1003 (Polygon)
# print(struct.unpack('I', wkb[14:18])[0]) # num lines
# print(struct.unpack('I', wkb[18:22])[0]) # num points
offset = 9
for _ in range(geom_nb):
offset += 5 # struct.unpack('bI', wkb[offset:offset + 5])[0]
# 1 (byteorder), 1003 (Polygon)
line_nb = struct.unpack(bo + "I", wkb[offset : offset + 4])[0]
offset += 4
polygon = []
for _ in range(line_nb):
point_nb = struct.unpack(bo + "I", wkb[offset : offset + 4])[0]
offset += 4
line = []
for _ in range(point_nb - 1):
pt = np.array(
struct.unpack(bo + pnt_unpack, wkb[offset : offset + pnt_offset]),
dtype=np.float32,
)
offset += pnt_offset
line.append(pt)
offset += pnt_offset # skip redundant point
polygon.append(line)
multipolygon.append(polygon)
return multipolygon
[docs]
def triangulate(
polygon: PolygonType, additional_polygons: list[PolygonType] | None = None
) -> list[PolygonAsTriangleType]:
"""
Triangulates 3D polygons
"""
# let's find out if the polygon is *mostly* clockwise or counter-clockwise
# and triangulate accordingly
# for 2D explanations:
# https://stackoverflow.com/a/1165943/1528985
# and https://www.element84.com/blog/determining-the-winding-of-a-polygon-given-as-a-set-of-ordered-points
#
# Quick explanation in case it goes down: for each edge we calculate the
# area of the polygon formed by this edge, the x axis and the 2 vertical.
# It's (x2-x1) / ((y2+y1 / 2) (draw it if you don't believe me). This
# results will be positive for a edge that goes toward positive x. Summing
# all these areas will give plus or minus the total polygon area. it would
# be positive for a clockwise polygon (upper edges contributing positively)
# and negative for counter-clockwise polygons (upper edges contributing
# negatively)
#
# Adaptations here:
# - we prefer to reason with counter-clockwise positive, hence the x1-x2 instead of x2-x1
# - in 3D, we calcule this value for each axis planes (xy, yz, zx),
# looking in the other axis negative direction.
# - comparing these 3 results actually give us the most interesting plane
# to triangulate in (the plane were the projected area is the biggest)
# - we drop the 1/2 factor because we are only interesting in the sign and relative comparison
vect_prod = np.array([0, 0, 0], dtype=np.float32)
for i in range(len(polygon[0])):
curr_edge = polygon[0][i]
next_edge = polygon[0][(i + 1) % len(polygon[0])]
vect_prod += np.array(
[
# yz plane, seen from negative x
(curr_edge[1] - next_edge[1]) * (next_edge[2] + curr_edge[2]),
# zx plane, seen from negative y
(curr_edge[2] - next_edge[2]) * (next_edge[0] + curr_edge[0]),
# xy plane, seen from negative z
(curr_edge[0] - next_edge[0]) * (next_edge[1] + curr_edge[1]),
],
dtype=np.float32,
)
if additional_polygons is None:
additional_polygons = []
polygon_2d = []
holes = []
delta = 0
for p in polygon[:-1]:
holes.append(delta + len(p))
delta += len(p)
# triangulation of the polygon projected on planes (xy) (zx) or (yz)
if math.fabs(vect_prod[0]) > math.fabs(vect_prod[1]) and math.fabs(
vect_prod[0]
) > math.fabs(vect_prod[2]):
# (yz) projection
for linestring in polygon:
for point in linestring:
polygon_2d.extend([point[1], point[2]])
elif math.fabs(vect_prod[1]) > math.fabs(vect_prod[2]):
# (zx) projection
for linestring in polygon:
for point in linestring:
polygon_2d.extend([point[0], point[2]])
else:
# (xy) projection
for linestring in polygon:
for point in linestring:
polygon_2d.extend([point[0], point[1]])
triangles_idx = earcut(polygon_2d, holes, 2)
arrays: list[PolygonAsTriangleType] = [
[] for _ in range(len(additional_polygons) + 1)
]
for i in range(0, len(triangles_idx), 3):
t = triangles_idx[i : i + 3]
p0 = unflatten(polygon, holes, t[0])
p1 = unflatten(polygon, holes, t[1])
p2 = unflatten(polygon, holes, t[2])
# triangulation may break triangle orientation, test it before
# adding triangles
# FIXME fix / change the triangulation code instead?
cross_product = np.cross(p1 - p0, p2 - p0)
invert = np.dot(vect_prod, cross_product) < 0
if invert:
arrays[0].append(np.array([p1, p0, p2], dtype=np.float32))
else:
arrays[0].append(np.array([p0, p1, p2], dtype=np.float32))
for array, additional_polygon in zip(arrays[1:], additional_polygons):
pp0 = unflatten(additional_polygon, holes, t[0])
pp1 = unflatten(additional_polygon, holes, t[1])
pp2 = unflatten(additional_polygon, holes, t[2])
if invert:
array.append(np.array([pp1, pp0, pp2], dtype=np.float32))
else:
array.append(np.array([pp0, pp1, pp2], dtype=np.float32))
return arrays
[docs]
def unflatten(array: PolygonType, lengths: list[int], index: int) -> CoordinateType:
for i in reversed(range(len(lengths))):
lgth = lengths[i]
if index >= lgth:
return array[i + 1][index - lgth]
return array[0][index]