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

import enum
import io
import obja.obja as obja
import numpy as np
import argparse

from rich.progress import Progress


def cot(x: float) -> float:
    """Cotangeante of x

    Args:
        x (float): angle

    Returns:
        float: cotangeante of x
    """
    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) -> list[tuple]:
    """Create a sliding window of size n

    Args:
        l (list): list to create the sliding window from
        n (int, optional): number of value. Defaults to 2.

    Returns:
        list[tuple]: sliding window
    """
    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:
    """Edge representation"""
    def __init__(self, a: int, b: int) -> None:
        """Create an edge

        Args:
            a (int): first vertex
            b (int): second vertex
        """
        self.a = min(a, b)
        self.b = max(a, b)

        self.face1: Face | None = None
        self.face2: Face | None = None

        self.fold: float = 0.0
        self.curvature: float = 0.0

    def __eq__(self, __o: object) -> bool:
        """Check if two edges are equal

        Args:
            __o (object): other edge

        Returns:
            bool: True if equal, False otherwise
        """
        if isinstance(__o, Edge):
            return self.a == __o.a and self.b == __o.b

        return False


class Face:
    """Face representation"""
    def __init__(self, a: int, b: int, c: int) -> None:
        """Face constructor

        Args:
            a (int): first vertex index
            b (int): second vertex index
            c (int): third vertex index
        """
        self.a = a
        self.b = b
        self.c = c

        self.edges: list[Edge] = []
        self.a = a
        self.b = b
        self.c = c

        self.normal = np.zeros(3)

    def to_obja(self) -> obja.Face:
        """Convert face to obja format

        Returns:
            obja.Face: face in obja format
        """
        return obja.Face(self.a, self.b, self.c)

    def __eq__(self, __o: object) -> bool:
        """Check if two faces are equal

        Args:
            __o (object): other face

        Returns:
            bool: True if equal, False otherwise
        """
        if isinstance(__o, Face):
            return set((__o.a, __o.b, __o.c)) == set((self.a, self.b, self.c))

        return False


class Vertex:
    """Vertex representation"""
    def __init__(self, pos: np.ndarray[int, float]) -> None:
        """Vertex constructor

        Args:
            pos (np.ndarray[int, float]): position of the vertex
        """
        self.pos = pos

        self.vertex_ring: list[int] = []
        self.face_ring: list[int] = []

        self.normal: np.ndarray = np.zeros(3)

        self.area: float = 0.0
        self.curvature: float = 0.0

    def to_obja(self) -> np.ndarray:
        """Convert vertex to obja format

        Returns:
            np.ndarray: vertex in obja format
        """
        return self.pos


class MAPS(obja.Model):
    """MAPS compression model"""

    def __init__(self):
        """MAPS constructor"""
        super().__init__()

    def parse_file(self, path: str) -> None:
        """Parse a file

        Args:
            path (str): path to the file
        """
        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) -> None:
        """Precompute things"""
        # Reset progress bars
        self.progress.reset(self.select_task, total=len(self.vertices))
        self.progress.reset(self.compress_task)

        # Compute usefull things
        self.update_rings()
        self.update_edges()
        self.update_normals()
        self.update_area_curvature()

    def fix(self) -> None:
        """Fix the model"""

        fixed = True

        # Find a vertex with less than 3 faces
        for i, vertex in enumerate(self.vertices):
            if vertex is None:
                continue

            if len(vertex.face_ring) < 3:
                # Remove the vertex and its faces
                for face in vertex.face_ring:
                    if not self.final_only:
                        self.operations.append(
                            ('af', face, self.faces[face].to_obja()))
                    self.faces[face] = None
                if not self.final_only:
                    self.operations.append(
                        ('av', i, vertex.to_obja()))
                self.vertices[i] = None

                # Indicate that the model has to be fixed again
                fixed = False
        
        return fixed

    def update_edges(self) -> None:
        """Update edges"""

        self.edges = {}

        for face in self.faces:
            if face is None:
                continue

            # Create all edges for the face
            for a, b in sliding_window([face.a, face.b, face.c], n=2):
                new_edge = Edge(a, b)

                # Check if the edge already exists
                if f"{new_edge.a}:{new_edge.b}" in self.edges.keys():
                    continue

                new_edge.face1 = face

                # Find the opoosite face
                for face2_i in self.vertices[new_edge.a].face_ring:
                    face2 = self.faces[face2_i]
                    if face2 == face:
                        continue

                    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

                # Add the new edge to the list
                self.edges[f"{new_edge.a}:{new_edge.b}"] = new_edge

    def update_rings(self) -> None:
        """Update vertex and face rings"""

        fixed = False

        # Wait till the model is fixed
        while not fixed:
            for vertex in self.vertices:
                # Reset vertex ring
                if vertex is None:
                    continue
                vertex.face_ring = []

            # Add faces to vertex ring
            for i, face in enumerate(self.faces):
                if face is None:
                    continue
                for vertex_i in (face.a, face.b, face.c):
                    self.vertices[vertex_i].face_ring.append(i)

            # Fix the model
            fixed = self.fix()

        for i, vertex in enumerate(self.vertices):
            vertex = self.vertices[i]
            if vertex is None:
                continue
            
            # Compute rings
            ring = self.one_ring(i)
            vertex.vertex_ring = ring

    def update_area_curvature(self) -> None:
        """Update area and curvature"""

        # Get the area and curvature of each vertex
        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

        # Find feature edges
        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) -> None:
        """Update normals"""
        for face in self.faces:
            if face is None:
                continue

            # Compute face normal
            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)
            norm = np.linalg.norm(n)
            if norm != 0:
                n /= np.linalg.norm(n)

            face.normal = n

            # Sum vertex normal
            self.vertices[face.a].normal += n
            self.vertices[face.b].normal += n
            self.vertices[face.c].normal += n

        # Normalize vertex normals
        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) -> 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
        """

        ring_faces = [self.faces[i] for i in self.vertices[index].face_ring]

        # 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 Exception('Ring not corrupted')

        return ring

    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
        """

        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.5, 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 = []

        while len(vertices) > 0:
            # Select prefered vertex
            vertex = vertices.pop(0)  # remove it from remaining vertices
            self.progress.advance(self.select_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)
                            self.progress.advance(self.select_task)

        return selected_vertices

    def project_polar(self, index: int) -> tuple[list[np.ndarray], list[int]]:
        """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) -> list[obja.Face]:
        """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 = self.is_convex(prev_vert, curr_vert, next_vert)
                is_ear = True
                # the triangle needs to be convext to be an ear
                if is_convex or cycle_counter > len(indices):
                    # 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 self.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(self,
                  prev_vert: np.ndarray[int, float],
                  curr_vert: np.ndarray[int, float],
                  next_vert: np.ndarray[int, float],
                  ) -> 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(self,
                  a: np.ndarray[int, float],
                  b: np.ndarray[int, float],
                  c: np.ndarray[int, float],
                  p: np.ndarray[int, float],
                  ) -> 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 debug(self, output):

        self.update()
        priorities = self.compute_priority()

        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)

        self.operations = []
        for i, face in enumerate(self.faces):
            if face is None:
                continue
            r, g, b = colors[i] / max_c, 1.0, 1.0

            for feature_edge in self.feature_edges:
                face_vertices = (face.a, face.b, face.c)
                if feature_edge.a in face_vertices and feature_edge.b in face_vertices:
                    r, g, b = 1.0, 0.0, 0.0
                    break

            self.operations.append(('fc', i, (r, g, b)))
            self.operations.append(('af', i, face.to_obja()))

        for i, vertex in enumerate(self.vertices):
            if vertex is None:
                continue
            self.operations.append(('av', i, vertex.to_obja()))
        self.operations.reverse()

        # Write the result in output file
        output_model = obja.Output(output)

        for (op, index, value) in self.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(
                    len(output_model.face_mapping),
                    value[0],
                    value[1],
                    value[2]),
                    file=output
                )

    def compress(self, output: io.TextIOWrapper, level: int, final_only: bool, debug: bool) -> None:
        """Compress the 3d model

        Args:
            output (io.TextIOWrapper): Output file descriptor
        """

        with Progress() as progress:
            self.global_task = progress.add_task('╓ Global compression')
            self.select_task = progress.add_task('╟── Vertex selection')
            self.compress_task = progress.add_task('╙── Compression')
            self.progress = progress

            self.operations = []
            self.final_only = final_only

            # while len(self.vertices) > 64:
            for _ in progress.track(range(level), task_id=self.global_task):
                self.update()
                selected_vertices = self.select_vertices()  # find the set of vertices to remove

                for v_index in self.progress.track(selected_vertices, task_id=self.compress_task):

                    # 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 self.final_only:
                            self.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 self.final_only:
                            self.operations.append(
                                ('af', ring_faces[i], self.faces[ring_faces[i]].to_obja()))
                        self.faces[ring_faces[i]] = None

                    # Remove the vertex
                    if not self.final_only:
                        self.operations.append(
                            ('av', v_index, self.vertices[v_index].to_obja()))
                    self.vertices[v_index] = None

        if debug:
            self.debug(output)
            return

        # Register remaining vertices and faces
        for i, face in enumerate(self.faces):
            if face is not None:
                self.operations.append(('af', i, face.to_obja()))
        for i, v_index in enumerate(self.vertices):
            if v_index is not None:
                self.operations.append(('av', i, v_index.to_obja()))

        # To rebuild the model, run operations in reverse order
        self.operations.reverse()

        # Write the result in output file
        output_model = obja.Output(output)

        for (op, index, value) in self.operations:
            if op == 'av':
                output_model.add_vertex(index, value)
            elif op == 'af':
                try:
                    output_model.add_face(index, value)
                except:
                    print(self.vertices[value.b])
            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.level,
                       args.final or args.debug, args.debug)


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('-l', '--level', type=int, required=True)
    parser.add_argument('-f', '--final', type=bool, default=False)
    parser.add_argument('-d', '--debug', type=bool, default=False)
    args = parser.parse_args()

    main(args)