fLOD_generation.py

import trimesh
import matplotlib.pyplot as plt
import napari
import numpy as np
import os
import triangle
import argparse

from collections import deque, Counter
from scipy.ndimage import label, binary_fill_holes, distance_transform_edt
from scipy.spatial import cKDTree
from shapely.geometry import MultiPoint, Point
from trimesh.util import concatenate


def visualize_results(binary_volume: np.ndarray,
                      skeleton: np.ndarray) -> None:
    """
    Visualize the original volume and its skeleton using napari.

    Parameters
    ----------
    binary_volume : np.ndarray
        Original binary volume
    skeleton : np.ndarray
        Skeleton volume
    """
    viewer = napari.Viewer(ndisplay=3)

    # Add layers
    viewer.add_image(binary_volume, name='Original Volume', opacity=0.3)
    viewer.add_image(skeleton, name='Skeleton', colormap='red', blending='additive')

    # Maximize window
    viewer.window._qt_window.showMaximized()

    napari.run()


def projection_fits_circle(mesh_part, radius):
    """
    Check if projections of mesh onto XY, XZ, YZ fit inside
    a circle of given radius centered at the projection's center.
    """
    projections = {
        'XY': mesh_part.vertices[:, [0, 1]],
        'XZ': mesh_part.vertices[:, [0, 2]],
        'YZ': mesh_part.vertices[:, [1, 2]],
    }

    for plane, verts_2d in projections.items():
        hull = MultiPoint(verts_2d).convex_hull
        centroid = np.array([hull.centroid.x, hull.centroid.y])

        circle = Point(centroid).buffer(radius)

        if not circle.within(hull):
            print(f"{plane} projection does not fit circle")
        else:
            print(f"{plane} projection fits circle")
            return True

    return False


def visualise_parts(mesh_parts):
    # Visualize parts with random colors
    scene = trimesh.Scene()
    for part in mesh_parts:
        # Make a copy so each part has its own visuals
        part_copy = part.copy()

        # Random color
        color = np.random.randint(0, 256, size=3, dtype=np.uint8)
        face_colors = np.tile(np.append(color, 255), (len(part_copy.faces), 1))  # RGBA

        part_copy.visual = trimesh.visual.ColorVisuals(
            mesh=part_copy,
            face_colors=face_colors
        )

        scene.add_geometry(part_copy)

    scene.show()


def bfs(slice_2d, start_i, start_j, visited):
    rows, cols = len(slice_2d), len(slice_2d[0])
    q = deque([(start_i, start_j)])
    region = []
    visited[start_i][start_j] = True

    while q:
        i, j = q.popleft()
        region.append((i, j))

        # 4-connected neighbors
        for di, dj in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
            ni, nj = i + di, j + dj
            if 0 <= ni < rows and 0 <= nj < cols:
                if not visited[ni][nj] and slice_2d[ni][nj] == 0:
                    visited[ni][nj] = True
                    q.append((ni, nj))
    return region


def find_zero_regions_2d(slice_2d):
    rows, cols = len(slice_2d), len(slice_2d[0])
    visited = [[False] * cols for _ in range(rows)]
    regions = []

    for i in range(rows):
        for j in range(cols):
            if slice_2d[i][j] == 0 and not visited[i][j]:
                regions.append(bfs(slice_2d, i, j, visited))

    return regions


def fill_big_holes(matrix3d, threshold_hole):
    filled_matrix = [[row[:] for row in slice_2d] for slice_2d in matrix3d]  # copy
    for z, slice_2d in enumerate(matrix3d):
        regions = find_zero_regions_2d(slice_2d)
        for region in regions:
            avg_i = round(sum(i for i, j in region) / len(region))
            avg_j = round(sum(j for i, j in region) / len(region))
            n_i_min, n_i_max = max(0, int(avg_i - threshold_hole)), min(slice_2d.shape[0], int(avg_i + threshold_hole + 1))
            n_j_min, n_j_max = max(0, int(avg_j - threshold_hole)), min(slice_2d.shape[1], int(avg_j + threshold_hole + 1))
            neighbouring_cells = slice_2d[n_i_min:n_i_max, n_j_min:n_j_max]
            if np.sum(neighbouring_cells) != 0 and avg_i != 0 and avg_i != slice_2d.shape[0] - 1:
                # fill region with 1s
                for (i, j) in region:
                    filled_matrix[z][i][j] = 1
    return np.array(filled_matrix)


def rebuild_mesh_with_patches(original_mesh, patch_meshes, tol=1e-8):
    """
    Rebuilds a Trimesh from the original mesh and a list of patch meshes.
    
    Parameters:
    - original_mesh: trimesh.Trimesh
    - patch_meshes: list of trimesh.Trimesh
    - tol: tolerance for merging vertices (default 1e-8)
    
    Returns:
    - new_mesh: trimesh.Trimesh with patches merged
    """
    merged_vertices = original_mesh.vertices.copy()
    all_faces = original_mesh.faces.copy()

    for patch in patch_meshes:
        patch_indices = []

        # Merge patch vertices into merged_vertices
        for v in patch.vertices:
            # Find existing vertex within tolerance
            diff = np.linalg.norm(merged_vertices - v, axis=1)
            idx = np.where(diff < tol)[0]
            if len(idx) > 0:
                patch_indices.append(idx[0])
            else:
                merged_vertices = np.vstack([merged_vertices, v])
                patch_indices.append(len(merged_vertices) - 1)

        patch_indices = np.array(patch_indices)
        # Remap patch faces to merged vertex indices
        new_faces = patch_indices[patch.faces]
        all_faces = np.vstack([all_faces, new_faces])

    # Build the new mesh
    new_mesh = trimesh.Trimesh(vertices=merged_vertices, faces=all_faces)
    return new_mesh


if __name__ == '__main__':
    # Step 1: Load STL
    parser = argparse.ArgumentParser()
    parser.add_argument("--fLOD", type=int, default=2, help="furniture Level of detail")
    parser.add_argument("--path", default=os.path.join("meshes", "table.stl"),
                        help="Path to the STL file")

    args = parser.parse_args()

    fLOD = args.fLOD
    path = args.path
    # fLOD = 2
    # path = os.path.join("meshes", "table.stl")
    mesh = trimesh.load(path)
    print("loaded STL")

    if fLOD == 1:
        # Step 2: Compute the oriented bounding box
        obb = mesh.bounding_box_oriented

        # The OBB returned is itself a Trimesh object (a box aligned with principal axes)
        # but still in local coordinates. To get the box in world coordinates:
        obb_mesh = obb.to_mesh()

        # Save
        # obb_mesh.export("obb_mesh.stl")

        # Or show OBB
        scene = trimesh.Scene([obb_mesh])
        scene.show()

    else:
        # Step 2: Voxelize
        voxel_size = 0.01  # adjust for detail
        voxelized = mesh.voxelized(pitch=voxel_size)
        volume = voxelized.matrix.astype(np.uint8)
        print("Voxelisation complete!")

        # fill enclosed small holes before generating skeleton
        volume_filled = binary_fill_holes(volume).astype(np.uint8)

        # Compute distance transform
        distance = distance_transform_edt(volume_filled)

        # Initialize skeleton as empty array
        skeleton = np.zeros_like(volume_filled, dtype=np.uint8)

        # Sort voxel coordinates by decreasing distance
        coords = np.argwhere(volume_filled)
        distances = distance[coords[:, 0], coords[:, 1], coords[:, 2]]
        sorted_indices = np.argsort(-distances)
        sorted_coords = coords[sorted_indices]

        for x, y, z in sorted_coords:
            if volume_filled[x, y, z] == 0:
                continue
            # Search radius is the distance value at this voxel
            r = distance[x, y, z]
            # If this voxel is the local maximum, keep it as skeleton
            # Check if any neighboring voxel with higher distance than r
            x_min, x_max = max(0, x - int(r)), min(volume_filled.shape[0], x + int(r) + 1)
            y_min, y_max = max(0, y - int(r)), min(volume_filled.shape[1], y + int(r) + 1)
            z_min, z_max = max(0, z - int(r)), min(volume_filled.shape[2], z + int(r) + 1)

            neighbouring_distances = distance[x_min:x_max, y_min:y_max, z_min:z_max]
            if np.max(neighbouring_distances) <= r:
                skeleton[x, y, z] = 1  # keep voxel as skeleton

        # # Visualize skeleton
        # fig = plt.figure()
        # ax = fig.add_subplot(111, projection='3d')
        # ax.scatter(skeleton[:, 0], skeleton[:, 1], skeleton[:, 2], s=1)
        # plt.show()

        labeled_skeleton, num_labels = label(skeleton)
        print(f"Skeleton branches detected: {num_labels}")
        # identify and label the corner parts
        unique_values, counts = np.unique(labeled_skeleton, return_counts=True)
        for i, unique_value in enumerate(unique_values):
            if counts[i] < 10:  # 10 should be a user defined parameter
                labeled_skeleton[labeled_skeleton == unique_value] = 0

        # Extract skeleton coordinates
        skeleton_coords = np.argwhere(labeled_skeleton > 0)

        # unique_values, counts = np.unique(labeled_skeleton, return_counts=True)

        # Get skeleton coordinates + their branch labels
        branch_labels = labeled_skeleton[skeleton_coords[:, 0],
                                         skeleton_coords[:, 1],
                                         skeleton_coords[:, 2]]

        # Convert voxel coords to world coords (mesh space)
        origin = voxelized.bounds[0]
        skeleton_points_world = origin + skeleton_coords * voxel_size

        # ==== Step 5: Nearest-neighbor mapping (vertex → skeleton) ====
        tree = cKDTree(skeleton_points_world)
        # Compute face centroids: average of the three vertices per face
        face_vertices = mesh.vertices[mesh.faces]  # shape (n_faces, 3, 3)
        face_centroids = face_vertices.mean(axis=1)  # shape (n_faces, 3)

        # Query KDTree with centroids instead of vertices
        distances, nearest_idx = tree.query(face_centroids)

        # Assign labels to faces
        final_face_labels = branch_labels[nearest_idx]

        # ==== Step 6b: bounding box relabeling check ====
        # Compute bounding boxes of each part based on current labels
        part_bboxes = {}
        unique_labels = np.unique(final_face_labels)
        for lbl in unique_labels:
            if lbl == 0:
                continue
            mask = final_face_labels == lbl
            part_mesh = mesh.submesh([mask], append=True, repair=False)
            min_corner, max_corner = part_mesh.bounds
            part_bboxes[lbl] = (min_corner, max_corner)

        # Check each face centroid against all part bounding boxes
        face_centroids = mesh.triangles_center  # (n_faces, 3)
        for f_idx, centroid in enumerate(face_centroids):
            current_label = final_face_labels[f_idx]
            potential_label = []
            for lbl, (min_c, max_c) in part_bboxes.items():
                if np.all(centroid >= min_c) and np.all(centroid <= max_c):
                    if not potential_label:
                        potential_label = lbl
                    else:
                        current_volume = np.prod((max_c - min_c))
                        potential_obb_volume = np.prod((part_bboxes[potential_label][1] - part_bboxes[potential_label][0]))
                        if current_volume < potential_obb_volume:
                            potential_label = lbl
            if potential_label:
                final_face_labels[f_idx] = potential_label

        # ==== Step 6: check facets - coplanar facets should belong to the same part (strong assumption though) ====
        test = mesh.facets_boundary

        for facet in mesh.facets:
            labels = final_face_labels[facet]
            most_common = Counter(labels).most_common(1)[0][0]
            final_face_labels[facet] = most_common

        # ==== Step 7: create parts from assigned meshes ====
        part_bboxes = {}
        parts = []
        for lbl in np.unique(final_face_labels):
            if lbl == 0:
                continue
            mask = final_face_labels == lbl
            part_mesh = mesh.submesh([mask], append=True)
            # another bounding box check
            part_bboxes[lbl] = part_mesh.bounds

        for lbl in part_bboxes.keys():
            for j, obbs in part_bboxes.items():
                if lbl == j:
                    continue
                else:
                    inside = np.all((obbs[0] >= part_bboxes[lbl][0]) & (obbs[1] <= part_bboxes[lbl][1]))
                    if np.all(inside):
                        final_face_labels[final_face_labels == j] = lbl

        for lbl in np.unique(final_face_labels):
            if lbl == 0:
                continue
            mask = final_face_labels == lbl
            part_mesh = mesh.submesh([mask], append=True)
            parts.append(part_mesh)

        print(f"Identified {len(parts)} parts after NN-based labeling")

        # visualise_parts(parts)

        # Filter parts
        if fLOD == 2:
            filtering_radius = 0.05
        elif fLOD == 3:
            filtering_radius = 0.005
        filtered_parts = []
        for part in parts:
            include_part = projection_fits_circle(part, radius=filtering_radius)
            if include_part:
                filtered_parts.append(part)

        print(f"Identified {len(filtered_parts)} parts after filtering")

        # visualise_parts(filtered_parts)

        # Shape healing
        if len(parts) == len(filtered_parts):
            # mesh.export("mesh_unchanged.stl")
            scene = trimesh.Scene([mesh])
            scene.show()
        else:
            cleaned_parts = []
            for part in filtered_parts:
                if not part.is_watertight:

                    # Boundary edges = edges that belong to only 1 triangle
                    boundary_edges = part.edges[trimesh.grouping.group_rows(part.edges_sorted, require_count=1)]

                    # Create a Path3D object from those edges so we can visualize them
                    boundary_path = trimesh.load_path(part.vertices[boundary_edges])

                    patches = []

                    # Iterate over each closed loop (entity) in the path
                    for entity in boundary_path.entities:
                        # Extract the vertices of this boundary loop
                        loop_vertices = boundary_path.vertices[entity.points]

                        # Remove duplicate vertices
                        loop_vertices, unique_idx = np.unique(loop_vertices.round(8), axis=0, return_index=True)
                        loop_vertices = loop_vertices[np.argsort(unique_idx)]

                        if len(loop_vertices) < 3:
                            print("Skipping degenerate loop")
                            continue

                        # Fit a plane to the loop vertices
                        plane_origin, plane_normal = trimesh.points.plane_fit(loop_vertices)

                        # Transform from 3D -> 2D (XY plane)
                        T = trimesh.geometry.plane_transform(origin=plane_origin, normal=plane_normal)
                        loop_2d = trimesh.transform_points(loop_vertices, T)[:, :2]  # drop z

                        # Build PSLG (planar straight line graph)
                        segments = np.column_stack([
                            np.arange(len(loop_2d)),
                            np.roll(np.arange(len(loop_2d)), -1)
                        ])
                        A = dict(vertices=loop_2d, segments=segments)

                        # Constrained Delaunay triangulation
                        t = triangle.triangulate(A, 'p')

                        tri_vertices_2d = t['vertices']
                        tri_faces = t['triangles']

                        # Back-project to 3D
                        tri_vertices_3d = trimesh.transform_points(
                            np.column_stack([tri_vertices_2d, np.zeros(len(tri_vertices_2d))]),
                            np.linalg.inv(T)
                        )

                        print("hi!")

                        patch_mesh = trimesh.Trimesh(vertices=tri_vertices_3d, faces=tri_faces)
                        # patch_mesh.show()
                        patches.append(patch_mesh)

                    # Merge patches into the original mesh
                    if patches:
                        part_filled = rebuild_mesh_with_patches(part, patches)
                        part_filled.fix_normals()
                        print(part_filled.is_watertight)
                        cleaned_parts.append(part)
                        part_filled.show()
                    else:
                        print("No patches created")
            if len(cleaned_parts) > 1:
                # Combine them into one mesh
                combined = concatenate(cleaned_parts)

                # Export to a single STL file
                # combined.export("combined_part.stl")
                combined.show()

            else:
                # Export to STL file
                cleaned_parts[0].export("combined_part.stl")
                cleaned_parts[0].show()