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):
    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:
        if isinstance(__o, Edge):
            return self.a == __o.a and self.b == __o.b

        return False


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 isinstance(__o, Face):
            return set((__o.a, __o.b, __o.c)) == set((self.a, self.b, self.c))

        return False


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.progress.reset(self.select_task, total=len(self.vertices))
        self.progress.reset(self.compress_task)
        self.update_rings()
        self.update_edges()
        self.update_normals()
        self.update_area_curvature()

    def fix(self):
        fixed = True
        for i, vertex in enumerate(self.vertices):
            if vertex is None:
                continue
            if len(vertex.face_ring) < 3:
                for face in vertex.face_ring:
                    self.faces[face] = None
                self.vertices[i] = None
                fixed = False
        return fixed
                

    def update_edges(self):
        self.edges = {}

        for face in self.faces:
            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 f"{new_edge.a}:{new_edge.b}" not in self.edges.keys():
                    new_edge.face1 = 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

                    self.edges[f"{new_edge.a}:{new_edge.b}"] = new_edge

    def update_rings(self):
        try:
            fixed = False
            while not fixed:
                for vertex in self.vertices:
                    if vertex is None:
                        continue
                    vertex.face_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)
                
                fixed = self.fix()

            for i, vertex in enumerate(self.vertices):
                vertex = self.vertices[i]
                if vertex is None:
                    continue
                if len(vertex.face_ring) == 0:
                    self.vertices[i] = None
                    continue
                ring = self.one_ring(i)
                vertex.vertex_ring = ring
        except ValueError:
            self.update_rings()
        

    def fail(self, index):
        print('fail')
        output_file = open('obja/example/fail.obja', 'w')
        output = obja.Output(output_file)
        used = []
        for i, x in enumerate(self.vertices[index].face_ring):
            face = self.faces[x]
            for y in (face.a, face.b, face.c):
                if y in used:
                    continue
                output.add_vertex(y, self.vertices[y].to_obja())
                used.append(y)
            output.add_face(x, face.to_obja())
            print('fc {} {} {} {}'.format(
                i + 1,
                np.random.rand(),
                np.random.rand(),
                np.random.rand()),
                file=output_file
            )
            print(x, (face.a, face.b, face.c))

    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():
            if edge.face2 is None:
                self.fail(edge.b)
            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) -> 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:
                self.fail(index)
                
                for i, face_i in enumerate(self.vertices[index].face_ring):
                    for face_j in self.vertices[index].face_ring[i+1:]:
                        face1 = self.faces[face_i]
                        face2 = self.faces[face_j]
                        if face1 == face2:
                            self.faces[face_i] = None
                            self.faces[face_j] = None
                            verts = (face1.a, face1.b, face1.c)
                            for vert in verts:
                                if vert == index:
                                    continue
                                to_remove = True
                                for face_k in self.vertices[vert].face_ring:
                                    face = self.faces[face_k]
                                    if face is None:
                                        continue
                                    if vert in (face.a, face.b, face.c):
                                        to_remove = False
                                        break
                                if to_remove:
                                    self.vertices[vert] = None
                                    break
                            break
                                        

                raise ValueError(
                    f"Vertex {prev_index} is not in the remaining faces {ring_faces}. Origin {ring} on {index}")

        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, np.dtype[np.float64]],
                  curr_vert: np.ndarray[int, np.dtype[np.float64]],
                  next_vert: np.ndarray[int, np.dtype[np.float64]]
                  ) -> 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, np.dtype[np.float64]],
                  b: np.ndarray[int, np.dtype[np.float64]],
                  c: np.ndarray[int, np.dtype[np.float64]],
                  p: np.ndarray[int, np.dtype[np.float64]]
                  ) -> 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)

        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

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

        for i, vertex in enumerate(self.vertices):
            if vertex is None:
                continue
            operations.append(('av', i, vertex.to_obja()))
        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(
                    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

            operations = []

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

        if debug:
            self.debug(output)
            return
        
        # 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':
                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)