"""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)