Script to merge triangular meshes from the surface of papyri sheets (Vesuvius Challenge)

mesh_merger.py

import open3d as o3d
import numpy as np
import igl
from PIL import Image
from scipy.spatial import Delaunay

Image.MAX_IMAGE_PIXELS = None

def find_affine_2dtransform(src_pts, dst_pts):
    """
    Calculate the affine transformation matrix from source points to destination points.
    
    Parameters:
    - src_pts (np.array): Nx2 array of source points.
    - dst_pts (np.array): Nx2 array of destination points.
    
    Returns:
    - np.array: The 3x3 affine transformation matrix.
    
    Raises:
    - ValueError: If there are fewer than three points or if any set of three points is collinear.
    """
    if src_pts.shape[0] < 3 or src_pts.shape[1] != 2:
        raise ValueError("Source points array must be at least 3x2.")
    if dst_pts.shape[0] < 3 or dst_pts.shape[1] != 2:
        raise ValueError("Destination points array must be at least 3x2.")
    
    # Function to augment points matrix
    def augment_points_matrix(pts):
        ones = np.ones((pts.shape[0], 1))
        return np.hstack([pts, ones])
    
    src_pts = src_pts + np.random.normal(0, 0.001, size=(src_pts.shape[0], 2))
    src_pts_aug = augment_points_matrix(src_pts)
    dst_pts = dst_pts + np.random.normal(0, 0.001, size=(dst_pts.shape[0], 2))
    dst_pts_aug = augment_points_matrix(dst_pts)

    # Solving the equation src_pts_aug * T = dst_pts_aug
    T, residuals, rank, s = np.linalg.lstsq(src_pts_aug, dst_pts_aug, rcond=None)

    # Form the full 3x3 affine transformation matrix
    T_aug = np.vstack([T.T, [0, 0, 1]])
    
    return T_aug

def find_affine_3dtransform(src_pts, dst_pts):
    
    # Function to augment points matrix
    def augment_points_matrix(pts):
        ones = np.ones((pts.shape[0], 1))
        return np.hstack([pts, ones])
    
    src_pts = src_pts + np.random.normal(0, 0.001, size=(src_pts.shape[0], 3))
    src_pts_aug = augment_points_matrix(src_pts)
    dst_pts = dst_pts + np.random.normal(0, 0.001, size=(dst_pts.shape[0], 3))
    dst_pts_aug = augment_points_matrix(dst_pts)

    # Solving the equation src_pts_aug * T = dst_pts_aug
    T, residuals, rank, s = np.linalg.lstsq(src_pts_aug, dst_pts_aug, rcond=None)

    T_aug = np.vstack([T.T, [0, 0, 0, 1]])
    
    return T_aug

def min_edge_length(mesh):
    # Get all edges from the triangles
    triangles = np.asarray(mesh.triangles)
    edges = np.vstack((triangles[:, [0, 1]], triangles[:, [1, 2]], triangles[:, [2, 0]]))
    
    # Remove duplicate edges
    edges = np.unique(np.sort(edges, axis=1), axis=0)
    
    # Calculate lengths of each edge
    vertices = np.asarray(mesh.vertices)
    edge_lengths = np.linalg.norm(vertices[edges[:, 0]] - vertices[edges[:, 1]], axis=1)
    
    return np.min(edge_lengths)

def find_closest_vertices(mesh1, mesh2, distance, multip=1, boundary=True):
    # Convert Open3D mesh to numpy arrays
    v1 = np.asarray(mesh1.vertices)
    f1 = np.asarray(mesh1.triangles)
    
    pcd1 = o3d.geometry.PointCloud()
    # Use libigl to find boundary vertices
    if boundary:
        print(f"Bondary {boundary}")
        boundary_vertices = igl.boundary_loop(f1)
        
        # Create a point cloud for boundary vertices of mesh1
        pcd1.points = o3d.utility.Vector3dVector(v1[boundary_vertices])  # Only include boundary vertices
    else:
        pcd1.points = o3d.utility.Vector3dVector(v1)
    # Create a point cloud for all vertices of mesh2
    pcd2 = o3d.geometry.PointCloud()
    pcd2.points = o3d.utility.Vector3dVector(np.asarray(mesh2.vertices))
    
    # Build a KDTree for mesh2 vertices
    tree = o3d.geometry.KDTreeFlann(pcd2)
    
    # Find closest vertices
    close_pairs = []
    used_indices = set()  # Track which indices in mesh2 are already paired
    
    for i, point in enumerate(pcd1.points):
        _, idxs, dists = tree.search_hybrid_vector_3d(point, multip*distance, 100)  # Search within a specified distance and limit
        closest_idx = None
        min_dist = float('inf')

        # Find the closest available index in mesh2 that hasn't been used yet
        for idx, dist in zip(idxs, dists):
            if idx not in used_indices and dist < min_dist:
                min_dist = dist
                closest_idx = idx
        
        # If a closest index was found and it's not used, add it to the pairs
        if closest_idx is not None:
            real_idx1 = boundary_vertices[i] if boundary else i  # The actual index in mesh1
            close_pairs.append((real_idx1, closest_idx))  # Store pair with original mesh1 index
            used_indices.add(closest_idx)  # Mark this index as used

    return close_pairs


def filter_triangles(mesh, vertex_indices):
    triangles_to_exclude = set()
    vertex_indices_set = set(vertex_indices)
    for triangle in mesh.triangles:
        if all(vertex in vertex_indices_set for vertex in triangle):
            triangles_to_exclude.add(tuple(triangle))
    return [triangle for triangle in mesh.triangles if tuple(triangle) not in triangles_to_exclude]

def extract_uvs(mesh1, mesh2):
    uv1 = np.asarray(mesh1.triangle_uvs)
    v1 = np.asarray(mesh1.vertices)
    t1 = np.asarray(mesh1.triangles)
    uv2 = np.asarray(mesh2.triangle_uvs)
    v2 = np.asarray(mesh2.vertices)
    t2 = np.asarray(mesh2.triangles)


    uuvv1 = np.zeros((v1.shape[0], 2), dtype=np.float64)
    uvs1 = uv1.reshape((t1.shape[0], t1.shape[1], 2))

    for t in range(t1.shape[0]):
        for v in range(t1.shape[1]):
            uuvv1[t1[t,v]] = uvs1[t,v]

    uuvv2 = np.zeros((v2.shape[0], 2), dtype=np.float64)
    uvs2 = uv2.reshape((t2.shape[0], t2.shape[1], 2))

    for t in range(t2.shape[0]):
        for v in range(t2.shape[1]):
            uuvv2[t2[t,v]] = uvs2[t,v]

    return uuvv1, uuvv2

def normalize_uv(uv):
    min_uv = uv.min(axis=0)
    max_uv = uv.max(axis=0)
    norm_uv = (uv - min_uv) / np.where(max_uv - min_uv == 0, 1, max_uv - min_uv)
    return norm_uv

def extract_dimensions(mask1_path, mask2_path):
    with Image.open(mask1_path) as img:
        size1 = img.size

    with Image.open(mask2_path) as img:
        size2 = img.size

    return size1, size2
    
def shift_mesh(mesh1, mesh2, close_pairs):
    two_vertices = np.asarray(mesh2.vertices)
    one_vertices = np.asarray(mesh1.vertices)
    src_points = np.array([two_vertices[j] for _, j in close_pairs])
    dst_points = np.array([one_vertices[i] for i, _ in close_pairs])
    affine = find_affine_3dtransform(src_points,dst_points)
    print("Affine 3D Transformation Matrix")
    print(f"{affine}")
    print("Applying affine transform")
    transformed_pts = np.einsum('ij,nj->ni', affine, np.hstack((two_vertices,np.ones(two_vertices.shape[0]).reshape(-1,1))))[:,:3]
    mesh2.vertices = o3d.utility.Vector3dVector(transformed_pts)
    return mesh2


def merge_meshes(mesh1, mesh2):
    max_index1 = len(mesh1.vertices)
    adjusted_vertices = np.asarray(mesh2.vertices)
    adjusted_triangles = np.asarray(mesh2.triangles) + max_index1
    one_vertices = np.asarray(mesh1.vertices)

    new_vertices = np.vstack((one_vertices, adjusted_vertices))
    new_triangles = np.vstack((np.asarray(mesh1.triangles), adjusted_triangles))

    merged_mesh = o3d.geometry.TriangleMesh()
    merged_mesh.vertices = o3d.utility.Vector3dVector(new_vertices)
    merged_mesh.triangles = o3d.utility.Vector3iVector(new_triangles)
    merged_mesh.compute_vertex_normals()

    return merged_mesh

def find_holes(mesh):
    # Get the vertex and face data from the mesh
    faces = np.asarray(mesh.triangles)

    # Use igl to compute the boundary edges
    boundary_edges = igl.boundary_facets(faces)

    # Organize boundary edges into loops
    loops = []  # This should hold lists of edges that form loops
    visited = set()

    for e in boundary_edges:
        if tuple(e) in visited or tuple(e[::-1]) in visited:
            continue

        # Walk through to find the complete loop
        loop = [tuple(e)]
        visited.add(tuple(e))

        while True:
            found_next = False
            for edge in boundary_edges:
                if tuple(edge) in visited or tuple(edge[::-1]) in visited:
                    continue

                if edge[0] == loop[-1][1]:
                    loop.append(tuple(edge))
                    visited.add(tuple(edge))
                    found_next = True
                    break

                elif edge[1] == loop[-1][1]:
                    loop.append(tuple(edge[::-1]))
                    visited.add(tuple(edge[::-1]))
                    found_next = True
                    break

            if not found_next or loop[0][0] == loop[-1][1]:
                break

        # Append to loops if it's a valid hole (loop with at least three edges)
        if len(loop) >= 3:
            loops.append(loop)

    # To distinguish between holes and the outer boundary, sort by loop length
    loops = sorted(loops, key=len)

    # Assume the largest loop (last one) is the outer boundary and ignore it
    return loops[:-1]

def fill_triangular_holes(mesh, loops):
    # Convert the mesh to vertices and faces
    faces = np.asarray(mesh.triangles).tolist()

    for loop in loops:
        if len(loop) == 3:
            # Get vertices from the loop edges
            v1 = loop[0][0]
            v2 = loop[0][1]
            v3 = loop[1][1]  # since loop[1] = (loop[0][1], loop[1][1])

            # Add a new face to the mesh
            faces.append([v1, v2, v3])

    # Convert faces back to numpy array
    mesh.triangles = o3d.utility.Vector3iVector(np.array(faces))
    return mesh

def orient_uvs(vertices):
        # Rotate vertices and calculate the needed area
        vertices[:, 0] = 1.0 - vertices[:, 0]
        u_range = np.max(vertices[:, 0]) - np.min(vertices[:, 0])
        v_range = np.max(vertices[:, 1]) - np.min(vertices[:, 1])
        u_longer_v = u_range > v_range
        u_return = vertices[:, 0]
        v_return = vertices[:, 1]
        area_return = u_range * v_range
        for angle in range(-70, 70, 5):
            u_prime = vertices[:, 0] * np.cos(np.deg2rad(angle)) - vertices[:, 1] * np.sin(np.deg2rad(angle))
            v_prime = vertices[:, 0] * np.sin(np.deg2rad(angle)) + vertices[:, 1] * np.cos(np.deg2rad(angle))
            u_prime_range = np.max(u_prime) - np.min(u_prime)
            v_prime_range = np.max(v_prime) - np.min(v_prime)
            if u_prime_range < v_prime_range and u_longer_v:
                continue
            elif u_prime_range > v_prime_range and not u_longer_v:
                continue
            area = u_prime_range * v_prime_range
            if area < area_return:
                u_return = u_prime
                v_return = v_prime
                area_return = area

        return np.stack((u_return, v_return), axis=-1) 


if __name__ == "__main__":
    import argparse
    from copy import deepcopy
    from matplotlib import pyplot as plt
    # Create the parser
    parser = argparse.ArgumentParser(description="Merge two partially overlapping triangular meshes.")

    # Add positional arguments
    parser.add_argument(
        "mesh1_path",
        type=str,
        help="Path to the first mesh file"
    )

    parser.add_argument(
        "mesh2_path",
        type=str,
        help="Path to the second mesh file"
    )

    # Add optional argument
    parser.add_argument(
        "-o", "--output",
        type=str,
        help="Path where the output should be stored",
        default="merged_mesh.obj"
    )

    # Parse the arguments
    args = parser.parse_args()

    # Use the arguments
    print(f"Mesh 1 file path: {args.mesh1_path}")
    print(f"Mesh 2 file path: {args.mesh2_path}")
    if args.output:
        print(f"Output will be stored in: {args.output}")
    else:
        print("No output path provided, default will be used.")
    
    print("Loading the mesh 1")
    mesh1 = o3d.io.read_triangle_mesh(args.mesh1_path)
    print("Loading the mesh 2")
    mesh2 = o3d.io.read_triangle_mesh(args.mesh2_path)

    #mesh0 = mesh1.subdivide_midpoint(number_of_iterations=2)
    #mesh3 = mesh2.subdivide_midpoint(number_of_iterations=2)
    #print("Computing min edge distance in meshes")
    max_dist = max(min_edge_length(mesh1),min_edge_length(mesh2))
    print(f"Min computed distance: {max_dist:.3f}")

    print("Finding pairs of close vertices between meshes")
    close_pairs1 = find_closest_vertices(deepcopy(mesh1), deepcopy(mesh2), max_dist, 2, boundary=False)
    close_pairs2_temp = find_closest_vertices(deepcopy(mesh2), deepcopy(mesh1), max_dist, 2, boundary=False)
    close_pairs2 = set()
    for i,j in close_pairs2_temp:
        close_pairs2.add((j,i))
    print(f"Found {len(close_pairs1)} pairs")
    print(f"Found {len(close_pairs2)} pairs")
    close_pairs = close_pairs1 + list(close_pairs2)
    close_pairs = list(set(close_pairs))
    print(f"Found {len(close_pairs)} pairs")

    print("Extracting UV parametrizations from meshes")
    try:
        uv1, uv2 = extract_uvs(mesh1, mesh2)
        dim1, dim2 = extract_dimensions(args.mesh1_path[:-3]+"png", args.mesh2_path[:-3]+"png")
        uv1[:,0] *= dim1[0]
        uv1[:,1] *= dim1[1]
        uv2[:,0] *= dim2[0]
        uv2[:,1] *= dim2[1]

        print(dim1, dim2)
        print("Computing affine transform between uvs")
        src_points = np.array([uv2[j] for _, j in close_pairs])
        dst_points = np.array([uv1[i] for i, _ in close_pairs])
        affine = find_affine_2dtransform(src_points,dst_points)
        print("Affine Transformation Matrix")
        print(f"{affine}")
        print("Applying affine transform")
        transformed_uv = np.einsum('ij,nj->ni', affine, np.hstack((uv2,np.ones(uv2.shape[0]).reshape(-1,1))))[:,:2]
        merged_uv = np.vstack((uv1, transformed_uv))
        print("Normalizing UV")
        plt.figure()
        plt.scatter(uv1[:,0], uv1[:,1], color='blue', alpha=0.5, label='uv1')
        plt.scatter(transformed_uv[:,0], transformed_uv[:,1], color='red', alpha=0.5, label='uv2')
        plt.savefig("merged_uvs.png")
        min_x, min_y = np.min(merged_uv, axis=0)
        shifted_coords = merged_uv - np.array([min_x, min_y])
        max_x, max_y = np.max(shifted_coords, axis=0)
        # Create a white image of the determined size
        image_size = (int(round(max_y)) + 1, int(round(max_x)) + 1)
        white_image = np.ones((image_size[0], image_size[1]), dtype=np.uint16) * 65535

        merged_uv = orient_uvs(merged_uv)
        print("Performing Delaunay", end="\n")
        tri = Delaunay(merged_uv)
        print("Done")
        delaunay_faces = tri.simplices

        # Save the grayscale image
        Image.fromarray(white_image, mode='L').save("merged_mesh.png")
        #merged_uv = normalize_uv(merged_uv)
        del uv1, uv2, transformed_uv, affine
    except:
        merged_uv = None
        delaunay_faces = None
        print("No parametrization found")
    print("Merging meshes, main step")
    merged_mesh = merge_meshes(deepcopy(mesh1), deepcopy(mesh2))
    print(f"Mesh 1 {mesh1}")
    print(f"Mesh 2 {mesh2}")
    print(f"Merged {merged_mesh}")

    del mesh1, mesh2

    print("Intrinsic Delaunay")
    V = np.asarray(merged_mesh.vertices)
    F = np.asarray(merged_mesh.triangles)
    if delaunay_faces is not None:
        _, _, F = igl.intrinsic_delaunay_cotmatrix(V, delaunay_faces)
    else:
        _, _, F = igl.intrinsic_delaunay_cotmatrix(V, F)
    print("Intrinsic Delaunay Computed")
    #m = igl.massmatrix(V, F, igl.MASSMATRIX_TYPE_BARYCENTRIC)
    #s = m - L
    print("Smoothing")
    #V = spsolve(s, m @ V)
    #merged_mesh.vertices = o3d.utility.Vector3dVector(V)
    merged_mesh.triangles = o3d.utility.Vector3iVector(F)
    merged_mesh = merged_mesh.filter_smooth_laplacian(number_of_iterations=2)
    merged_mesh.compute_vertex_normals()
    
    print(f"Removing non manifold edges")
    merged_mesh.remove_non_manifold_edges()
    merged_mesh.remove_duplicated_triangles()
    merged_mesh.remove_degenerate_triangles()

    F = np.asarray(merged_mesh.triangles)
    if merged_uv is not None:
        filtered_uvs = np.zeros((F.shape[0], F.shape[1], 2))
        for i in range(filtered_uvs.shape[0]):
            for j in range(filtered_uvs.shape[1]):
                filtered_uvs[i,j] = merged_uv[F[i,j]]
        filtered_uvs = filtered_uvs.reshape(-1,2)
        merged_mesh.triangle_uvs = o3d.utility.Vector2dVector(filtered_uvs)

    print(f"Removing unreferenced vertices")
    merged_mesh = merged_mesh.remove_unreferenced_vertices()
    merged_mesh.compute_vertex_normals()
    print(f"Filtered {merged_mesh}")
    

    vmanifold = merged_mesh.is_vertex_manifold()
    emanifold = merged_mesh.is_edge_manifold()
    print(f"Is vertex manifold: {vmanifold}")
    print(f"Is edge manifold: {emanifold}")

    print("Cluster connected triangles")
    with o3d.utility.VerbosityContextManager(
            o3d.utility.VerbosityLevel.Debug) as cm:
        triangle_clusters, cluster_n_triangles, cluster_area = (
            merged_mesh.cluster_connected_triangles())
    

    print(f"Coherent reflattening")
    V = np.asarray(merged_mesh.vertices)
    F = np.asarray(merged_mesh.triangles)
    UV = np.asarray(merged_mesh.triangle_uvs)  # Ensure your mesh has UV coordinates
    
    uv = np.zeros((V.shape[0], 2), dtype=np.float64)
    uvs = UV.reshape((F.shape[0], F.shape[1], 2))

    for t in range(F.shape[0]):
        for v in range(F.shape[1]):
            uv[F[t,v]] = uvs[t,v]
    bnd = igl.boundary_loop(F)
    bnd_uv = np.zeros((bnd.shape[0], 2), dtype=np.float64)
    for i in range(bnd.shape[0]):
        bnd_uv[i] = uv[bnd[i]]
    
    slim = igl.SLIM(V, F, v_init=uv, b=bnd, bc=bnd_uv, energy_type=igl.SLIM_ENERGY_TYPE_SYMMETRIC_DIRICHLET, soft_penalty=0)
    threshold = 1e-4
    converged = False
    iteration = 0
    while (not converged) and (iteration < 20000):
        print(iteration)
        temp_energy = slim.energy()
        slim.solve(1)
        new_energy = slim.energy()
        iteration += 1
        print(f"{temp_energy:.3f} {new_energy:.3f}")
        if new_energy >= float("inf") or new_energy == float("nan") or np.isnan(new_energy) or np.isinf(new_energy):
            continue
        elif new_energy < temp_energy:
            if abs(new_energy - temp_energy) < threshold:
                converged = True
            else:
                converged = False
        else:
            converged = True
    
    for t in range(F.shape[0]):
        for v in range(F.shape[1]):
            uvs[t,v] = slim.vertices()[F[t,v]]

    UV = uvs.reshape((-1,2))
    
    UV = orient_uvs(UV)
    UV = normalize_uv(UV)
    merged_mesh.triangle_uvs = o3d.utility.Vector2dVector(UV)





    print(f"Saving mesh")
    o3d.io.write_triangle_mesh(args.output, merged_mesh)
    print("Done")

Example

merge segments 3334 and 5109,

requirement: 3334.png and 5109.png files of the appropriate sizes in the same folder of the relative .obj. For these images, you can use the _mask.png output generated by the Volume Cartographer software. Make sure to double-check the paths.

example usage:

python mesh_merger.py "D:/vesuvius/Scroll1.volpkg/working/merging/3334/3334.obj" "D:/vesuvius/Scroll1.volpkg/working/merging/5109/5109.obj"