projet-compression-streaming-interaction-3D/main.py

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)