pixal3d_mrp_mlx/pixal3d_mlx/cleanup.py

"""Raw cascade output -> a mesh you can actually put in a game.

The decoder emits ~4M vertices / 8M faces at 1024^3, and on a measured run that was
one good body (96.9% of faces) plus **52,838 disconnected fragments** under 100 faces
each, inconsistent winding, and 1.5% non-manifold edges. Those floaters are the
"stray shells" failure Pixal3D is known for. They are also why feeding the raw mesh
straight into o_voxel's remesher hangs: it is being asked to remesh fifty thousand
objects, most of them noise.

So the order matters and it is: **strip floaters -> fix winding -> decimate**. Doing
it the other way round spends all the time on geometry that gets thrown away.

Nothing here is Pixal3D-specific; it is ordinary mesh hygiene, deliberately kept out
of the model code.
"""

from __future__ import annotations

import numpy as np


def largest_components(mesh, keep_ratio: float = 0.01, min_faces: int = 100):
    """Drop disconnected fragments. Returns (mesh, stats).

    A component survives if it has at least `min_faces` faces AND at least
    `keep_ratio` of the largest component's faces. Two thresholds because neither
    alone is safe: an absolute floor keeps genuine small parts on a big model, and the
    relative test drops proportionally-large noise on a small one.

    Labelling runs as a union-find over the VERTEX graph rather than through
    `trimesh.face_adjacency`. Same answer, very different cost: building face
    adjacency on this mesh took ~240s, which is unaffordable when an operator does it
    on every job. Vertices must already be welded — the decoder emits per-voxel
    vertices, so without a merge every triangle is its own island.
    """
    import trimesh
    from scipy.sparse import coo_matrix
    from scipy.sparse.csgraph import connected_components as cc

    faces = np.asarray(mesh.faces)
    n_v = len(mesh.vertices)
    if not len(faces):
        return mesh, {"components": 0, "kept": 0, "faces_removed": 0}

    e = np.vstack([faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]])
    g = coo_matrix((np.ones(len(e), np.int8), (e[:, 0], e[:, 1])), shape=(n_v, n_v))
    n_comp, vert_label = cc(g, directed=False)

    face_label = vert_label[faces[:, 0]]
    sizes = np.bincount(face_label, minlength=n_comp)
    threshold = max(min_faces, int(sizes.max() * keep_ratio))
    keep_labels = np.where(sizes >= threshold)[0]

    mask = np.isin(face_label, keep_labels)
    removed = int((~mask).sum())
    out = mesh.submesh([np.where(mask)[0]], append=True) if removed else mesh
    return out, {"components": int((sizes > 0).sum()), "kept": len(keep_labels),
                 "faces_removed": removed, "threshold": threshold}


def manifold_remesh(mesh, divisions: int = 256):
    """Voxelise -> fill -> marching cubes. Trades the dual grid for a manifold surface.

    WHY THIS EXISTS. The decoder's output cannot be decimated below ~214k faces, and
    the cause is NOT boundary edges (32,370) as first assumed — it is **non-manifold
    edges (81,112)**, which quadric decimation refuses to collapse. Filling holes
    barely moved either number (boundaries 32,370 -> 30,990, floor 214k -> 210k).
    Those edges are the Flexible Dual Grid working as designed: it represents open and
    non-manifold surfaces on purpose. So getting past the floor means giving that up,
    deliberately, rather than tuning a decimator that was never going to win.

    Measured on the sample: 500k non-manifold faces at IoU 0.969 becomes 20k
    watertight, winding-consistent faces at **IoU 0.956** — 25x smaller for 1.3%.

    THE TRAP: `marching_cubes` returns vertices in VOXEL INDEX space (0..divisions),
    not world units. Translating without scaling leaves the mesh ~292x too large,
    which still exports and renders as a plausible object — it just silhouettes at
    IoU 0.08. Scale by pitch first.
    """
    import trimesh

    pitch = float(np.asarray(mesh.extents).max()) / divisions
    out = mesh.voxelized(pitch=pitch).fill().marching_cubes
    out.apply_scale(pitch)                                  # index space -> world
    out.apply_translation(np.asarray(mesh.bounds[0]) - np.asarray(out.bounds[0]))
    out.fix_normals()
    return out


def decimate(mesh, target_faces: int, max_passes: int = 8, step: float = 0.7):
    """Quadric decimation to a face budget, iteratively.

    A single `simplify` call will not reduce past roughly 4.4% of its input no matter
    what `target_reduction` (or `agg`) is asked for: from 7.99M faces, targets of
    200k, 50k and 20k all returned the same 351,535. Repeated smaller passes get
    further — 2.4M -> 719k -> 270k -> 238k — because each pass re-evaluates the
    quadrics on the already-collapsed mesh.

    It still converges to a floor near ~240k on this geometry, set by the ~180k
    BOUNDARY edges the Flexible Dual Grid produces for open surfaces; quadric
    decimation will not collapse those. Getting below that needs a remesh, not a
    decimator. The loop therefore stops when a pass stops making progress rather than
    spinning, and the caller gets whatever floor the mesh actually has.
    """
    import fast_simplification
    import trimesh

    v = np.asarray(mesh.vertices, np.float32)
    f = np.asarray(mesh.faces, np.int32)

    for _ in range(max_passes):
        if len(f) <= target_faces:
            break
        reduction = min(step, 1.0 - target_faces / len(f))
        nv, nf = fast_simplification.simplify(v, f, target_reduction=reduction)
        if len(nf) >= len(f) * 0.98:      # no meaningful progress: at the floor
            break
        v, f = np.asarray(nv, np.float32), np.asarray(nf, np.int32)

    return trimesh.Trimesh(v, f, process=False)


def clean(vertices, faces, target_faces: int | None = 100_000,
          keep_ratio: float = 0.01, min_faces: int = 100, manifold: bool = False,
          divisions: int = 256, log=print):
    """The whole hygiene pass. Returns (mesh, stats).

    `manifold=True` inserts a voxel remesh after floater removal, which is the only
    way past the ~214k non-manifold decimation floor. It is lossy and opt-in.
    """
    import trimesh

    mesh = trimesh.Trimesh(np.asarray(vertices), np.asarray(faces), process=False)
    before = (len(mesh.vertices), len(mesh.faces))

    # WELD FIRST. The decoder emits per-voxel vertices, so coincident corners are
    # duplicated and neither component labelling nor winding propagation can cross a
    # shared edge. Run before either and every triangle looks like its own island.
    mesh.merge_vertices()
    log(f"  welded     {before[0]:,} -> {len(mesh.vertices):,} verts")

    mesh, comp = largest_components(mesh, keep_ratio, min_faces)
    log(f"  floaters   {comp['components']:6} components -> {comp['kept']} kept, "
        f"{comp['faces_removed']:,} faces dropped")

    if manifold:
        mesh = manifold_remesh(mesh, divisions)
        log(f"  remeshed   manifold @{divisions}: {len(mesh.faces):,} faces, "
            f"watertight={mesh.is_watertight}")

    if target_faces:
        mesh = decimate(mesh, target_faces)
        log(f"  decimated  {before[1]:,} -> {len(mesh.faces):,} faces")

        # SECOND floater pass. Quadric decimation collapses thin features and pinches
        # the surface apart, so a mesh that was ONE component before decimation came
        # back as 7,257 at a 500k budget. Cleaning only beforehand therefore ships a
        # fragmented result while reporting a clean intermediate.
        mesh.merge_vertices()
        mesh, post = largest_components(mesh, keep_ratio, min_faces)
        if post["faces_removed"]:
            log(f"  re-cleaned {post['components']} components after decimation -> "
                f"{post['kept']} kept, {post['faces_removed']:,} faces dropped")
        comp = {**comp, "post_decimation_components": post["components"],
                "post_decimation_removed": post["faces_removed"]}

    # Normals LAST, deliberately. fix_normals walks face adjacency, which costs
    # minutes on the raw 7.7M-face mesh and seconds on a decimated one — and it is
    # the shipped mesh whose winding actually matters. Doing it before decimation
    # spends the time and then throws most of the result away.
    mesh.fix_normals()
    log(f"  normals    winding consistent={mesh.is_winding_consistent}")

    return mesh, {"before_verts": before[0], "before_faces": before[1],
                  "after_verts": len(mesh.vertices), "after_faces": len(mesh.faces),
                  **comp}