utils/mesh.py

import argparse
from pathlib import Path
import os 
from tqdm import tqdm
import potpourri3d as pp3d
import scipy.io as sio
import numpy as np
import re
import sys
from utils.surfaces import Surface, centroid
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from shape_data import get_data_dirs
import vtk
from vtk.util import numpy_support
# List of file extensions to consider as "mesh" files.
# Kudos to chatgpt!
# Add or remove extensions here as needed.
MESH_EXTENSIONS = {".ply", ".obj", ".off", ".stl", ".fbx", ".gltf", ".glb"}

def sorted_alphanum(file_list_ordered):
    def convert(text):
        return int(text) if text.isdigit() else text
    def alphanum_key(key):
        return [convert(c) for c in re.split('([0-9]+)', str(key)) if len(c) > 0]
    return sorted(file_list_ordered, key=alphanum_key)


def list_files(folder_path, name_filter, alphanum_sort=False):
    file_list = [p.name for p in list(Path(folder_path).glob(name_filter))]
    if alphanum_sort:
        return sorted_alphanum(file_list)
    else:
        return sorted(file_list)

def find_mesh_files(directory, extensions, alphanum_sort=False): # Path, extensions: set=MESH_EXTENSIONS, alphanum_sort=False) -> list[Path]:
    """
    Recursively find all files in 'directory' whose suffix (lowercased) is in 'extensions'.
    Returns a list of Path objects.
    """
    matches = []
    for path in directory.rglob("*"):
        if path.is_file() and path.suffix.lower() in extensions:
            matches.append(path)
    if alphanum_sort:
        return sorted_alphanum(matches)
    else:
        return sorted(matches)

def save_ply(file_name, V, F, Rho=None, color=None):
    """Save mesh information either as an ASCII ply file.
    https://github.com/emmanuel-hartman/BaRe-ESA/blob/main/utils/input_output.py
    Input:
        - file_name: specified path for saving mesh [string]
        - V: vertices of the triangulated surface [nVx3 numpy ndarray]
        - F: faces of the triangulated surface [nFx3 numpy ndarray]
        - Rho: weights defined on the vertices of the triangulated surface [nVx1 numpy ndarray, default=None]
        - color: colormap [nVx3 numpy ndarray of RGB triples]

    Output:
        - file_name.mat or file_name.ply file containing mesh information
    """

    # Save as .ply file
    nV = V.shape[0]
    nF = F.shape[0]   
    if not ".ply" in file_name:
        file_name += ".ply"
    file = open(file_name, "w")
    lines = ("ply","\n","format ascii 1.0","\n", "element vertex {}".format(nV),"\n","property float x","\n","property float y","\n","property float z","\n")
    
    if color is not None:
        lines += ("property uchar red","\n","property uchar green","\n","property uchar blue","\n")
        if Rho is not None:
            lines += ("property uchar alpha","\n")
    
    lines += ("element face {}".format(nF),"\n","property list uchar int vertex_index","\n","end_header","\n")

    file.writelines(lines)
    lines = []
    for i in range(0,nV):
        for j in range(0,3):
            lines.append(str(V[i][j]))
            lines.append(" ")
        if color is not None:
            for j in range(0,3):
                lines.append(str(color[i][j]))
                lines.append(" ")
            if Rho is not None:
                lines.append(str(Rho[i]))
                lines.append(" ")
                    
        lines.append("\n")
    for i in range(0,nF):
        l = len(F[i,:])
        lines.append(str(l))
        lines.append(" ")

        for j in range(0,l):
            lines.append(str(F[i,j]))
            lines.append(" ")
        lines.append("\n")

    file.writelines(lines)
    file.close()



def load_mesh(filepath, scale=True):
    V, F = pp3d.read_mesh(filepath)
    down_V, down_F = decimation(V, F)
    down_surf = Surface(FV=[down_F,down_V])
    down_file = os.path.join("tmp/plys", os.path.basename(filepath)[:-4])
    save_ply(down_file, down_surf.vertices, down_surf.faces)
    surf = Surface(FV=[F,V])
    center, area = centroid(surf)
    if scale:
        new_vertices = (surf.vertices - center)/area
        new_vertices_down = (down_surf.vertices - center) / area
    else:
        new_vertices = (surf.vertices - center)
        new_vertices_down = (down_surf.vertices - center) / area
    surf.updateVertices(np.array(new_vertices, dtype=np.float32))
    down_surf.updateVertices(np.array(new_vertices_down, dtype=np.float32))
    return surf, down_surf, down_file+".ply"


def add_vectors(vertices: np.ndarray, faces: np.ndarray, vtk_vertices, vtk_faces):
    assert len(vertices.shape) == 2
    assert len(faces.shape) == 2
    assert vertices.shape[1] == 3
    assert faces.shape[1] == 3

    # Add points
    [num_vertices, _] = vertices.shape
    for vertex_idx in range(num_vertices):
        vtk_vertices.InsertNextPoint(vertices[vertex_idx, 0], vertices[vertex_idx, 1], vertices[vertex_idx, 2])
    [num_faces, _] = faces.shape
    for face_idx in range(num_faces):
        vtk_faces.InsertNextCell(3)
        for corner_idx in range(3):
            vtk_faces.InsertCellPoint(faces[face_idx, corner_idx])
    # Allocate additional memory
    vtk_vertices.Resize(num_vertices)
    # self._faces.Resize(self.num_faces + num_faces*3)

    vtk_vertices.Modified()
    vtk_faces.Modified()

def decimation(verts, faces, target_n=1500):
    vtk_vertices = vtk.vtkPoints()

    # VTK polygone(surface) representation
    vtk_faces = vtk.vtkCellArray()
    surface_data = vtk.vtkPolyData()
    surface_data.SetPoints(vtk_vertices)
    surface_data.SetPolys(vtk_faces)
    add_vectors(verts, faces, vtk_vertices, vtk_faces)


    ratio = target_n /verts.shape[0]
    print(ratio)
    print("Before decimation\n"
          "-----------------\n"
          "There are " + str(surface_data.GetNumberOfPoints()) + "points.\n"
          "There are " + str(surface_data.GetNumberOfPolys()) + "polygons.\n")

    decimate = vtk.vtkQuadricDecimation()
    decimate.SetInputData(surface_data)
    decimate.SetTargetReduction(1-ratio)
    decimate.Update()

    decimatedPoly = vtk.vtkPolyData()
    decimatedPoly.ShallowCopy(decimate.GetOutput())

    print("After decimation \n"
          "-----------------\n"
          "There are " + str(decimatedPoly.GetNumberOfPoints()) + "points.\n"
          "There are " + str(decimatedPoly.GetNumberOfPolys()) + "polygons.\n")
    points = decimatedPoly.GetPoints().GetData()
    vertices = numpy_support.vtk_to_numpy(points)  # shape: (n_points, 3)

    # --- Faces (connectivity) ---
    faces = decimatedPoly.GetPolys().GetData()
    faces_np = numpy_support.vtk_to_numpy(faces)
    faces_np = faces_np.reshape(-1, 4)[:, 1:]
    return vertices, faces_np


def mesh_geod_matrix(vertices, faces, do_tqdm=False, verbose=False):
    if verbose:
        print("Setting Geodesic matrix bw vertices")
    n_vertices = vertices.shape[0]
    distmat = np.zeros((n_vertices, n_vertices))
    solver = pp3d.MeshHeatMethodDistanceSolver(vertices, faces)
    if do_tqdm:
        iterable = tqdm(range(n_vertices))
    else:
        iterable = range(n_vertices)
    for vertind in iterable:
        distmat[vertind] = np.maximum(solver.compute_distance(vertind), 0)
    geod_mat = distmat
    return geod_mat


def prepare_geod_mats(shapes_folder, out, basename=None):
    if basename is None:
        basename = os.path.basename(os.path.dirname(shapes_folder)) #+ "_" + os.path.basename(shapes_folder)
    case = basename
    case_folder_out = os.path.join(out, case)
    os.makedirs(case_folder_out, exist_ok=True)
    all_shapes = [f for f in os.listdir(shapes_folder) if (".ply" in f) or (".off" in f) or ('.obj' in f)]
    for shape in tqdm(all_shapes, "Processing " + os.path.basename(shapes_folder)):
        vertices, faces = pp3d.read_mesh(os.path.join(shapes_folder, shape))
        areas = pp3d.face_areas(vertices, faces)
        mat = mesh_geod_matrix(vertices, faces, verbose=False)
        dict_save = {
            'geod_dist': mat,
            'areas_f': areas
        }
        sio.savemat(os.path.join(case_folder_out, shape[:-4]+'.mat'), dict_save)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="What to do.")
    parser.add_argument('--make_geods', required=True, type=int, default=0, help='launch computation of geod matrices')
    parser.add_argument('--data', required=False, type=str, default=None)
    parser.add_argument('--datadir', type=str, default="data", help='path where datasets are store')
    parser.add_argument('--basename', required=False, type=str, default=None)
    args = parser.parse_args()
    if args.make_geods:
        # from config import get_geod_path, get_dataset_path, get_template_path
        # output = get_geod_path()
        output = os.path.join(args.datadir, "geomats")
        if args.data == "humans":
            all_folders = [get_data_dirs(args.datadir, "faust", 'test')[0], get_data_dirs(args.datadir, "scape", 'test')[0], get_data_dirs(args.datadir, "shrec19", 'test')[0]]
            for folder in all_folders:
                prepare_geod_mats(folder, output)
        elif args.data == "dt4d":
            data_dir, _, corr_dir = get_data_dirs(args.datadir, args.data, 'test')
            all_folders = sorted([f for f in os.listdir(data_dir) if "cross" not in f])
            for folder in all_folders:
                prepare_geod_mats(os.path.join(data_dir, folder), os.path.join(output, "DT4D"), basename=folder)
        elif args.data is not None:
            data_dir, _, corr_dir = get_data_dirs(args.datadir, args.data, 'test')
            prepare_geod_mats(data_dir, output, args.basename)
    # parser = argparse.ArgumentParser(description="Find all mesh files in a folder and list their paths.")
    # parser.add_argument("--folder", type=Path, help="Path to the folder to search (will search recursively).")
    # args = parser.parse_args()

    # search_folder = args.folder
    # if not search_folder.is_dir():
    #     print(f"Error: '{search_folder}' is not a valid directory.")
    # else:
    #     # Find all matching files
    #     mesh_files = find_mesh_files(search_folder)

    #     # Sort the results for consistency
    #     mesh_files.sort()
    #     for p in mesh_files:
    #         print(p.resolve())