Upload vgnet.py

vgnet.py

@@ -0,0 +1,88 @@
+"""
+VGNetwork — Vertex Generator Network (MLP-only, no PointTransformerV3).
+Inputs: sample points + normals. Outputs: 3D displacement.
+"""
+import torch
+import torch.nn as nn
+import numpy as np
+from .embedder import get_embedder
+
+
+class VGNetwork(nn.Module):
+    def __init__(self,
+                 d_in=3,
+                 d_out=3,
+                 d_hidden=256,
+                 n_layers=8,
+                 skip_in=(4,),
+                 multires=8,
+                 scale=1.0,
+                 geometric_init=True,
+                 weight_norm=True):
+        super(VGNetwork, self).__init__()
+
+        dims = [d_in] + [d_hidden for _ in range(n_layers)] + [d_out]
+
+        self.embed_fn_fine = None
+        if multires > 0:
+            embed_fn, input_ch = get_embedder(multires, input_dims=d_in)
+            self.embed_fn_fine = embed_fn
+            dims[0] = input_ch + 3   # positional encoding + original xyz + normals
+        else:
+            dims[0] += 3  # add normals
+
+        self.num_layers = len(dims)
+        self.skip_in = skip_in
+        self.scale = scale
+
+        for l in range(0, self.num_layers - 1):
+            if l + 1 in self.skip_in:
+                out_dim = dims[l + 1] - dims[0]
+            else:
+                out_dim = dims[l + 1]
+
+            lin = nn.Linear(dims[l], out_dim)
+
+            if geometric_init:
+                if multires > 0 and l == 0:
+                    torch.nn.init.constant_(lin.bias, 0.0)
+                    torch.nn.init.constant_(lin.weight[:, 3:], 0.0)
+                    torch.nn.init.normal_(lin.weight[:, :3], 0.0, np.sqrt(2) / np.sqrt(out_dim))
+                elif multires > 0 and l in self.skip_in:
+                    torch.nn.init.constant_(lin.bias, 0.0)
+                    torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))
+                    torch.nn.init.constant_(lin.weight[:, -(dims[0] - 3):], 0.0)
+                else:
+                    torch.nn.init.constant_(lin.bias, 0.0)
+                    torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))
+
+            if weight_norm:
+                lin = nn.utils.weight_norm(lin)
+            setattr(self, "lin" + str(l), lin)
+
+        self.activation = nn.ReLU()
+
+    def forward(self, samples, normals):
+        """
+        Args:
+            samples:  (B, 3) query points
+            normals:  (B, 3) estimated normals at samples
+        Returns:
+            moving_pcd: (B, 3) displaced points = samples + delta
+        """
+        inputs = samples * self.scale
+        if self.embed_fn_fine is not None:
+            inputs = self.embed_fn_fine(inputs)
+        inputs = torch.cat((inputs, normals), dim=-1)
+
+        x = inputs
+        for l in range(0, self.num_layers - 1):
+            lin = getattr(self, "lin" + str(l))
+            if l in self.skip_in:
+                x = torch.cat([x, inputs], 1) / np.sqrt(2)
+            x = lin(x)
+            if l < self.num_layers - 2:
+                x = self.activation(x)
+
+        moving_pcd = samples + x / self.scale
+        return moving_pcd