add main optimization loop

point2mesh/optimize.py

@@ -0,0 +1,292 @@
+"""
+Core Point2Mesh optimisation loop.
+
+Implements the coarse-to-fine self-prior optimisation: at each level
+the CNN weights are re-initialised, a fixed random input is drawn,
+and the network learns to deform the mesh surface to match the target
+point cloud.  Between levels the mesh is subdivided and remeshed.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from dataclasses import dataclass, field
+from typing import Optional, Callable
+
+import numpy as np
+import torch
+
+from .mesh import Mesh, PartMesh, edge_to_vertex_displacement
+from .network import Point2MeshNet
+from .losses import sample_surface, chamfer_loss, beam_gap_loss, normal_loss
+from .io_utils import (
+    load_pointcloud,
+    build_initial_mesh,
+    remesh,
+    save_mesh,
+    estimate_normals,
+)
+
+logger = logging.getLogger("point2mesh")
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  Configuration
+# ──────────────────────────────────────────────────────────────────────
+@dataclass
+class Point2MeshConfig:
+    # Coarse-to-fine levels
+    n_levels: int = 4
+    iters_per_level: int = 1000
+
+    # Mesh resolution schedule
+    init_faces: int = 2000
+    face_growth: float = 1.5
+    max_faces: int = 20000
+
+    # Sampling
+    samples_start: int = 15000
+    samples_end: int = 50000
+
+    # Loss weights
+    lambda_beam: float = 1.0
+    lambda_normal: float = 0.1
+    beam_epsilon: float = 0.5
+
+    # Network
+    in_channels: int = 6
+    enc_channels: list = field(default_factory=lambda: [64, 128, 256, 256])
+    lr: float = 2e-4
+
+    # PartMesh (set > 0 to enable spatial partitioning)
+    part_threshold: int = 10000  # enable parts when #faces > this
+    n_parts: int = 2
+
+    # Misc
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    log_every: int = 50
+    save_intermediates: bool = False
+    output_dir: str = "."
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  Optimisation loop
+# ──────────────────────────────────────────────────────────────────────
+def run_point2mesh(
+    input_path: str,
+    output_path: str,
+    cfg: Optional[Point2MeshConfig] = None,
+    progress_callback: Optional[Callable] = None,
+) -> str:
+    """
+    Full Point2Mesh pipeline.
+
+    Parameters
+    ----------
+    input_path  : path to .ply / .pcd / .xyz / .obj point cloud
+    output_path : path to write final mesh (.obj / .ply / .stl)
+    cfg         : configuration dataclass (defaults are sensible)
+    progress_callback : optional fn(level, iteration, loss_val)
+
+    Returns
+    -------
+    output_path : the path where the mesh was saved
+    """
+    if cfg is None:
+        cfg = Point2MeshConfig()
+
+    device = torch.device(cfg.device)
+
+    # ── 1. Load point cloud ──────────────────────────────────────────
+    logger.info(f"Loading point cloud from {input_path}")
+    points_np, normals_np = load_pointcloud(input_path)
+    logger.info(f"  {len(points_np)} points loaded")
+
+    # Centre and scale to unit sphere
+    centroid = points_np.mean(axis=0)
+    points_np -= centroid
+    scale = np.abs(points_np).max()
+    if scale > 0:
+        points_np /= scale
+
+    X = torch.tensor(points_np, dtype=torch.float32, device=device)
+
+    # Normals
+    if normals_np is None:
+        logger.info("  Estimating normals …")
+        normals_np = estimate_normals(points_np)
+    else:
+        normals_np = normals_np.copy()
+        # Renormalize
+        norms = np.linalg.norm(normals_np, axis=1, keepdims=True)
+        normals_np /= np.clip(norms, 1e-8, None)
+
+    X_normals = torch.tensor(normals_np, dtype=torch.float32, device=device)
+
+    # ── 2. Build initial mesh (convex hull) ──────────────────────────
+    logger.info("Building initial mesh (convex hull) …")
+    init_verts, init_faces = build_initial_mesh(points_np, target_faces=cfg.init_faces)
+    logger.info(f"  Initial mesh: {len(init_verts)} verts, {len(init_faces)} faces")
+
+    mesh = Mesh(init_verts, init_faces, device=str(device))
+
+    # ── 3. Coarse-to-fine optimisation ───────────────────────────────
+    net = Point2MeshNet(
+        in_ch=cfg.in_channels,
+        enc_channels=cfg.enc_channels,
+    ).to(device)
+
+    for level in range(cfg.n_levels):
+        logger.info(f"\n{'='*60}")
+        logger.info(f"Level {level}  |  {mesh.n_faces} faces, {mesh.n_edges} edges")
+        logger.info(f"{'='*60}")
+
+        # Re-initialise network weights and random input
+        net.reset_weights()
+        C_l = torch.rand(1, cfg.in_channels, mesh.n_edges, device=device)
+
+        optimiser = torch.optim.Adam(net.parameters(), lr=cfg.lr)
+
+        # Optionally use PartMesh for large meshes
+        use_parts = mesh.n_faces > cfg.part_threshold
+        pmesh = PartMesh(mesh, cfg.n_parts) if use_parts else None
+        if use_parts:
+            logger.info(f"  Using PartMesh with {len(pmesh.parts)} parts")
+
+        t0 = time.time()
+        for it in range(cfg.iters_per_level):
+            optimiser.zero_grad()
+
+            # Linear ramp of sample count
+            frac = it / max(cfg.iters_per_level - 1, 1)
+            n_samples = int(cfg.samples_start + (cfg.samples_end - cfg.samples_start) * frac)
+
+            if use_parts:
+                loss_total = _forward_partmesh(
+                    pmesh, net, C_l, X, X_normals, n_samples, cfg, mesh, device
+                )
+            else:
+                loss_total = _forward_single(
+                    net, C_l, mesh, X, X_normals, n_samples, cfg
+                )
+
+            loss_total.backward()
+            optimiser.step()
+
+            if progress_callback:
+                progress_callback(level, it, loss_total.item())
+
+            if it % cfg.log_every == 0 or it == cfg.iters_per_level - 1:
+                elapsed = time.time() - t0
+                logger.info(
+                    f"  [{level}][{it:4d}/{cfg.iters_per_level}]  "
+                    f"loss={loss_total.item():.6f}  "
+                    f"({elapsed:.1f}s)"
+                )
+
+        # ── Apply final displacements ────────────────────────────────
+        with torch.no_grad():
+            delta_edges = net(C_l, mesh)
+            delta_v = edge_to_vertex_displacement(delta_edges, mesh)
+            new_verts = mesh.vs + delta_v
+
+        # Save intermediate
+        if cfg.save_intermediates:
+            import os
+            intermediate_path = os.path.join(
+                cfg.output_dir, f"level_{level}.obj"
+            )
+            save_mesh(
+                intermediate_path,
+                (new_verts.cpu().numpy() * scale) + centroid,
+                mesh.faces.cpu().numpy(),
+            )
+            logger.info(f"  Saved intermediate → {intermediate_path}")
+
+        # ── Remesh for next level ────────────────────────────────────
+        if level < cfg.n_levels - 1:
+            target = min(
+                int(mesh.n_faces * cfg.face_growth), cfg.max_faces
+            )
+            logger.info(f"  Remeshing {mesh.n_faces} → {target} faces …")
+            new_v_np = new_verts.cpu().numpy()
+            new_f_np = mesh.faces.cpu().numpy()
+            rem_v, rem_f = remesh(new_v_np, new_f_np, target)
+            mesh = Mesh(rem_v, rem_f, device=str(device))
+        else:
+            # Final level — just update vertex positions
+            mesh.vs = new_verts
+
+    # ── 4. Write output ──────────────────────────────────────────────
+    final_verts = mesh.vs.detach().cpu().numpy() * scale + centroid
+    final_faces = mesh.faces.cpu().numpy()
+
+    save_mesh(output_path, final_verts, final_faces)
+    logger.info(f"\nDone! Mesh saved to {output_path}")
+    logger.info(f"  {len(final_verts)} vertices, {len(final_faces)} faces")
+
+    return output_path
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  Forward pass helpers
+# ──────────────────────────────────────────────────────────────────────
+def _forward_single(net, C_l, mesh, X, X_normals, n_samples, cfg):
+    """Standard forward on the full mesh."""
+    delta_edges = net(C_l, mesh)
+    delta_v = edge_to_vertex_displacement(delta_edges, mesh)
+    V_new = mesh.vs + delta_v
+
+    Y, face_ids = sample_surface(V_new, mesh.faces, n_samples)
+
+    # Face normals at sampled points
+    fnormals = mesh.face_normals(V_new)
+    Y_normals = fnormals[face_ids]
+
+    loss = chamfer_loss(X, Y)
+
+    if cfg.lambda_beam > 0:
+        loss = loss + cfg.lambda_beam * beam_gap_loss(
+            Y, Y_normals, X, epsilon=cfg.beam_epsilon
+        )
+
+    if cfg.lambda_normal > 0 and X_normals is not None:
+        loss = loss + cfg.lambda_normal * normal_loss(Y, Y_normals, X, X_normals)
+
+    return loss
+
+
+def _forward_partmesh(pmesh, net, C_l_full, X, X_normals, n_samples, cfg, full_mesh, device):
+    """
+    Forward pass with PartMesh splitting.
+    Process each spatial part independently, accumulate gradients.
+    """
+    n_parts = len(pmesh.parts)
+    samples_per_part = max(n_samples // n_parts, 1000)
+
+    total_loss = torch.tensor(0.0, device=device, requires_grad=True)
+
+    for part_idx, part_mesh in enumerate(pmesh.parts):
+        # Create random input for this part's edge count
+        C_part = torch.rand(1, cfg.in_channels, part_mesh.n_edges, device=device)
+
+        delta_edges = net(C_part, part_mesh)
+        delta_v = edge_to_vertex_displacement(delta_edges, part_mesh)
+        V_new = part_mesh.vs + delta_v
+
+        Y, face_ids = sample_surface(V_new, part_mesh.faces, samples_per_part)
+        fnormals = part_mesh.face_normals(V_new)
+        Y_normals = fnormals[face_ids]
+
+        loss = chamfer_loss(X, Y)
+        if cfg.lambda_beam > 0:
+            loss = loss + cfg.lambda_beam * beam_gap_loss(
+                Y, Y_normals, X, epsilon=cfg.beam_epsilon
+            )
+        if cfg.lambda_normal > 0 and X_normals is not None:
+            loss = loss + cfg.lambda_normal * normal_loss(Y, Y_normals, X, X_normals)
+
+        total_loss = total_loss + loss / n_parts
+
+    return total_loss