projet-compression-streamin.../main.py
2022-10-17 15:51:37 +02:00

701 lines
23 KiB
Python

import io
from math import floor
import obja.obja as obja
import numpy as np
import argparse
from time import time
from rich.progress import track, Progress
def cot(x: float):
sin_x = np.sin(x)
if sin_x == 0:
return 1e16
return np.cos(x) / sin_x
def sliding_window(l: list, n: int = 2):
k = n - 1
l2 = l + [l[i] for i in range(k)]
res = [(x for x in l2[i:i+n]) for i in range(len(l2)-k)]
return res
class Edge:
def __init__(self, a, b):
self.a = min(a, b)
self.b = max(a, b)
self.face1 = None
self.face2 = None
self.fold = 0.0
self.curvature = 0.0
def __eq__(self, __o: object) -> bool:
return self.a == __o.a and self.b == __o.b
class Face:
def __init__(self, a, b, c):
self.a = a
self.b = b
self.c = c
self.normal = np.zeros(3)
def to_obja(self):
return obja.Face(self.a, self.b, self.c)
def __eq__(self, __o: object) -> bool:
if __o is None:
return False
return self.a == __o.a and self.b == __o.b and self.c == __o.c
class Vertex:
def __init__(self, pos):
self.pos = pos
self.vertex_ring = []
self.face_ring = []
self.normal = np.zeros(3)
self.area = 0.0
self.curvature = 0.0
def to_obja(self):
return self.pos
class MAPS(obja.Model):
"""_summary_
Args:
obja (_type_): _description_
"""
def __init__(self):
super().__init__()
def parse_file(self, path):
super().parse_file(path)
for i, vertex in enumerate(self.vertices):
self.vertices[i] = Vertex(vertex)
for i, face in enumerate(self.faces):
self.faces[i] = Face(face.a, face.b, face.c)
def update(self):
self.update_edges()
self.update_rings()
self.update_normals()
self.update_area_curvature()
def update_edges(self):
self.edges = {}
remaining_faces = self.faces.copy()
while None in remaining_faces:
remaining_faces.remove(None)
for face in track(self.faces, description='Update edges'):
if face is None:
continue
for a, b in sliding_window([face.a, face.b, face.c], n=2):
new_edge = Edge(a, b)
if self.edges.get(f"{new_edge.a}:{new_edge.b}") is None:
new_edge.face1 = face
if face in remaining_faces:
remaining_faces.remove(face)
for face2 in remaining_faces:
face2_vertices = (face2.a, face2.b, face2.c)
if not (a in face2_vertices and b in face2_vertices):
continue
new_edge.face2 = face2
break
if new_edge.face2 is None:
print('ooooooooooooooooooooooo')
self.edges[f"{new_edge.a}:{new_edge.b}"] = new_edge
def update_rings(self):
for i, vertex in enumerate(self.vertices):
if vertex is None:
continue
vertex_ring, face_ring = self.one_ring(i)
vertex.vertex_ring = vertex_ring
vertex.face_ring = face_ring
def update_area_curvature(self):
for i, vertex in enumerate(self.vertices):
if vertex is None:
continue
area, curvature = self.compute_area_curvature(i)
vertex.area = area
vertex.curvature = curvature
self.feature_edges = []
for edge in self.edges.values():
edge.fold = np.dot(edge.face1.normal, edge.face2.normal)
if edge.fold < 0.5:
self.feature_edges.append(edge)
def update_normals(self):
for face in self.faces:
if face is None:
continue
p1 = self.vertices[face.a].pos
p2 = self.vertices[face.b].pos
p3 = self.vertices[face.c].pos
u = p2 - p1
v = p3 - p1
n = np.cross(u, v)
n /= np.linalg.norm(n)
face.normal = n
self.vertices[face.a].normal += n
self.vertices[face.b].normal += n
self.vertices[face.c].normal += n
for vertex in self.vertices:
if vertex is None:
continue
norm = np.linalg.norm(vertex.normal)
if norm != 0:
vertex.normal /= norm
def one_ring(self, index: int) -> tuple[list[int], list[int]]:
""" Return the corresponding 1-ring
Args:
index (int): index of the 1-ring's main vertex
Returns:
list[int]: ordered list of the 1-ring vertices
"""
if self.vertices[index] is None:
return None, None
# Find the 1-ring faces
ring_faces, ring_face_indices = [], []
for face_index, face in enumerate(self.faces):
if face is None:
continue
if index in (face.a, face.b, face.c):
ring_faces.append(face)
ring_face_indices.append(face_index)
# Initialize the ring
start_index = (ring_faces[0].a if ring_faces[0].a != index and ring_faces[0].c != index else
ring_faces[0].b if ring_faces[0].a != index and ring_faces[0].b != index else
ring_faces[0].c)
ring = [start_index]
ring_faces.pop(0)
# Select the indexes of the ring in the right order
while len(ring_faces) > 0:
broke = False
prev_index = ring[-1]
for i, face in enumerate(ring_faces):
if prev_index in (face.a, face.b, face.c):
# Found the face that correspond to the next vertex
current_index = ( # select the next vertex from the face
face.a if face.a != index and face.a != prev_index else
face.b if face.b != index and face.b != prev_index else
face.c
)
ring.append(current_index)
ring_faces.pop(i)
broke = True
break
if not broke:
raise ValueError(
f"Vertex {prev_index} is not in the remaining faces {ring_faces}. Origin {ring} on {index}")
return ring, ring_face_indices
def compute_area_curvature(self, index: int) -> tuple[float, float]:
""" Compute area and curvature the corresponding 1-ring
Args:
index (int): index of the 1-ring's main vertex
Returns:
tuple[float, float]: area and curvature
"""
if self.vertices[index] is None:
return None, None
ring = self.vertices[index].vertex_ring
p1 = self.vertices[index].pos
n1 = self.vertices[index].normal
area_sum = 0
curvature = 0
for index1, index2 in sliding_window(ring, n=2):
# the second vertice of the triangle
p2 = self.vertices[index1].pos
p3 = self.vertices[index2].pos # the third vertice of the triangle
n2 = self.vertices[index1].normal
M = np.array([ # build the matrix, used to compute the area
[p1[0], p2[0], p3[0]],
[p1[1], p2[1], p3[1]],
[p1[2], p2[2], p3[2]],
])
area = abs(np.linalg.det(M) / 2) # compute the area
area_sum += area
edge_curvature = np.dot(n2 - n1, p2 - p1) / \
np.linalg.norm(p2 - p1)**2
edge_curvature = abs(edge_curvature)
edge_key = f"{min(index, index1)}:{max(index, index1)}"
self.edges[edge_key].curvature = edge_curvature
curvature += edge_curvature
curvature /= len(ring)
return area_sum, curvature
def compute_priority(self, lamb: float = 0.0, max_length: int = 12) -> list[float]:
""" Compute selection priority of vertices (0.0 -> hight priority ; 1.0 -> low priority)
Args:
lamb (float, optional): convex combination factor. Defaults to 0.5.
max_length (int, optional): 1-ring maximum length to be prioritary. Defaults to 12.
Returns:
list[float]: priority values
"""
max_area = max(
[vertex.area for vertex in self.vertices if vertex is not None])
max_curvature = max(
[vertex.curvature for vertex in self.vertices if vertex is not None])
# Compute priorities
priorities = []
for vertex in self.vertices:
if vertex is not None and len(vertex.vertex_ring) < max_length:
# Compute priority
priority = (
lamb * vertex.area / max_area +
(1.0 - lamb) * vertex.curvature / max_curvature
)
else:
# Vertex with low priority
priority = 2.0
priorities.append(priority)
return priorities
def select_vertices(self) -> list[int]:
""" Select vertices for the current level reduction
Returns:
list[int]: selected vertices
"""
# Order vertices by priority
priorities = self.compute_priority()
vertices = [i[0]
for i in sorted(enumerate(priorities), key=lambda p: p[1])]
selected_vertices = []
with Progress() as progress:
task = progress.add_task('Select vertices', total=len(vertices))
while not progress.finished:
# Select prefered vertex
vertex = vertices.pop(0) # remove it from remaining vertices
progress.advance(task)
if priorities[vertex] == 2.0:
continue
incident_count = 0
for feature_edge in self.feature_edges:
if vertex in (feature_edge.a, feature_edge.b):
incident_count += 1
if incident_count > 2:
continue
selected_vertices.append(vertex)
# Remove neighbors
# for face in remaining_faces:
for face in self.faces:
if face is None:
continue
face_vertices = (face.a, face.b, face.c)
if vertex in face_vertices:
# Remove face and face's vertices form remainings
# remaining_faces.remove(face)
for face_vertex in face_vertices:
if face_vertex in vertices:
vertices.remove(face_vertex)
progress.advance(task)
return selected_vertices[:floor(1.0 * len(selected_vertices))]
def project_polar(self, index: int) -> list[np.ndarray]:
""" Flatten the 1-ring to retriangulate
Args:
index (int): main vertex of the 1-ring
Returns:
list[np.ndarray]: list the cartesian coordinates of the flattened 1-ring projected in the plane
"""
ring = self.vertices[index].vertex_ring
radius, angles = [], []
teta = 0.0 # cumulated angles
for index1, index2 in sliding_window(ring):
r = np.linalg.norm(
self.vertices[index].pos - self.vertices[index1].pos)
teta += self.compute_angle(index1, index, index2) # add new angle
radius.append(r)
angles.append(teta)
angles = [2 * np.pi * a / teta for a in angles] # normalize angles
coordinates = [np.array([r * np.cos(a), r * np.sin(a)])
for r, a in zip(radius, angles)] # parse polar to cartesian
return coordinates, ring
def compute_angle(self, i: int, j: int, k: int) -> float:
""" Calculate the angle defined by three points
Args:
i (int): previous index
j (int): central index
k (int): next index
Returns:
float: angle defined by the three points
"""
a = self.vertices[i].pos
b = self.vertices[j].pos
c = self.vertices[k].pos
u = a - b
v = c - b
u /= np.linalg.norm(u)
v /= np.linalg.norm(v)
res = np.dot(u, v)
return np.arccos(np.clip(res, -1, 1))
def clip_ear(self, index: int) -> tuple[list[obja.Face], int]:
""" Retriangulate a polygon using the ear clipping algorithm
Args:
index (int): index of 1-ring
Returns:
tuple[list[obja.Face], int]: list the triangles
"""
polygon_, ring_ = self.project_polar(index)
main_v = []
for i, r in enumerate(ring_):
for feature_edge in self.feature_edges:
feat_edge_vertices = (feature_edge.a, feature_edge.b)
if r in feat_edge_vertices and index in feat_edge_vertices:
main_v.append(i)
if len(main_v) < 2:
polygons_rings = [(polygon_, ring_)]
else:
v1 = ring_[main_v[0]]
v2 = ring_[main_v[1]]
ring1, ring2 = [], []
polygon1, polygon2, = [], []
start = ring_.index(v1)
while ring_[start] != v2:
ring1.append(ring_[start])
polygon1.append(polygon_[start])
start += 1
start %= len(ring_)
ring1.append(ring_[start])
polygon1.append(polygon_[start])
start = ring_.index(v2)
while ring_[start] != v1:
ring2.append(ring_[start])
polygon2.append(polygon_[start])
start += 1
start %= len(ring_)
ring2.append(ring_[start])
polygon2.append(polygon_[start])
polygons_rings = [(polygon1, ring1), (polygon2, ring2)]
faces = [] # the final list of faces
for polygon, ring in polygons_rings:
indices = [(local_i, global_i)
for local_i, global_i in enumerate(ring)] # remainging vertices
node_index = 0
cycle_counter = 0
while len(indices) > 2:
# Extract indices
local_i, global_i = indices[node_index - 1]
local_j, global_j = indices[node_index]
local_k, global_k = indices[node_index + 1]
# Extract verticies
prev_vert = polygon[local_i]
curr_vert = polygon[local_j]
next_vert = polygon[local_k]
is_convex = MAPS.is_convex(prev_vert, curr_vert, next_vert)
is_ear = True
if is_convex or cycle_counter > len(indices): # the triangle needs to be convext to be an ear
# Begin with the point next to the triangle
test_node_index = (node_index + 2) % len(indices)
while indices[test_node_index][0] != local_i and is_ear:
test_vert = polygon[indices[test_node_index][0]]
is_ear = not MAPS.is_inside(prev_vert,
curr_vert,
next_vert,
test_vert)
test_node_index = (test_node_index + 1) % len(indices)
else:
is_ear = False
cycle_counter += 1
if is_ear:
faces.append(Face(global_i, global_j, global_k))
indices.pop(node_index) # remove the point from the ring
cycle_counter = 0
node_index = (node_index + 2) % len(indices) - 1
return faces
def is_convex(prev_vert: np.ndarray, curr_vert: np.ndarray, next_vert: np.ndarray) -> bool:
""" Check if the angle less than pi
Args:
prev_vert (np.ndarray): first point
curr_vert (np.ndarray): middle point
next_vert (np.ndarray): last point
Returns:
bool: angle smaller than pi
"""
a = prev_vert - curr_vert
b = next_vert - curr_vert
dot = a[0] * b[0] + a[1] * b[1]
det = a[0] * b[1] - a[1] * b[0]
angle = np.arctan2(det, dot)
if angle < 0.0:
angle = 2.0 * np.pi + angle
internal_angle = angle
return internal_angle >= np.pi
def is_inside(a: np.ndarray, b: np.ndarray, c: np.ndarray, p: np.ndarray) -> bool:
""" Check if p is in the triangle a b c
Args:
a (np.ndarray): point one
b (np.ndarray): point two
c (np.ndarray): point three
p (np.ndarray): point to check
Returns:
bool: if the point to check is in a b c
"""
# Compute vectors
v0 = c - a
v1 = b - a
v2 = p - a
# Compute dot products
dot00 = np.dot(v0, v0)
dot01 = np.dot(v0, v1)
dot02 = np.dot(v0, v2)
dot11 = np.dot(v1, v1)
dot12 = np.dot(v1, v2)
# Compute barycentric coordinates
denom = dot00 * dot11 - dot01 * dot01
if abs(denom) < 1e-20:
return True
invDenom = 1.0 / denom
u = (dot11 * dot02 - dot01 * dot12) * invDenom
v = (dot00 * dot12 - dot01 * dot02) * invDenom
# Check if point is in triangle
return (u >= 0) and (v >= 0) and (u + v < 1)
def truc(self, output):
self.update()
priorities = self.compute_priority()
# min_p = min(priorities)
# priorities = [p - min_p for p in priorities]
# max_p = max(priorities)
colors = [priorities[face.a] + priorities[face.b] +
priorities[face.c] if face is not None else 0.0 for face in self.faces]
min_c = min(colors)
colors = [c - min_c for c in colors]
max_c = max(colors)
operations = []
for i, face in enumerate(self.faces):
if face != None:
r, g, b = colors[i] / max_c, 1.0, 1.0
c = 0
for x in (face.a, face.b, face.c):
for feature_edge in self.feature_edges:
if x in feature_edge:
c += 1
break
if c > 1:
r, g, b = 1.0, 0.0, 0.0
operations.append(('fc', i, (r, g, b), 0, 0, 0))
operations.append(('af', i, face, 0, 0, 0))
for i, vertex in enumerate(self.vertices):
if vertex is None:
# r, g, b = priorities[i] / max_p , 1.0, 1.0
operations.append(('av', i, vertex, 1.0, 1.0, 1.0))
operations.reverse()
# Write the result in output file
output_model = obja.Output(output)
for (op, index, value, r, g, b) in operations:
if op == 'av':
output_model.add_vertex_rgb(index, value, r, g, b)
elif op == 'af':
output_model.add_face(index, value)
elif op == 'ev':
output_model.edit_vertex(index, value)
elif op == 'ef':
output_model.edit_face(index, value)
elif op == 'fc':
print('fc {} {} {} {}'.format(
len(output_model.face_mapping),
value[0],
value[1],
value[2]),
file=output
)
def compress(self, output: io.TextIOWrapper, final_only: bool) -> None:
""" Compress the 3d model
Args:
output (io.TextIOWrapper): Output file descriptor
"""
operations = []
# while len(self.vertices) > 64:
for _ in range(2):
self.update()
selected_vertices = self.select_vertices() # find the set of vertices to remove
for v_index in track(selected_vertices, description="Compression"):
# Extract ring faces
ring_faces = self.vertices[v_index].face_ring
# Apply retriangulation algorithm
faces = self.clip_ear(v_index)
# Edit the first faces
for i in range(len(faces)):
if not final_only:
operations.append(
('ef', ring_faces[i], self.faces[ring_faces[i]].to_obja()))
self.faces[ring_faces[i]] = faces[i]
# Remove the last faces
for i in range(len(faces), len(ring_faces)):
if not final_only:
operations.append(
('af', ring_faces[i], self.faces[ring_faces[i]].to_obja()))
self.faces[ring_faces[i]] = None
# Remove the vertex
if not final_only:
operations.append(
('av', v_index, self.vertices[v_index].to_obja()))
self.vertices[v_index] = None
# Register remaining vertices and faces
for i, face in enumerate(self.faces):
if face is not None:
operations.append(('af', i, face.to_obja()))
for i, v_index in enumerate(self.vertices):
if v_index is not None:
operations.append(('av', i, v_index.to_obja()))
# To rebuild the model, run operations in reverse order
operations.reverse()
# Write the result in output file
output_model = obja.Output(output)
for (op, index, value) in operations:
if op == 'av':
output_model.add_vertex(index, value)
elif op == 'af':
output_model.add_face(index, value)
elif op == 'ev':
output_model.edit_vertex(index, value)
elif op == 'ef':
output_model.edit_face(index, value)
elif op == 'fc':
print('fc {} {} {} {}'.format(
index,
value[0],
value[1],
value[2]),
file=output
)
def main(args):
""" Run MAPS model compression
Args:
args (Namespace): arguments (input and output path)
"""
model = MAPS()
model.parse_file(args.input)
with open(args.output, 'w') as output:
model.compress(output, args.final)
# with open(args.output, 'w') as output:
# model.truc(output)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input', type=str, required=True)
parser.add_argument('-o', '--output', type=str, required=True)
parser.add_argument('-f', '--final', type=bool, default=False)
args = parser.parse_args()
main(args)