Upload point_sam/model/common.py

point_sam/model/common.py

@@ -0,0 +1,602 @@
+# https://github.com/baaivision/Uni3D/blob/main/models/point_encoder.py
+# Modified: torkit3d dependencies replaced with pure-PyTorch implementations
+from typing import Union
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def sample_farthest_points(points: torch.Tensor, num_samples: int):
+    """Pure PyTorch farthest point sampling (FPS).
+
+    Args:
+        points: [B, N, 3]. Input point clouds.
+        num_samples: int. Number of points to sample.
+
+    Returns:
+        torch.Tensor: [B, num_samples]. Indices of sampled points.
+    """
+    device = points.device
+    batch_size, num_points, _ = points.shape
+    indices = torch.zeros(batch_size, num_samples, dtype=torch.long, device=device)
+    distances = torch.ones(batch_size, num_points, device=device) * float('inf')
+
+    # Start from a random point
+    farthest = torch.randint(0, num_points, (batch_size,), dtype=torch.long, device=device)
+
+    for i in range(num_samples):
+        indices[:, i] = farthest
+        centroid = points[torch.arange(batch_size, device=device), farthest, :].view(batch_size, 1, 3)
+        dist = torch.sum((points - centroid) ** 2, -1)
+        mask = dist < distances
+        distances[mask] = dist[mask]
+        farthest = torch.max(distances, dim=1)[1]
+
+    return indices
+
+
+def batch_index_select(data: torch.Tensor, index: torch.Tensor, dim: int):
+    """Batch index select — pure PyTorch implementation.
+
+    Args:
+        data: [B, N, C] tensor.
+        index: [B, K] indices.
+        dim: dimension to index along (after batch dim).
+
+    Returns:
+        torch.Tensor: [B, K, C] selected values.
+    """
+    batch_size = data.shape[0]
+    view_shape = [1] * data.dim()
+    view_shape[0] = batch_size
+    view_shape[dim] = -1
+    index = index.view(view_shape)
+    shape = list(data.shape)
+    shape[dim] = index.shape[dim]
+    index = index.expand(shape)
+    return torch.gather(data, dim, index)
+
+
+def chamfer_distance(x: torch.Tensor, y: torch.Tensor):
+    """Compute chamfer distance between two point clouds.
+
+    Args:
+        x: [B, N, 3]
+        y: [B, M, 3]
+
+    Returns:
+        min_dists_x: [B, N] minimum distances from x to y
+        min_idx_x: [B, N] indices of nearest neighbors in y
+    """
+    # x: [B, N, 3], y: [B, M, 3]
+    dist = torch.cdist(x, y)  # [B, N, M]
+    min_dists, min_idx = torch.min(dist, dim=2)
+    return min_dists, min_idx
+
+
+def fps(points: torch.Tensor, num_samples: int):
+    """A wrapper of farthest point sampling (FPS).
+
+    Args:
+        points: [B, N, 3]. Input point clouds.
+        num_samples: int. The number of points to sample.
+
+    Returns:
+        torch.Tensor: [B, num_samples, 3]. Sampled points.
+    """
+    idx = sample_farthest_points(points, num_samples)
+    sampled_points = batch_index_select(points, idx, dim=1)
+    return sampled_points
+
+
+def knn_points(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    k: int,
+    sorted: bool = False,
+    transpose: bool = False,
+):
+    """Compute k nearest neighbors.
+
+    Args:
+        query: [B, N1, D], query points. [B, D, N1] if @transpose is True.
+        key:  [B, N2, D], key points. [B, D, N2] if @transpose is True.
+        k: the number of nearest neighbors.
+        sorted: whether to sort the results
+        transpose: whether to transpose the last two dimensions.
+
+    Returns:
+        torch.Tensor: [B, N1, K], distances to the k nearest neighbors in the key.
+        torch.Tensor: [B, N1, K], indices of the k nearest neighbors in the key.
+    """
+    if transpose:
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+    # Compute pairwise distances, [B, N1, N2]
+    distance = torch.cdist(query, key)
+    if k == 1:
+        knn_dist, knn_ind = torch.min(distance, dim=2, keepdim=True)
+    else:
+        knn_dist, knn_ind = torch.topk(distance, k, dim=2, largest=False, sorted=sorted)
+    return knn_dist, knn_ind
+
+
+class KNNGrouper(nn.Module):
+    """Group points based on K nearest neighbors.
+
+    A number of points are sampled as centers by farthest point sampling (FPS).
+    Each group is formed by the center and its k nearest neighbors.
+    """
+
+    def __init__(self, num_groups, group_size, radius=None, centralize_features=False):
+        super().__init__()
+        self.num_groups = num_groups
+        self.group_size = group_size
+        self.radius = radius
+        self.centralize_features = centralize_features
+
+    def forward(self, xyz: torch.Tensor, features: torch.Tensor, use_fps=True):
+        """
+        Args:
+            xyz: [B, N, 3]. Input point clouds.
+            features: [B, N, C]. Point features.
+            use_fps: bool. Whether to use farthest point sampling.
+                If not, `xyz` should already be sampled by FPS.
+
+        Returns:
+            dict: {
+                features: [B, G, K, 3 + C]. Group features.
+                centers: [B, G, 3]. Group centers.
+                knn_idx: [B, G, K]. The indices of k nearest neighbors.
+            }
+        """
+        batch_size, num_points, _ = xyz.shape
+        with torch.no_grad():
+            if use_fps:
+                fps_idx = sample_farthest_points(xyz.float(), self.num_groups)
+                centers = batch_index_select(xyz, fps_idx, dim=1)
+            else:
+                fps_idx = torch.arange(self.num_groups, device=xyz.device)
+                fps_idx = fps_idx.expand(batch_size, -1)
+                centers = xyz[:, : self.num_groups]
+            _, knn_idx = knn_points(centers, xyz, self.group_size)  # [B, G, K]
+
+        batch_offset = torch.arange(batch_size, device=xyz.device) * num_points
+        batch_offset = batch_offset.reshape(-1, 1, 1)
+        knn_idx_flat = (knn_idx + batch_offset).reshape(-1)  # [B * G * K]
+
+        nbr_xyz = xyz.reshape(-1, 3)[knn_idx_flat]
+        nbr_xyz = nbr_xyz.reshape(batch_size, self.num_groups, self.group_size, 3)
+        nbr_xyz = nbr_xyz - centers.unsqueeze(2)  # [B, G, K, 3]
+        # NOTE: Follow PointNext to normalize the relative position
+        if self.radius is not None:
+            nbr_xyz = nbr_xyz / self.radius
+
+        nbr_feats = features.reshape(-1, features.shape[-1])[knn_idx_flat]
+        nbr_feats = nbr_feats.reshape(
+            batch_size, self.num_groups, self.group_size, features.shape[-1]
+        )
+
+        group_feats = [nbr_xyz, nbr_feats]
+        if self.centralize_features:
+            center_feats = batch_index_select(features, fps_idx, dim=1)
+            group_feats.append(nbr_feats - center_feats.unsqueeze(2))
+
+        group_feats = torch.cat(group_feats, dim=-1)
+        return dict(
+            features=group_feats, centers=centers, knn_idx=knn_idx, fps_idx=fps_idx
+        )
+
+
+def group_with_centers_and_knn(
+    xyz: torch.Tensor,
+    features: torch.Tensor,
+    centers: torch.Tensor,
+    knn_idx: torch.Tensor,
+    radius: float = None,
+    centralize_features: bool = False,
+    center_idx: torch.Tensor = None,
+):
+    """Group points based on K nearest neighbors.
+
+    Args:
+        xyz: [B, N, 3]. Input point clouds.
+        features: [B * M, N, C]. Point features. Support multiple features for the same point cloud.
+        centers: [B, L, 3]. Group centers.
+        knn_idx: [B, L, K]. The indices of k nearest neighbors.
+
+    Returns:
+        torch.Tensor: [B * M, L, K, 3 + C]. Group features.
+    """
+    assert xyz.dim() == features.dim(), (xyz.shape, features.shape)
+    assert xyz.shape[1] == features.shape[1], (xyz.shape, features.shape)
+    assert xyz.shape[0] == centers.shape[0] == knn_idx.shape[0]
+    assert knn_idx.shape[:2] == centers.shape[:2], (knn_idx.shape, centers.shape)
+
+    # 1. Compute neighborhood coordinates
+    batch_size, num_points, _ = xyz.shape
+    _, num_patches, patch_size = knn_idx.shape
+
+    batch_offset = torch.arange(batch_size, device=xyz.device) * num_points
+    batch_offset = batch_offset.reshape(-1, 1, 1)
+    knn_idx_flat = (knn_idx + batch_offset).reshape(-1)  # [B * L * K]
+
+    nbr_xyz = xyz.reshape(-1, 3)[knn_idx_flat]
+    nbr_xyz = nbr_xyz.reshape(batch_size, num_patches, patch_size, 3)
+    nbr_xyz = nbr_xyz - centers.unsqueeze(2)  # [B, L, K, 3]
+    if radius is not None:
+        nbr_xyz = nbr_xyz / radius
+
+    # 2. Compute neighborhood features
+    batch_size2 = features.shape[0]
+    repeats = features.shape[0] // xyz.shape[0]
+    knn_idx2 = torch.repeat_interleave(knn_idx, repeats, dim=0)  # [B*M,L,K]
+
+    batch_offset = torch.arange(batch_size2, device=xyz.device) * num_points
+    batch_offset = batch_offset.reshape(-1, 1, 1)
+    knn_idx_flat = (knn_idx2 + batch_offset).reshape(-1)  # [B*M*L*K]
+    nbr_feats = features.reshape(-1, features.shape[-1])[knn_idx_flat]
+    nbr_feats = nbr_feats.reshape(
+        batch_size2, num_patches, patch_size, features.shape[-1]
+    )
+
+    # 3. Concatenate features
+    nbr_xyz = torch.repeat_interleave(nbr_xyz, repeats, dim=0)
+    group_feats = [nbr_xyz, nbr_feats]
+    if centralize_features:
+        center_idx = torch.repeat_interleave(center_idx, repeats, dim=0)
+        center_feats = batch_index_select(features, center_idx, dim=1)
+        group_feats.append(nbr_feats - center_feats.unsqueeze(2))
+    return torch.cat(group_feats, dim=-1)
+
+
+class NNGrouper(nn.Module):
+    """Group points based on the nearest neighbors."""
+
+    def __init__(self, num_groups: int):
+        super().__init__()
+        self.num_groups = num_groups
+
+    def forward(self, xyz: torch.Tensor, features: torch.Tensor):
+        with torch.no_grad():
+            fps_idx = sample_farthest_points(xyz.float(), self.num_groups)
+            centers = batch_index_select(xyz, fps_idx, dim=1)
+            _, nn_idx = knn_points(xyz, centers, 1)  # [B, N, 1]
+
+        # Compute the relative position of each point to its nearest center
+        nn_idx = nn_idx.squeeze(-1)
+        nbr_xyz = xyz - batch_index_select(centers, nn_idx, dim=1)  # [B, N, 3]
+
+        # Normalize the relative position
+        dist = torch.linalg.norm(nbr_xyz, dim=-1, keepdim=True, ord=2)
+        nbr_xyz = nbr_xyz / torch.clamp(dist, min=1e-8)
+
+        group_feats = torch.cat([nbr_xyz, dist, features], dim=-1)
+        return dict(features=group_feats, centers=centers, nn_idx=nn_idx)
+
+
+def group_with_centers_and_nn(
+    xyz: torch.Tensor,
+    features: torch.Tensor,
+    centers: torch.Tensor,
+    nn_idx: torch.Tensor,
+):
+    """
+    Group points based on the voronoi diagram.
+
+    Args:
+        xyz: [B, N, 3]. Input point clouds.
+        features: [B, N, C]. Point features.
+        centers: [B, L, 3]. Group centers.
+        nn_idx: [B, N]. The indices of the nearest neighbors.
+
+    Returns:
+        torch.Tensor: [B, L, 3 + C]. Group features.
+    """
+    nbr_xyz = xyz - batch_index_select(centers, nn_idx, dim=1)  # [B, N, 3]
+    dist = torch.linalg.norm(nbr_xyz, dim=-1, keepdim=True, ord=2)
+    nbr_xyz = nbr_xyz / torch.clamp(dist, min=1e-8)
+    group_feats = torch.cat([nbr_xyz, dist, features], dim=-1)
+    return group_feats
+
+
+def compute_interp_weights(query: torch.Tensor, key: torch.Tensor, k=3, eps=1e-8):
+    """Compute interpolation weights for each query point.
+
+    Args:
+        query: [B, Nq, 3]. Query points.
+        key: [B, Nk, 3]. Key points.
+        k: int. The number of nearest neighbors.
+        eps: float. A small value to avoid division by zero.
+
+    Returns:
+        torch.Tensor: [B, Nq, K], indices of the k nearest neighbors in the key.
+        torch.Tensor: [B, Nq, K], interpolation weights.
+    """
+    dist, idx = knn_points(query, key, k)
+    inv_dist = 1.0 / torch.clamp(dist.square(), min=eps)
+    normalizer = torch.sum(inv_dist, dim=2, keepdim=True)
+    weight = inv_dist / normalizer  # [B, Nq, K]
+    return idx, weight
+
+
+def interpolate_features(x: torch.Tensor, index: torch.Tensor, weight: torch.Tensor):
+    """
+    Interpolates features based on the given index and weight.
+
+    Args:
+        x (torch.Tensor): The input tensor of shape (batch_size, num_keys, num_features).
+        index (torch.Tensor): The index tensor of shape (batch_size, num_queries, K).
+        weight (torch.Tensor): The weight tensor of shape (batch_size, num_queries, K).
+
+    Returns:
+        torch.Tensor: The interpolated features tensor of shape (batch_size, num_queries, num_features).
+    """
+    B, Nq, K = index.shape
+    batch_offset = torch.arange(B, device=x.device).reshape(-1, 1, 1) * x.shape[1]
+    index_flat = (index + batch_offset).flatten()  # [B*Nq*K]
+    _x = x.flatten(0, 1)[index_flat].reshape(B, Nq, K, x.shape[-1])
+    return (_x * weight.unsqueeze(-1)).sum(-2)
+
+
+def repeat_interleave(x: torch.Tensor, repeats: int, dim: int):
+    if repeats == 1:
+        return x
+    shape = list(x.shape)
+    shape.insert(dim + 1, 1)
+    shape[dim + 1] = repeats
+    x = x.unsqueeze(dim + 1).expand(shape).flatten(dim, dim + 1)
+    return x
+
+
+@torch.no_grad()
+def sample_prompts_adapter(
+    points: torch.Tensor,
+    gt_masks: torch.Tensor,
+    pred_logits: Union[torch.Tensor, None],
+    threshold: float = None,
+    is_eval = False,
+):
+    """Select prompt sampler based on iou."""
+    if pred_logits is None:
+        return sample_fixed_points(
+            points, gt_masks, pred_logits, threshold, from_error_region=True
+        )
+    else:
+        batch_size, num_masks, _ = gt_masks.shape
+
+        # if the batch iou is less than 0.5, use fixed sampler
+        gt_masks_copy = gt_masks.reshape(batch_size * num_masks, -1)
+        if threshold is None:
+            pred_masks = pred_logits > 0
+        else:
+            pred_masks = pred_logits.sigmoid() > threshold
+
+        iou = (gt_masks_copy & pred_masks).sum() / (gt_masks_copy | pred_masks).sum()
+        if iou < 1 or is_eval:
+            return sample_fixed_points(
+                points, gt_masks, pred_logits, threshold, from_error_region=False
+            )
+        else:
+            return sample_prompts(points, gt_masks, pred_logits, threshold)
+
+
+@torch.no_grad()
+def sample_prompts(
+    points: torch.Tensor,
+    gt_masks: torch.Tensor,
+    pred_logits: Union[torch.Tensor, None],
+    threshold: float = None,
+):
+    """Sample prompts from point clouds given ground-truth and predicted masks.
+
+    Args:
+        points: [B, N, 3]. Input point clouds.
+        gt_masks: [B, M, N], bool. Ground-truth (binary) masks.
+        pred_logits: A float tensor of shape [B*M, N]. Predicted logits.
+            If None, the prompt points will be sampled from the ground-truth masks.
+
+    Returns:
+        torch.Tensor: [B*M, 1, 3]. Prompt points.
+        torch.Tensor: [B*M, 1], bool. Prompt labels.
+    """
+    batch_size, num_masks, _ = gt_masks.shape
+
+    # The prompt point will be sampled from the error region.
+    if pred_logits is None:
+        diff_masks = gt_masks
+    else:
+        pred_logits = pred_logits.reshape(batch_size, num_masks, -1)
+        assert gt_masks.shape == pred_logits.shape, (gt_masks.shape, pred_logits.shape)
+        if threshold is None:
+            pred_masks = pred_logits > 0
+        else:
+            pred_masks = pred_logits.sigmoid() > threshold
+        diff_masks = gt_masks != pred_masks
+
+    prompt_coords, prompt_labels = [], []
+    for i in range(batch_size):
+        for j in range(num_masks):
+            diff_inds = torch.nonzero(diff_masks[i, j])  # [?, 1]
+            if len(diff_inds) == 0:
+                diff_inds = torch.nonzero(gt_masks[i, j])
+            diff_inds = diff_inds.squeeze(1)  # [?]
+            idx = diff_inds[torch.randint(0, len(diff_inds), [1])]
+            prompt_coords.append(points[i][idx])
+            prompt_labels.append(gt_masks[i, j][idx])
+
+    prompt_coords = torch.stack(prompt_coords)
+    prompt_labels = torch.stack(prompt_labels)
+    return prompt_coords, prompt_labels
+
+
+@torch.no_grad()
+def sample_fixed_points(
+    points: torch.Tensor,
+    gt_masks: torch.Tensor,
+    pred_logits: Union[torch.Tensor, None],
+    threshold: float = None,
+    from_error_region: bool = False,
+):
+    """Sample prompts from point clouds given ground-truth and predicted masks.
+
+    Args:
+        points: [B, N, 3]. Input point clouds.
+        gt_masks: [B, M, N], bool. Ground-truth (binary) masks.
+        pred_logits: A float tensor of shape [B*M, N]. Predicted logits.
+            If None, the prompt points will be sampled from the ground-truth masks.
+
+    Returns:
+        torch.Tensor: [B*M, 1, 3]. Prompt points.
+        torch.Tensor: [B*M, 1], bool. Prompt labels.
+    """
+    batch_size, num_masks, _ = gt_masks.shape
+
+    # The prompt point will be sampled from the error region.
+    if pred_logits is None:
+        fn = gt_masks
+        fp = torch.zeros_like(fn)
+    else:
+        pred_logits = pred_logits.reshape(batch_size, num_masks, -1)
+        assert gt_masks.shape == pred_logits.shape, (gt_masks.shape, pred_logits.shape)
+        if threshold is None:
+            pred_masks = pred_logits > 0
+        else:
+            pred_masks = pred_logits.sigmoid() > threshold
+        fn = gt_masks & ~pred_masks
+        fp = ~gt_masks & pred_masks
+
+    prompt_points, prompt_labels = [], []
+    if from_error_region:
+        mask = fn | fp
+        for i in range(batch_size):
+            for j in range(num_masks):
+                coords, label, _ = sample_furthest_points_from_border(
+                    points[i], mask[i, j], gt_masks[i, j]
+                )
+                prompt_points.append(coords)
+                prompt_labels.append(label)
+    else:
+        for i in range(batch_size):
+            for j in range(num_masks):
+                pprompt_coord, pprompt_label, pdist = (
+                    sample_furthest_points_from_border(
+                        points[i], fn[i, j], gt_masks[i, j]
+                    )
+                )
+                nprompt_coord, nprompt_label, ndist = (
+                    sample_furthest_points_from_border(
+                        points[i], fp[i, j], gt_masks[i, j]
+                    )
+                )
+                if pdist > ndist:
+                    prompt_points.append(pprompt_coord)
+                    prompt_labels.append(pprompt_label)
+                elif ndist == -1:
+                    pprompt_coord, pprompt_label, pdist = (
+                        sample_furthest_points_from_border(
+                            points[i], gt_masks[i, j], gt_masks[i, j]
+                        )
+                    )
+                    prompt_points.append(pprompt_coord)
+                    prompt_labels.append(pprompt_label)
+                else:
+                    prompt_points.append(nprompt_coord)
+                    prompt_labels.append(nprompt_label)
+
+    prompt_points = torch.stack(prompt_points)
+    prompt_labels = torch.stack(prompt_labels)
+    return prompt_points, prompt_labels
+
+
+def sample_furthest_points_from_border(
+    coords: torch.Tensor, lables: torch.Tensor, gt: torch.Tensor
+):
+    """
+    Sample points from the border of the mask.
+
+    Args:
+        coords: [N, 3]. Input point clouds.
+        lables: [N]. Point labels.
+        gt: [N]. Ground-truth labels.
+    """
+    bg_inds = lables == 0
+    fg_inds = lables == 1
+
+    # if bg_inds or fg_inds is empty, return None
+    if bg_inds.sum() == 0 or fg_inds.sum() == 0:
+        return None, None, -1
+
+    # All distances from foreground points to background points
+    min_dists, _ = chamfer_distance(coords[fg_inds][None, ...], coords[bg_inds][None, ...])
+
+    # Sample the farthest points from the border
+    center_idx = torch.argmax(min_dists)
+    center_coords = coords[fg_inds][center_idx]
+    center_dist = torch.max(min_dists)
+    center_label = gt[fg_inds][center_idx]
+
+    return center_coords[None, ...], center_label[None, ...], center_dist
+
+
+class PatchEncoder(nn.Module):
+    """Encode point patches following the PointNet structure for segmentation."""
+
+    def __init__(self, in_channels, out_channels, hidden_dims: list[int]):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        # NOTE: The original Uni3D implementation uses BatchNorm1d, while we use LayerNorm.
+        self.conv1 = nn.Sequential(
+            nn.Linear(in_channels, hidden_dims[0]),
+            nn.LayerNorm(hidden_dims[0]),
+            nn.GELU(),
+            nn.Linear(hidden_dims[0], hidden_dims[0]),
+        )
+        self.conv2 = nn.Sequential(
+            nn.Linear(hidden_dims[0] * 2, hidden_dims[1]),
+            nn.LayerNorm(hidden_dims[1]),
+            nn.GELU(),
+            nn.Linear(hidden_dims[1], out_channels),
+        )
+
+    def forward(self, point_patches: torch.Tensor):
+        # point_patches: [B, L, K, C_in]
+        x = self.conv1(point_patches)
+        y = torch.max(x, dim=-2, keepdim=True).values
+        x = torch.cat([y.expand_as(x), x], dim=-1)
+        x = self.conv2(x)  # [B, L, K, C_out]
+        y = torch.max(x, dim=-2).values  # [B, L, C_out]
+        return y
+
+
+class PatchEncoderNN(nn.Module):
+    def __init__(self, in_channels, out_channels, hidden_dims: list[int]) -> None:
+        super().__init__()
+        self.conv1 = nn.Sequential(
+            nn.Linear(in_channels, hidden_dims[0]),
+            nn.LayerNorm(hidden_dims[0]),
+            nn.GELU(),
+            nn.Linear(hidden_dims[0], hidden_dims[0]),
+        )
+        self.conv2 = nn.Sequential(
+            nn.Linear(hidden_dims[0] * 2, hidden_dims[1]),
+            nn.LayerNorm(hidden_dims[1]),
+            nn.GELU(),
+            nn.Linear(hidden_dims[1], out_channels),
+        )
+
+    def forward(self, point_patches: torch.Tensor, nn_idx: torch.Tensor, center_number: int) -> torch.Tensor:
+        # point_patches: [B, N, C_in]
+        x = self.conv1(point_patches)
+        y = torch.zeros([x.shape[0], center_number, x.shape[-1]], device=x.device, dtype=x.dtype)
+        y = torch.scatter_reduce(y, 1, nn_idx, x, "max")
+        x_max = torch.zeros_like(x)
+        x_max = torch.gather(y, 1, nn_idx.unsqueeze(-1).expand_as(y))
+        x = torch.cat([x_max, x], dim=-1)
+        x = self.conv2(x)
+        y = torch.zeros([x.shape[0], center_number, x.shape[-1]], device=x.device, dtype=x.dtype)
+        y = torch.scatter_reduce(y, 1, nn_idx, x, "max")
+        return y