utils.py

import numpy as np
import os

# color constants for polyscope consistent pics
PS_COLOR_SURFACE_MESH = (74/255, 161/255, 199/255)
PS_COLOR_CURVE_NETWORK = (255/255, 157/255, 0)

def read_planespec_file(fname):
    with open(fname, "r") as specfile:
        lines = specfile.read().splitlines()
    scan_in_direction = {"x":0, "y": 1, "z": 2}[lines[0]]
    n_planes = int(lines[1])
    min_coord = float(lines[2])
    max_coord = float(lines[3])
    return scan_in_direction, n_planes, min_coord, max_coord

def pad(arr: np.ndarray, pad_to_size: int, dim: int = 0, pad_value=0):
    pad_widths = [(0,0) for _ in range(len(arr.shape))]
    assert pad_to_size >= arr.shape[dim]
    pad_widths[dim] = (0, pad_to_size - arr.shape[dim])
    return np.pad(arr, pad_width=pad_widths, constant_values=pad_value)

def vstack_with_padding(arrays, pad_value=0):
    """ numpy.vstack but with automatic padding on the 2nd dimension
    (the 1st dimension is the stacking dimension)
    """
    max_dim1_size = max(map(lambda arr: arr.shape[1], arrays))
    arrays = list(map(lambda arr: pad(arr, max_dim1_size, dim=1, pad_value=pad_value), arrays))
    return np.vstack(arrays)

def write_oriented_pcloud_to_plyfile(fname:str, 
    orig_mesh_name:str, points: np.array, normals: np.array):

    assert points.shape == normals.shape # (n_points, 3)
    ply_lines = \
        [ "ply"
        , "format ascii 1.0"
        , f"comment Generated by slicing {orig_mesh_name} with some planes and taking its normals"
        , f"element vertex {len(points)}"
        , "property float x"
        , "property float y"
        , "property float z"
        , "property float nx"
        , "property float ny"
        , "property float nz"
        , "element face 0"
        , "property list uchar int vertex_indices"
        , "end_header"
        ]
    ply_lines = [l + '\n' for l in ply_lines]
    # now write the vertices and their normals
    ply_lines.extend([f"{p[0]} {p[1]} {p[2]} {n[0]} {n[1]} {n[2]}\n"
        for p, n in zip(points, normals)])
    
    with open(fname, "wt") as f:
        f.writelines(ply_lines)


# dummy state holder
class StatefulValue:
    def __init__(self, value):
        """ stateful holder of some values. Also includes a dict to store
        extra values if you wish, other than a 'primary value'. """
        self.value = value
        self.storage = dict()

    def __call__(self):
        return self.value

    def put(self, new_value):
        self.value = new_value

    def modify(self, func):
        self.value = func(self.value)
    
    def store(self, store_key, store_value):
        self.storage[store_key] = store_value
    
    def retrieve(self, store_key):
        return self.storage.get(store_key)


class PolygonSoup():
    """
    We define a triangular polygon soup as a collection of spatial points (vertices),
    and their connectivity information.
    The fields are:
        vertices: [N, 3]
        indices: [M, 3] where each row is the indices of the 3 vertices that make up a face

    This can be read from and written to in various file formats e.g obj, stl
    """
    def __init__(self, vertices, indices):
        vertices = np.asarray(vertices, dtype=np.float32)
        indices = np.asarray(indices, dtype=np.uint32)
        assert vertices.shape[1] == 3, "`vertices` must be an Nx3 array"
        assert indices.shape[1] == 3, "`faces` must be an Mx3 array"
        self.vertices = vertices
        self.indices = indices
    

    @classmethod
    def from_obj(cls, fname):
        """
        # An obj file looks like:
        v 0.123 0.234 0.345
        vn 0.707 0.000 0.707
        f 1 2 3
        # each line could be vertex_index/texture_index. Only need vertex_index
        f 3/1 4/2 5/3

        We can recompute normals "vn" ourselves
        """
        if os.path.splitext(fname)[1] == '.off':
            return PolygonSoup.from_off(fname)

        vertices = []
        indices = []
        with open(fname, "r") as f_handle:
            for line in f_handle:
                line = line.strip()
                tokens = line.split(" ")
                identifier = tokens[0]

                if identifier == "v":
                    vertices.append(
                        [float(tokens[1]), float(tokens[2]), float(tokens[3])]
                    )
                elif identifier == "f":
                    assert len(tokens) == 4,\
                        f"only triangle meshes are supported, got face index {line}"

                    face_indices = []
                    for i in range(3):
                        inx = tokens[1 + i].split("/")[0]  # throw away texture index, etc
                        inx = int(inx)
                        # NOTE obj index is 1-based
                        # theoretically negatives are allowed in the spec but
                        # that's not implemented here
                        assert (inx > 0), "index should be positive"
                        face_indices.append(inx - 1)
                    indices.append(face_indices)

        return cls(vertices, indices)