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

import io
from types import NoneType
import obja.obja as obja
import numpy as np
import argparse

from rich.progress import track

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 MAPS(obja.Model):
    """_summary_

    Args:
        obja (_type_): _description_
    """

    def __init__(self):
        super().__init__()
        self.deleted_faces = set()

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

        # Find the 1-ring faces
        ring_faces, ring_face_indices = [], []
        for face_index, face in enumerate(self.faces):
            if face == 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 else ring_faces[0].b
        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, int]:
        """ 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
        """
        area_sum = 0
        laplace_sum = 0
        one_ring_vertices, _ = self.one_ring(index)

        teta = 0.0
        p1 = self.vertices[index]  # the center of the one-ring
        for index1, index2, index3 in sliding_window(one_ring_vertices, n=3):
            p2 = self.vertices[index1]  # the second vertice of the triangle
            p3 = self.vertices[index2]  # the third vertice of the triangle
            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

            teta += self.compute_angle(index1, index, index2)

            laplace = self.compute_laplace(index, index1, index2, index3)
            laplace_sum += laplace

        K = (2 * np.pi - teta) / area * 3
        H = np.linalg.norm(laplace_sum) / 4 / area * 3
        curvature = abs(H - np.sqrt(H*H - K)) + abs(H + np.sqrt(H*H - K))
        curvature = K

        return area_sum, curvature, len(one_ring_vertices)

    def compute_laplace(self, i: int, j: int, a: int, b: int) -> np.ndarray:
        alpha = self.compute_angle(i, a, j)
        beta = self.compute_angle(i, b, j)
        cot_sum = cot(alpha) + cot(beta)
        vec = self.vertices[j] - self.vertices[i]
        return cot_sum * vec

    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
        """
        # Compute area and curvature for each vertex
        areas, curvatures, ring_lengths = [], [], []
        for i in range(len(self.vertices)):
            if type(self.vertices[i]) != NoneType:
                area, curvature, ring_length = self.compute_area_curvature(i)
            else:
                area, curvature, ring_length = -1.0, -1.0, None
            areas.append(area)
            curvatures.append(curvature)
            ring_lengths.append(ring_length)
        # Get maxes to normalize
        max_area = max(areas)
        max_curvature = max(curvatures)

        # Compute priorities
        priorities = []
        for a, k, l in zip(areas, curvatures, ring_lengths):
            if l != None and l <= max_length:
                # Compute priority
                priority = lamb * a / max_area + \
                    (1.0 - lamb) * k / max_curvature
            else:
                # Vertex with low priority
                priority = 2.0
            priorities.append(priority)

        # return np.random.rand(len(priorities))
        return priorities

    def select_vertices(self) -> list[int]:
        """ Select vertices for the current level reduction

        Returns:
            list[int]: selected vertices
        """
        print("Selecting 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
            if priorities[vertex] == 2.0:
                continue
            selected_vertices.append(vertex)

            # Remove neighbors
            # for face in remaining_faces:
            for face in self.faces:
                if face == 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)

        print("Vertices selected.")
        return 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.one_ring(index)
        radius, angles = [], []
        teta = 0.0  # cumulated angles
        for index1, index2 in sliding_window(ring):
            r = np.linalg.norm(self.vertices[index] - self.vertices[index1])
            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]
        b = self.vertices[j]
        c = self.vertices[k]
        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)

        faces = []  # the final list of faces

        indices = [(local_i, global_i)
                   for local_i, global_i in enumerate(ring)]  # remainging vertices

        node_index = 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:  # 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

            if is_ear:
                faces.append(obja.Face(global_i, global_j, global_k))
                indices.pop(node_index)  # remove the point from the ring

            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):
        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] 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
                # operations.append(('fc', i, (r, g, b)))
                operations.append(('af', i, face, 0, 0, 0))
        for i, vertex in enumerate(self.vertices):
            if type(vertex) != NoneType:
                r, g, b = priorities[i] / max_p , 1.0, 1.0
                operations.append(('av', i, vertex, r, g, b))
        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 contract(self, output: io.TextIOWrapper) -> None:
        """ Compress the 3d model

        Args:
            output (io.TextIOWrapper): Output file descriptor
        """
        operations = []

        # while len(self.vertices) > 64:
        for i in range(1):
            selected_vertices = self.select_vertices()  # find the set of vertices to remove

            for vertex in track(selected_vertices, description="compression"):
                # print(f"     {len(selected_vertices)}      ", end='\r')

                # Extract ring faces
                _, ring_faces = self.one_ring(vertex)

                # Apply retriangulation algorithm
                faces = self.clip_ear(vertex)

                # Edit the first faces
                for i in range(len(faces)):
                    # operations.append(
                    #     ('ef', ring_faces[i], self.faces[ring_faces[i]]))
                    self.faces[ring_faces[i]] = faces[i]

                # Remove the last faces
                for i in range(len(faces), len(ring_faces)):
                    # operations.append(
                    #     ('af', ring_faces[i], self.faces[ring_faces[i]]))
                    self.faces[ring_faces[i]] = None

                # Remove the vertex
                # operations.append(('av', vertex, self.vertices[vertex]))
                self.vertices[vertex] = None

        # Register remaining vertices and faces
        for i, face in enumerate(self.faces):
            if face != None:
                operations.append(('af', i, face))
        for i, vertex in enumerate(self.vertices):
            if type(vertex) != NoneType:
                operations.append(('av', i, vertex))

        # 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.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)
    args = parser.parse_args()

    main(args)