experiments/ptv3/ptv3_code/encoder_wrapper.py

"""Adapter that lets PointTransformerV3 plug into our existing model pipeline.

Bridges our [B, T, ·] batched format with PT v3's flat [Σ T_i, ·] + batch-indices
format. Handles voxel deduplication via inverse mapping so output point count
matches input exactly, even when multiple input points fall in the same voxel.

Drop-in replacement for the Perceiver latent bottleneck:
  - Input: per-token features [B, T, in_dim] and coord [B, T, 3]
  - Output: per-point features [B, T, hidden_out]
  - Output then fed to the existing DETR-style segment decoder (cross-attn over T tokens).
"""
from __future__ import annotations

import torch
import torch.nn as nn

from .model import PointTransformerV3


class PTv3Encoder(nn.Module):
    """Wrap PointTransformerV3 in a [B, T, ·] interface."""

    def __init__(
        self,
        in_channels: int,
        hidden: int = 256,
        grid_size: float = 0.005,
        enc_channels: tuple = (32, 64, 128, 256, 256),
        enc_depths: tuple = (2, 2, 2, 6, 2),
        enc_num_head: tuple = (2, 4, 8, 16, 16),
        enc_patch_size: tuple = (1024, 1024, 1024, 1024, 1024),
        dec_channels: tuple = (256, 64, 128, 256),
        dec_depths: tuple = (2, 2, 2, 2),
        dec_num_head: tuple = (4, 4, 8, 16),
        dec_patch_size: tuple = (1024, 1024, 1024, 1024),
        stride: tuple = (2, 2, 2, 2),
        order: tuple = ("z", "hilbert"),
        enable_flash: bool = False,
        shuffle_orders: bool = True,
        drop_path: float = 0.3,
    ):
        super().__init__()
        if dec_channels[0] != hidden:
            raise ValueError(
                f"dec_channels[0]={dec_channels[0]} must equal hidden={hidden} "
                f"so PT v3 output dim matches the rest of the model."
            )
        self.hidden = hidden
        self.grid_size = grid_size

        # Gracefully fall back to standard attention if flash_attn isn't installed.
        # Same weights work in both modes (mathematically equivalent).
        if enable_flash:
            try:
                import flash_attn  # noqa: F401
            except ImportError:
                enable_flash = False

        self.ptv3 = PointTransformerV3(
            in_channels=in_channels,
            order=order,
            stride=stride,
            enc_depths=enc_depths,
            enc_channels=enc_channels,
            enc_num_head=enc_num_head,
            enc_patch_size=enc_patch_size,
            dec_depths=dec_depths,
            dec_channels=list(dec_channels),
            dec_num_head=dec_num_head,
            dec_patch_size=dec_patch_size,
            drop_path=drop_path,
            enable_flash=enable_flash,
            shuffle_orders=shuffle_orders,
            cls_mode=False,
        )

    def forward(self, coord: torch.Tensor, feat: torch.Tensor,
                mask: torch.Tensor | None = None) -> torch.Tensor:
        """
        Args:
            coord: [B, T, 3] xyz (normalized to roughly [-1, 1])
            feat:  [B, T, in_channels] per-token features
            mask:  [B, T] bool; True=valid. If None, all valid.

        Returns:
            [B, T, hidden] per-point features. Invalid positions are zeroed.
        """
        from addict import Dict

        B, T, _ = coord.shape
        device = coord.device

        if mask is None:
            valid = torch.ones(B, T, dtype=torch.bool, device=device)
        else:
            valid = mask.bool()

        flat_mask = valid.reshape(-1)                          # [B*T]
        coord_flat = coord.reshape(-1, 3)[flat_mask]           # [N_valid, 3]
        feat_flat = feat.reshape(-1, feat.shape[-1])[flat_mask]  # [N_valid, in_channels]
        batch_per_point = torch.arange(B, device=device).repeat_interleave(T)
        batch = batch_per_point[flat_mask]                     # [N_valid]
        N_valid = coord_flat.shape[0]

        # Deduplicate (batch, voxel) cells. Multiple input points falling in the
        # same voxel get merged into one PT v3 token; we map the output back to
        # all original points via the inverse mapping.
        coord_min = coord_flat.min(dim=0).values
        grid_coord = torch.div(
            coord_flat - coord_min, self.grid_size, rounding_mode="trunc"
        ).long()                                               # [N_valid, 3]
        gmax = grid_coord.max(dim=0).values + 1                # [3]
        stride_y = gmax[2].item()
        stride_x = gmax[1].item() * stride_y
        stride_b = gmax[0].item() * stride_x

        voxel_id = (
            batch * stride_b
            + grid_coord[:, 0] * stride_x
            + grid_coord[:, 1] * stride_y
            + grid_coord[:, 2]
        )
        unique_ids, inverse_idx = torch.unique(voxel_id, return_inverse=True)
        N_unique = unique_ids.shape[0]

        # For each unique cell, pick its first appearance to define the cell's coord/feat.
        # (PT v3 will internally rebuild structure, this is just initialization.)
        perm = torch.argsort(inverse_idx, stable=True)
        first_in_unique = torch.empty(N_unique, dtype=torch.long, device=device)
        # scatter: for each i in perm in order, last write wins; but with stable sort
        # the first occurrence has lowest position index, so we need to write in reverse.
        first_in_unique.scatter_(
            0, inverse_idx[perm].flip(0), perm.flip(0)
        )

        coord_u = coord_flat[first_in_unique]                  # [N_unique, 3]
        feat_u = feat_flat[first_in_unique].contiguous()       # [N_unique, in_channels]
        batch_u = batch[first_in_unique].contiguous()          # [N_unique]
        # Sort by batch (PT v3 expects ascending batch)
        sort_perm = torch.argsort(batch_u, stable=True)
        coord_u = coord_u[sort_perm].contiguous()
        feat_u = feat_u[sort_perm].contiguous()
        batch_u = batch_u[sort_perm].contiguous()
        # We need inverse mapping that says: for each original point, where is its unique-cell after sort?
        inv_sort = torch.empty_like(sort_perm)
        inv_sort[sort_perm] = torch.arange(N_unique, device=device)
        sorted_unique_idx_for_orig = inv_sort[inverse_idx]     # [N_valid]

        # Build the Point dict and run PT v3.
        data_dict = Dict(
            coord=coord_u,
            feat=feat_u,
            batch=batch_u,
            grid_size=self.grid_size,
        )
        out_point = self.ptv3(data_dict)
        out_feat_unique = out_point.feat                       # [N_unique_out, hidden]

        # PT v3 with cls_mode=False restores original (unique-cell) point count, so
        # N_unique_out == N_unique.
        if out_feat_unique.shape[0] != N_unique:
            raise RuntimeError(
                f"PT v3 output count {out_feat_unique.shape[0]} != input unique count {N_unique}. "
                f"Did you set cls_mode=False?"
            )

        # Map unique-cell features back to all original valid points (duplicates share features).
        per_point_feat = out_feat_unique[sorted_unique_idx_for_orig]  # [N_valid, hidden]

        # Scatter back into [B*T, hidden] then reshape to [B, T, hidden]
        out = torch.zeros(B * T, self.hidden,
                          device=device, dtype=per_point_feat.dtype)
        out[flat_mask] = per_point_feat
        return out.reshape(B, T, self.hidden)