add Point2Mesh U-Net network

point2mesh/network.py

@@ -0,0 +1,169 @@
+"""
+Point2Mesh network — MeshCNN-based U-Net that regresses per-edge
+vertex displacements from a fixed random input signal.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from typing import List, Optional
+from .mesh import Mesh
+from .layers import MeshConv, MeshPool, MeshUnpool
+
+
+class ConvBlock(nn.Module):
+    """MeshConv + GroupNorm + ReLU."""
+
+    def __init__(self, in_ch: int, out_ch: int, n_groups: int = 16):
+        super().__init__()
+        self.conv = MeshConv(in_ch, out_ch)
+        # Ensure n_groups divides out_ch
+        ng = min(n_groups, out_ch)
+        while out_ch % ng != 0:
+            ng -= 1
+        self.norm = nn.GroupNorm(ng, out_ch)
+        self.act = nn.ReLU(inplace=True)
+
+    def forward(self, x: torch.Tensor, mesh: Mesh) -> torch.Tensor:
+        x = self.conv(x, mesh)
+        x = self.norm(x)
+        x = self.act(x)
+        return x
+
+
+class Point2MeshNet(nn.Module):
+    """
+    U-Net encoder-decoder on mesh edges.
+
+    Encoder: Conv → Pool → Conv → Pool → Conv → Pool
+    Decoder: Unpool + skip → Conv → Unpool + skip → Conv → Unpool + skip → Conv
+    Output : Linear → [N_e, 6] → reshaped to [N_e, 2, 3] edge displacements
+
+    The input is a *fixed* random tensor C_l ∈ [0,1)^{N_e × in_ch} that is
+    NOT optimised; only the network weights are.
+    """
+
+    def __init__(
+        self,
+        in_ch: int = 6,
+        enc_channels: List[int] = None,
+        pool_targets: Optional[List[int]] = None,
+    ):
+        """
+        Parameters
+        ----------
+        in_ch        : feature dim of the random input (default 6)
+        enc_channels : channel widths for encoder stages (default [64,128,256,256])
+        pool_targets : #edges after each pool stage — **set at runtime** from mesh
+        """
+        super().__init__()
+        if enc_channels is None:
+            enc_channels = [64, 128, 256, 256]
+
+        self.in_ch = in_ch
+        self.enc_channels = enc_channels
+
+        # ── Encoder ──
+        self.enc_convs = nn.ModuleList()
+        ch = in_ch
+        for out in enc_channels:
+            self.enc_convs.append(ConvBlock(ch, out))
+            ch = out
+
+        # ── Pools (instantiated lazily when pool_targets is known) ──
+        self.pools: Optional[nn.ModuleList] = None
+
+        # ── Decoder (symmetric) ──
+        dec_channels = list(reversed(enc_channels[:-1])) + [enc_channels[0]]
+        self.dec_convs = nn.ModuleList()
+        for i, out in enumerate(dec_channels):
+            # skip connection doubles the channel count
+            skip_ch = enc_channels[-(i + 2)] if i < len(enc_channels) - 1 else in_ch
+            self.dec_convs.append(ConvBlock(ch + skip_ch, out))
+            ch = out
+
+        self.unpools = nn.ModuleList([MeshUnpool() for _ in dec_channels])
+
+        # Output head: edge displacement [N_e, 6] → reshaped [N_e, 2, 3]
+        self.head = nn.Conv1d(ch, 6, 1)
+        # Initialise head to zero so the first forward pass predicts no displacement
+        nn.init.zeros_(self.head.weight)
+        nn.init.zeros_(self.head.bias)
+
+    # ------------------------------------------------------------------
+    def _ensure_pools(self, n_edges: int):
+        """Create pool layers with target counts proportional to the mesh."""
+        if self.pools is not None:
+            return
+        n_pools = len(self.enc_channels) - 1  # one fewer pool than conv
+        targets = []
+        ne = n_edges
+        for _ in range(n_pools):
+            ne = max(ne // 2, 64)
+            targets.append(ne)
+        self.pools = nn.ModuleList([MeshPool(t) for t in targets])
+
+    # ------------------------------------------------------------------
+    def forward(
+        self,
+        C_l: torch.Tensor,   # (1, in_ch, N_e) — random input, fixed
+        mesh: Mesh,
+    ) -> torch.Tensor:
+        """
+        Returns
+        -------
+        delta_edges : (N_e, 2, 3) vertex displacements at each edge endpoint
+        """
+        self._ensure_pools(mesh.n_edges)
+
+        x = C_l  # (1, in_ch, N_e)
+
+        # ── Encoder ──
+        skips = [x]
+        histories = []
+        mesh_levels = [mesh]
+        for i, conv in enumerate(self.enc_convs):
+            x = conv(x, mesh_levels[-1])
+            if i < len(self.pools):
+                skips.append(x)
+                x, mesh_pooled, hist = self.pools[i](x, mesh_levels[-1])
+                histories.append(hist)
+                mesh_levels.append(mesh_pooled)
+
+        # ── Decoder ──
+        for i, (conv, unpool) in enumerate(zip(self.dec_convs, self.unpools)):
+            hist = histories[-(i + 1)]
+            x = unpool(x, hist)
+            skip = skips[-(i + 1)]
+            # Match sizes (in case of rounding during pool/unpool)
+            if x.shape[-1] != skip.shape[-1]:
+                min_e = min(x.shape[-1], skip.shape[-1])
+                x = x[:, :, :min_e]
+                skip = skip[:, :, :min_e]
+            x = torch.cat([x, skip], dim=1)
+            # Use the corresponding mesh level for this decoder stage
+            mesh_idx = len(mesh_levels) - 2 - i
+            x = conv(x, mesh_levels[max(mesh_idx, 0)])
+
+        # ── Head ──
+        x = self.head(x)  # (1, 6, N_e)
+        x = x.squeeze(0).t()  # (N_e, 6)
+        delta_edges = x.reshape(-1, 2, 3)  # (N_e, 2, 3)
+        return delta_edges
+
+    def reset_weights(self):
+        """Re-initialize all weights (called at each coarse-to-fine level)."""
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.Conv1d)):
+                nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.GroupNorm):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+        # Always zero-init the output head
+        nn.init.zeros_(self.head.weight)
+        nn.init.zeros_(self.head.bias)
+        # Reset pool targets so they are re-created for new mesh
+        self.pools = None