Upload learn_region_grow/lrg_net.py

learn_region_grow/lrg_net.py

@@ -0,0 +1,181 @@
+"""
+LrgNet — Dual-branch 1D PointNet for learnable region growing.
+
+This is a direct PyTorch port of the TensorFlow 1.x model described in:
+
+    LRGNet: Learnable Region Growing for Class-Agnostic Point Cloud Segmentation
+    Jingdao Chen, Zsolt Kira, Yong K. Cho
+    IEEE Robotics and Automation Letters (RAL), 2021
+    arXiv:2103.09160
+
+Architecture overview
+-------------------
+The network takes two point sets as input:
+  1. **Inlier branch**: the current region (points already assigned to the object).
+  2. **Neighbor branch**: candidate points lying on the region boundary.
+
+Each branch runs an independent 1D PointNet (shared weights between conv layers
+in the original TensorFlow code, but kept independent here for clarity).
+After local per-point convolutions, a global max-pool extracts a single
+feature vector summarising the whole set. That global vector is tiled back to
+match the point counts and concatenated with the per-point features.
+
+Two classification heads then predict:
+  - **add_head**    : per-neighbor binary logits  (should this point join the region?)
+  - **remove_head** : per-inlier  binary logits  (should this point leave the region?)
+
+Both heads are trained jointly with cross-entropy.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Tuple
+
+
+class PointNetBranch(nn.Module):
+    """
+    Single 1D PointNet branch: 1D conv layers + global max pool.
+
+    In the original TF code this is a sequence of:
+        Conv1D(13 -> 64), Conv1D(64 -> 64), Conv1D(64 -> 64),
+        Conv1D(64 -> 128), Conv1D(128 -> 512)
+    followed by max-pooling over the spatial (point) dimension.
+    """
+
+    def __init__(self, in_channels: int = 13):
+        super().__init__()
+        self.conv1 = nn.Conv1d(in_channels, 64, 1)
+        self.conv2 = nn.Conv1d(64, 64, 1)
+        self.conv3 = nn.Conv1d(64, 64, 1)
+        self.conv4 = nn.Conv1d(64, 128, 1)
+        self.conv5 = nn.Conv1d(128, 512, 1)
+
+        self.bn1 = nn.BatchNorm1d(64)
+        self.bn2 = nn.BatchNorm1d(64)
+        self.bn3 = nn.BatchNorm1d(64)
+        self.bn4 = nn.BatchNorm1d(128)
+        self.bn5 = nn.BatchNorm1d(512)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Parameters
+        ----------
+        x : torch.Tensor, shape (B, C, N)
+
+        Returns
+        -------
+        local_feat : torch.Tensor, shape (B, 512, N)
+            Per-point features from the deepest conv layer.
+        global_feat : torch.Tensor, shape (B, 512, 1)
+            Max-pooled global vector.
+        """
+        # x: (B, C, N)
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = F.relu(self.bn3(self.conv3(x)))
+        x = F.relu(self.bn4(self.conv4(x)))
+        x = F.relu(self.bn5(self.conv5(x)))
+        local_feat = x  # (B, 512, N)
+        global_feat = torch.max(x, dim=2, keepdim=True)[0]  # (B, 512, 1)
+        return local_feat, global_feat
+
+
+class LrgNet(nn.Module):
+    """
+    LrgNet — Dual-branch network for learned region growing.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of feature channels per point (default 13 from the paper).
+    lite : int
+        0 = full channels, 1 = half, 2 = quarter.
+        Lite variants run faster on edge devices with negligible accuracy loss.
+    """
+
+    def __init__(self, in_channels: int = 13, lite: int = 0):
+        super().__init__()
+        factor = 1 / (2 ** lite)  # 1, 0.5, 0.25
+        c0 = int(64 * factor)
+        c1 = int(64 * factor)
+        c2 = int(64 * factor)
+        c3 = int(128 * factor)
+        c4 = int(512 * factor)
+
+        # Independent branches (original TF code shares conv weights conceptually,
+        # but we keep them separate to avoid accidental information leakage).
+        self.inlier_branch = self._make_branch(in_channels, c0, c1, c2, c3, c4)
+        self.neighbor_branch = self._make_branch(in_channels, c0, c1, c2, c3, c4)
+
+        # Classification heads
+        # Input: 512 (local) + 512 (global inlier) + 512 (global neighbor) = 1536
+        self.add_head = self._make_head(c4 * 3, 256, 128, 1)
+        self.remove_head = self._make_head(c4 * 3, 256, 128, 1)
+
+    def _make_branch(self, cin, c0, c1, c2, c3, c4):
+        layers = [
+            nn.Conv1d(cin, c0, 1), nn.BatchNorm1d(c0), nn.ReLU(),
+            nn.Conv1d(c0, c1, 1), nn.BatchNorm1d(c1), nn.ReLU(),
+            nn.Conv1d(c1, c2, 1), nn.BatchNorm1d(c2), nn.ReLU(),
+            nn.Conv1d(c2, c3, 1), nn.BatchNorm1d(c3), nn.ReLU(),
+            nn.Conv1d(c3, c4, 1), nn.BatchNorm1d(c4), nn.ReLU(),
+        ]
+        return nn.Sequential(*layers)
+
+    def _make_head(self, cin, h1, h2, out):
+        return nn.Sequential(
+            nn.Conv1d(cin, h1, 1),
+            nn.BatchNorm1d(h1),
+            nn.ReLU(),
+            nn.Conv1d(h1, h2, 1),
+            nn.BatchNorm1d(h2),
+            nn.ReLU(),
+            nn.Conv1d(h2, out, 1),
+        )
+
+    def forward(self, inliers: torch.Tensor, neighbors: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Parameters
+        ----------
+        inliers  : torch.Tensor, shape (B, C, Ni)
+            Current region points.
+        neighbors : torch.Tensor, shape (B, C, Nn)
+            Candidate boundary points.
+
+        Returns
+        -------
+        add_logits    : torch.Tensor, shape (B, 1, Nn)
+            Log-odds for adding each neighbor.
+        remove_logits : torch.Tensor, shape (B, 1, Ni)
+            Log-odds for removing each inlier.
+        """
+        # Run branches
+        inlier_local = self.inlier_branch(inliers)   # (B, c4, Ni)
+        neighbor_local = self.neighbor_branch(neighbors)  # (B, c4, Nn)
+
+        # Global max-pool
+        inlier_global = torch.max(inlier_local, dim=2, keepdim=True)[0]      # (B, c4, 1)
+        neighbor_global = torch.max(neighbor_local, dim=2, keepdim=True)[0]  # (B, c4, 1)
+
+        # Tile globals to match point counts
+        inlier_global_tiled = inlier_global.expand(-1, -1, inliers.shape[2])    # (B, c4, Ni)
+        neighbor_global_tiled = neighbor_global.expand(-1, -1, neighbors.shape[2])  # (B, c4, Nn)
+
+        # Fuse for add head: neighbor local + neighbor global + inlier global
+        add_input = torch.cat([
+            neighbor_local,
+            neighbor_global_tiled,
+            inlier_global.expand(-1, -1, neighbors.shape[2])
+        ], dim=1)  # (B, c4*3, Nn)
+        add_logits = self.add_head(add_input)  # (B, 1, Nn)
+
+        # Fuse for remove head: inlier local + inlier global + neighbor global
+        remove_input = torch.cat([
+            inlier_local,
+            inlier_global_tiled,
+            neighbor_global.expand(-1, -1, inliers.shape[2])
+        ], dim=1)  # (B, c4*3, Ni)
+        remove_logits = self.remove_head(remove_input)  # (B, 1, Ni)
+
+        return add_logits, remove_logits