Spaces:
Running on Zero
Running on Zero
| """ | |
| Transform layer for registering previous frame to current frame using Kabsch algorithm. | |
| Handles batched rigid body registration with known vertex correspondences. | |
| Supports two modes: | |
| - Standard Kabsch: Non-differentiable SVD (requires STE for gradient flow) | |
| - RoMa Kabsch: Fully differentiable (gradients flow through transformation) | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| try: | |
| import roma | |
| ROMA_AVAILABLE = True | |
| except ImportError: | |
| ROMA_AVAILABLE = False | |
| roma = None | |
| class TransformLayer(nn.Module): | |
| """ | |
| Registers previous frame positions to current frame positions using optimal rigid transformation. | |
| Uses Kabsch algorithm for finding the best rotation and translation. | |
| Input shapes: | |
| batch_last_position: (B, N, 3) - positions from previous frame | |
| batch_current_position: (B, N, 3) - positions from current frame (target) | |
| Output shape: | |
| (B, N, 3) - transformed previous positions aligned to current frame | |
| """ | |
| def __init__(self, use_scale=False, regularize_rotation=False, if_roma=False): | |
| """ | |
| Args: | |
| use_scale: If True, use Umeyama algorithm (similarity transform with scale). | |
| If False, use Kabsch algorithm (rigid transform only). | |
| regularize_rotation: If True, add small regularization to avoid numerical issues. | |
| if_roma: If True, use RoMa's differentiable rigid_points_registration. | |
| If False, use standard SVD-based Kabsch (non-differentiable). | |
| """ | |
| super().__init__() | |
| self.use_scale = use_scale | |
| self.regularize_rotation = regularize_rotation | |
| self.if_roma = if_roma | |
| if self.if_roma: | |
| if not ROMA_AVAILABLE: | |
| raise ImportError( | |
| "RoMa library is not installed. Install with: pip install roma" | |
| ) | |
| if self.use_scale: | |
| raise ValueError("RoMa mode does not support scale (Umeyama). Set use_scale=False.") | |
| def forward(self, batch_last_position, batch_current_position): | |
| """ | |
| Register batch_last_position to batch_current_position. | |
| Args: | |
| batch_last_position: (B, N, 3) - source positions to be transformed | |
| batch_current_position: (B, N, 3) - target positions to align to | |
| Returns: | |
| transformed_positions: (B, N, 3) - batch_last_position after rigid transformation | |
| """ | |
| B, N, D = batch_last_position.shape | |
| assert D == 3, f"Expected 3D points, got {D}D" | |
| assert batch_last_position.shape == batch_current_position.shape | |
| # RoMa mode: fully differentiable Kabsch | |
| if self.if_roma: | |
| # roma.rigid_points_registration: X -> Y | |
| # Returns R (B,3,3) and t_vec (B,3) such that: Y ≈ (R @ X.T).T + t | |
| R, t_vec = roma.rigid_points_registration( | |
| batch_last_position, # Source (X) | |
| batch_current_position # Target (Y) | |
| ) | |
| # Apply transformation: Y = X @ R.T + t | |
| # R: (B,3,3), batch_last_position: (B,N,3) | |
| # X @ R.T: (B,N,3) @ (B,3,3).T = (B,N,3) | |
| transformed = torch.bmm(batch_last_position, R.transpose(1, 2)) # (B,N,3) | |
| transformed = transformed + t_vec.unsqueeze(1) # (B,N,3) + (B,1,3) | |
| return transformed | |
| # Compute centroids | |
| centroid_last = batch_last_position.mean(dim=1, keepdim=True) # (B, 1, 3) | |
| centroid_current = batch_current_position.mean(dim=1, keepdim=True) # (B, 1, 3) | |
| # Center the point clouds | |
| last_centered = batch_last_position - centroid_last # (B, N, 3) | |
| current_centered = batch_current_position - centroid_current # (B, N, 3) | |
| # Compute covariance matrix H = X^T @ Y | |
| # where X is source (last) and Y is target (current) | |
| H = torch.bmm(last_centered.transpose(1, 2), current_centered) # (B, 3, 3) | |
| # SVD decomposition with autocast disabled (SVD requires float32) | |
| with torch.amp.autocast('cuda', enabled=False): | |
| H_float32 = H.float() # Cast to float32 for SVD | |
| U, S, Vt = torch.linalg.svd(H_float32) # U: (B, 3, 3), S: (B, 3), Vt: (B, 3, 3) | |
| # Cast back to original dtype | |
| U, S, Vt = U.to(H.dtype), S.to(H.dtype), Vt.to(H.dtype) | |
| # Compute rotation matrix R = V @ U^T | |
| R = torch.bmm(Vt.transpose(1, 2), U.transpose(1, 2)) # (B, 3, 3) | |
| # Handle reflection case (ensure proper rotation with det(R) = 1) | |
| # Compute determinant in float32 to avoid half precision issues | |
| with torch.amp.autocast('cuda', enabled=False): | |
| R_float32 = R.float() | |
| det_R = torch.det(R_float32) # (B,) | |
| det_R = det_R.to(R.dtype) | |
| # Create a diagonal matrix for flipping if needed | |
| flip_matrix = torch.eye(3, device=R.device).unsqueeze(0).repeat(B, 1, 1) # (B, 3, 3) | |
| need_flip = det_R < 0 # (B,) | |
| if need_flip.any(): | |
| # Flip the sign of the last singular vector | |
| flip_matrix[need_flip, 2, 2] = -1 | |
| Vt_corrected = torch.bmm(flip_matrix, Vt) | |
| R = torch.bmm(Vt_corrected.transpose(1, 2), U.transpose(1, 2)) | |
| # Apply regularization if requested (helps with numerical stability) | |
| if self.regularize_rotation: | |
| # Re-orthogonalize using another SVD | |
| with torch.amp.autocast('cuda', enabled=False): | |
| R_float32 = R.float() | |
| U_reg, S_reg, Vt_reg = torch.linalg.svd(R_float32) | |
| R = torch.bmm(U_reg, Vt_reg).to(R.dtype) | |
| # Compute scale if using Umeyama algorithm | |
| if self.use_scale: | |
| # Compute optimal scale | |
| # s = trace(S) / trace(X^T @ X) | |
| numerator = S.sum(dim=1) # (B,) | |
| denominator = (last_centered ** 2).sum(dim=[1, 2]) # (B,) | |
| scale = (numerator / (denominator + 1e-8)).unsqueeze(1).unsqueeze(2) # (B, 1, 1) | |
| else: | |
| scale = torch.ones(B, 1, 1, device=R.device) | |
| # Compute translation | |
| # t = centroid_current - scale * R @ centroid_last | |
| rotated_centroid = torch.bmm(R, centroid_last.transpose(1, 2)).transpose(1, 2) # (B, 1, 3) | |
| translation = centroid_current - scale * rotated_centroid # (B, 1, 3) | |
| # Apply transformation: y = scale * R @ x + t | |
| transformed = scale * torch.bmm(batch_last_position, R.transpose(1, 2)) + translation | |
| return transformed | |
| class DifferentiableTransformLayer(TransformLayer): | |
| """ | |
| Learnable transform layer that predicts rotation (6D) and translation. | |
| Modes: | |
| - absolute: predict the full R/t mapping last -> current | |
| - residual: reserved for composing a residual with rigid alignment (not implemented) | |
| """ | |
| def __init__(self, use_scale=False, hidden_dim=256, use_6d_rotation=True, | |
| learnable_mode='absolute', return_params=True): | |
| super().__init__(use_scale=use_scale) | |
| self.use_6d_rotation = use_6d_rotation | |
| self.learnable_mode = learnable_mode # 'absolute' | 'residual' | |
| self.return_params = return_params | |
| # Shallow PointNet over all vertices for each body (last/current) | |
| # Encode centered coordinates with 1x1 convs + BN + ReLU, global max-pool | |
| # last_feat, curr_feat each -> (B, 256) | |
| self.pn_conv1 = nn.Conv1d(3, 64, 1) | |
| self.pn_bn1 = nn.BatchNorm1d(64) | |
| self.pn_conv2 = nn.Conv1d(64, 128, 1) | |
| self.pn_bn2 = nn.BatchNorm1d(128) | |
| self.pn_conv3 = nn.Conv1d(128, 256, 1) | |
| self.pn_bn3 = nn.BatchNorm1d(256) | |
| # Head MLP maps fused global features (+ centroids for translation prior) to rot6d + t | |
| # Input dim: 256 (last) + 256 (current) + 6 (centroid_last/current) | |
| head_in_dim = 256 + 256 + 6 | |
| out_dim = 9 # 6D rotation + 3D translation | |
| self.head_mlp = nn.Sequential( | |
| nn.Linear(head_in_dim, hidden_dim), | |
| nn.ReLU(), | |
| nn.Linear(hidden_dim, hidden_dim), | |
| nn.ReLU(), | |
| nn.Linear(hidden_dim, out_dim) | |
| ) | |
| def rotation_6d_to_matrix(rot6d: torch.Tensor) -> torch.Tensor: | |
| a1 = rot6d[:, 0:3] | |
| a2 = rot6d[:, 3:6] | |
| b1 = F.normalize(a1, dim=1) | |
| a2_proj = a2 - (b1 * (a2 * b1).sum(dim=1, keepdim=True)) | |
| b2 = F.normalize(a2_proj, dim=1) | |
| b3 = torch.cross(b1, b2, dim=1) | |
| return torch.stack([b1, b2, b3], dim=2) # (B,3,3) | |
| def forward(self, batch_last_position, batch_current_position): | |
| B, N, D = batch_last_position.shape | |
| assert D == 3, f"Expected 3D points, got {D}D" | |
| # Compute centroids for stability and translation prior | |
| centroid_last = batch_last_position.mean(dim=1) # (B,3) | |
| centroid_current = batch_current_position.mean(dim=1) # (B,3) | |
| # Center the point clouds (invariant to global translation) | |
| last_centered = batch_last_position - centroid_last.unsqueeze(1) # (B,N,3) | |
| current_centered = batch_current_position - centroid_current.unsqueeze(1) # (B,N,3) | |
| # PointNet feature extraction for both clouds | |
| def encode_points(x_centered: torch.Tensor) -> torch.Tensor: | |
| # x_centered: (B,N,3) -> (B,3,N) | |
| x = x_centered.transpose(1, 2) | |
| x = F.relu(self.pn_bn1(self.pn_conv1(x))) | |
| x = F.relu(self.pn_bn2(self.pn_conv2(x))) | |
| x = F.relu(self.pn_bn3(self.pn_conv3(x))) | |
| # Global max pooling over points -> (B,256) | |
| x = torch.amax(x, dim=2) | |
| return x | |
| last_feat = encode_points(last_centered) | |
| curr_feat = encode_points(current_centered) | |
| # Fuse global features and centroids | |
| fused = torch.cat([last_feat, curr_feat, centroid_last, centroid_current], dim=1) # (B, 518) | |
| # Predict rotation (6D) and translation (3D) | |
| pred = self.head_mlp(fused) # (B,9) | |
| rot6d = pred[:, :6] | |
| t = pred[:, 6:9].unsqueeze(1) # (B,1,3) | |
| R_pred = self.rotation_6d_to_matrix(rot6d) # (B,3,3) | |
| if self.learnable_mode == 'absolute': | |
| corrected = torch.bmm(batch_last_position, R_pred.transpose(1, 2)) + t | |
| return corrected, { 'rot6d': rot6d, 'trans': t.squeeze(1) } | |
| class RobustTransformLayer(TransformLayer): | |
| """ | |
| Robust variant that downweights outliers via iteratively reweighted | |
| Kabsch alignment (weights decay exponentially with residual distance). | |
| """ | |
| def __init__(self, use_scale=False, num_iterations=3, outlier_threshold=0.1): | |
| super().__init__(use_scale=use_scale) | |
| self.num_iterations = num_iterations | |
| self.outlier_threshold = outlier_threshold | |
| def forward(self, batch_last_position, batch_current_position): | |
| B, N, _ = batch_last_position.shape | |
| # Initialize weights (all points equally weighted) | |
| weights = torch.ones(B, N, 1, device=batch_last_position.device) | |
| transformed = batch_last_position.clone() | |
| for iteration in range(self.num_iterations): | |
| # Weighted centroids | |
| weights_norm = weights / (weights.sum(dim=1, keepdim=True) + 1e-8) | |
| centroid_last = (transformed * weights_norm).sum(dim=1, keepdim=True) | |
| centroid_current = (batch_current_position * weights_norm).sum(dim=1, keepdim=True) | |
| # Center with weighted centroids | |
| last_centered = transformed - centroid_last | |
| current_centered = batch_current_position - centroid_current | |
| # Weighted covariance | |
| last_weighted = last_centered * weights_norm.sqrt() | |
| current_weighted = current_centered * weights_norm.sqrt() | |
| H = torch.bmm(last_weighted.transpose(1, 2), current_weighted) | |
| # SVD and compute transformation with autocast disabled | |
| with torch.amp.autocast('cuda', enabled=False): | |
| H_float32 = H.float() | |
| U, S, Vt = torch.linalg.svd(H_float32) | |
| U, S, Vt = U.to(H.dtype), S.to(H.dtype), Vt.to(H.dtype) | |
| R = torch.bmm(Vt.transpose(1, 2), U.transpose(1, 2)) | |
| # Fix reflection | |
| # Compute determinant in float32 to avoid half precision issues | |
| with torch.amp.autocast('cuda', enabled=False): | |
| R_float32 = R.float() | |
| det_R = torch.det(R_float32) | |
| det_R = det_R.to(R.dtype) | |
| flip_matrix = torch.eye(3, device=R.device).unsqueeze(0).repeat(B, 1, 1) | |
| need_flip = det_R < 0 | |
| if need_flip.any(): | |
| flip_matrix[need_flip, 2, 2] = -1 | |
| Vt_corrected = torch.bmm(flip_matrix, Vt) | |
| R = torch.bmm(Vt_corrected.transpose(1, 2), U.transpose(1, 2)) | |
| # Translation | |
| translation = centroid_current - torch.bmm(R, centroid_last.transpose(1, 2)).transpose(1, 2) | |
| # Apply transformation | |
| transformed = torch.bmm(batch_last_position, R.transpose(1, 2)) + translation | |
| # Update weights based on distances (downweight outliers) | |
| if iteration < self.num_iterations - 1: | |
| distances = torch.norm(transformed - batch_current_position, dim=2, keepdim=True) | |
| weights = torch.exp(-distances / self.outlier_threshold) | |
| return transformed | |
| class RoMaTransformLayer(nn.Module): | |
| """ | |
| Differentiable Kabsch registration using RoMa library. | |
| Key difference from TransformLayer: | |
| - Fully differentiable (gradients flow through R and t) | |
| - No need for STE (Straight-Through Estimator) | |
| - Uses roma.rigid_points_registration() which handles SVD gradient | |
| Note: Requires 'roma' library to be installed. | |
| """ | |
| def __init__(self): | |
| super().__init__() | |
| if not ROMA_AVAILABLE: | |
| raise RuntimeError( | |
| "RoMa library not available. Install with: pip install roma\n" | |
| "RoMa is required for differentiable Kabsch registration." | |
| ) | |
| def forward(self, batch_last_position, batch_current_position): | |
| """ | |
| Register batch_last_position to batch_current_position using RoMa. | |
| Args: | |
| batch_last_position: (B, N, 3) - source positions to be transformed | |
| batch_current_position: (B, N, 3) - target positions to align to | |
| Returns: | |
| transformed_positions: (B, N, 3) - transformed positions (differentiable) | |
| """ | |
| B, N, D = batch_last_position.shape | |
| assert D == 3, f"Expected 3D points, got {D}D" | |
| assert batch_last_position.shape == batch_current_position.shape | |
| # RoMa's rigid_points_registration returns (R, t_vec) | |
| # R: (B, 3, 3) rotation matrix | |
| # t_vec: (B, 3) translation vector | |
| # Formula: y = R @ x + t | |
| R, t_vec = roma.rigid_points_registration( | |
| batch_last_position, # source | |
| batch_current_position # target | |
| ) | |
| # Apply transformation: y = x @ R^T + t | |
| # (B, N, 3) @ (B, 3, 3) -> (B, N, 3) | |
| transformed = torch.bmm(batch_last_position, R.transpose(1, 2)) + t_vec.unsqueeze(1) | |
| return transformed | |
| # Convenience function for creating the layer | |
| def create_transform_layer(method='kabsch', if_roma=False, **kwargs): | |
| """ | |
| Factory function to create transform layer. | |
| Args: | |
| method: Only 'kabsch' is supported | |
| if_roma: If True, use RoMa's differentiable Kabsch; otherwise use standard SVD Kabsch | |
| **kwargs: Additional arguments (e.g., regularize_rotation) | |
| Returns: | |
| TransformLayer instance with Kabsch algorithm | |
| """ | |
| if method != 'kabsch': | |
| raise ValueError(f"Only method='kabsch' is supported, got '{method}'") | |
| return TransformLayer(use_scale=False, if_roma=if_roma, **kwargs) |