CPD-Pytorch

https://github.com/yuliangxiaoylx/cpd-pytorch
Admin
CPD-Pytorch is a PyTorch implementation of the Coherent Point Drift algorithm for point cloud...
Tokens:5,238
Snippets:30
Trust Score:8.1
License:Apache-2.0
Update:5 months ago
Show doc for...
Tokens:
Raw
### Install CPD-Pytorch

Source: https://github.com/yuliangxiaoylx/cpd-pytorch/blob/main/README.md

Clone the repository and install the library in editable mode. Ensure you are in the cloned directory.

```bash
git clone https://github.com/mikami520/CPD-Pytorch.git
cd CPD-Pytorch
pip install -e .
```

--------------------------------

### Rigid Point Cloud Registration Example

Source: https://github.com/yuliangxiaoylx/cpd-pytorch/blob/main/README.md

Example demonstrating rigid registration using CPD-Pytorch. Loads target and source point clouds, sets up a callback for visualization, and performs registration. Uses CUDA if available, otherwise falls back to CPU.

```python
from functools import partial
import matplotlib.pyplot as plt
from torchcpd import RigidRegistration
import numpy as np
import torch as th

device = 'cuda:0' if th.cuda.is_available() else 'cpu'
X = np.loadtxt('data/bunny_target.txt')
# synthetic data, equaivalent to X + 1
Y = np.loadtxt('data/bunny_source.txt')

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

callback = partial(visualize, ax=ax, fig=fig, save_fig=False)

reg = RigidRegistration(**{'X': X, 'Y': Y, 'device': device})
reg.register(callback)
plt.show()
```

--------------------------------

### Rigid Registration with Callback

Source: https://context7.com/yuliangxiaoylx/cpd-pytorch/llms.txt

Performs rigid point cloud registration using the EM algorithm. Supports GPU acceleration and custom callbacks for tracking convergence. Initialize sigma2 automatically or provide a value.

```python
import numpy as np
import torch as th
from torchcpd import RigidRegistration

X = np.loadtxt('data/bunny_target.txt')
Y = np.loadtxt('data/bunny_source.txt')
device = 'cuda:0' if th.cuda.is_available() else 'cpu'

# Track convergence manually via callback
history = []
def log_callback(iteration, error, X, Y):
    history.append({'iteration': iteration, 'q': float(error[0])})

reg = RigidRegistration(
    X=X, Y=Y, device=device,
    max_iterations=200,   # stop after at most 200 EM steps
    tolerance=1e-5,       # stop when |q_t - q_{t-1}| < 1e-5
    w=0.1,                # 10% uniform outlier weight
    sigma2=None,          # auto-initialize variance from data
    scale=True            # allow isotropic scaling
)
TY, params = reg.register(log_callback)

print(f"Converged in {len(history)} iterations")
print(f"Final objective Q: {history[-1]['q']:.6f}")
# TY: aligned source point cloud (torch.Tensor on device)
# params: (s, R, t) for rigid; (B, t) for affine; (G, W) for deformable
```

--------------------------------

### Rigid Point Cloud Registration with Visualization

Source: https://context7.com/yuliangxiaoylx/cpd-pytorch/llms.txt

Performs rigid registration on 3D point clouds using RigidRegistration. Includes an optional callback for per-iteration visualization. Requires loading point clouds and setting up a matplotlib figure for plotting.

```python
from functools import partial
import numpy as np
import torch as th
import matplotlib.pyplot as plt
from torchcpd import RigidRegistration

device = 'cuda:0' if th.cuda.is_available() else 'cpu'

# Load 3D point clouds (NxD numpy arrays)
X = np.loadtxt('data/bunny_target.txt')  # target
Y = np.loadtxt('data/bunny_source.txt')  # source (to be aligned)

# Optional: per-iteration visualization callback
def visualize(iteration, error, X, Y, ax, fig):
    plt.cla()
    X_np, Y_np = X.detach().cpu().numpy(), Y.detach().cpu().numpy()
    ax.scatter(X_np[:, 0], X_np[:, 1], X_np[:, 2], color='red', label='Target')
    ax.scatter(Y_np[:, 0], Y_np[:, 1], Y_np[:, 2], color='blue', label='Source')
    ax.text2D(0.87, 0.92, f'Iter: {iteration}  Q: {error[0]:.4f}',
              transform=ax.transAxes)
    ax.legend(); plt.draw(); plt.pause(0.001)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
callback = partial(visualize, ax=ax, fig=fig)

# Construct and run registration
reg = RigidRegistration(X=X, Y=Y, device=device,
                        max_iterations=150, tolerance=1e-5, w=0.0, scale=True)
TY, (s, R, t) = reg.register(callback)

# TY: transformed source (torch.Tensor, MxD)
# s: scale factor, R: DxD rotation matrix, t: 1xD translation vector
print(f"Scale: {s.item():.4f}")
print(f"Rotation:\n{R.detach().cpu().numpy()}")
print(f"Translation: {t.detach().cpu().numpy()}")

# Apply transform to new unseen points
new_points = th.tensor(np.loadtxt('data/bunny_source.txt'), dtype=th.float64).to(device)
transformed = reg.transform_point_cloud(new_points)  # returns torch.Tensor
plt.show()
```

--------------------------------

### initialize_sigma2

Source: https://context7.com/yuliangxiaoylx/cpd-pytorch/llms.txt

Utility function that computes the initial variance for the GMM from the mean squared pairwise distances between all source and target points. This is called automatically during EMRegistration.__init__ when sigma2=None is passed.

```APIDOC
## initialize_sigma2

### Description
Utility function that computes the initial variance for the GMM from the mean squared pairwise distances between all source and target points. Called automatically during `EMRegistration.__init__` when `sigma2=None` is passed.

### Method Signature
```python
initialize_sigma2(X, Y)
```

### Parameters
- **X** (torch.Tensor): The target point cloud.
- **Y** (torch.Tensor): The source point cloud.

### Returns
- **sigma2_init** (float): The initial variance.
```

--------------------------------

### RigidRegistration

Source: https://context7.com/yuliangxiaoylx/cpd-pytorch/llms.txt

Performs rigid point cloud registration (rotation, translation, optional uniform scaling) for 2D and 3D point clouds. It returns the scale factor, rotation matrix, and translation vector. The `register` method can accept a callback for per-iteration updates, and the `transform_point_cloud` method applies the learned transformation to new points.

```APIDOC
## RigidRegistration

### Description
Performs rigid (rotation + translation + optional uniform scaling) point cloud registration. Supports 2D and 3D point clouds only. The rotation matrix is computed via SVD as per Lemma 1 of the CPD paper. Returns the scale `s`, rotation matrix `R`, and translation vector `t`.

### Method Signature
`RigidRegistration(X, Y, device, max_iterations, tolerance, w, scale)`

### Parameters
- **X** (numpy.ndarray or torch.Tensor) - Target point cloud.
- **Y** (numpy.ndarray or torch.Tensor) - Source point cloud to be aligned.
- **device** (str) - Device to perform computation on ('cuda:0' or 'cpu').
- **max_iterations** (int) - Maximum number of EM iterations.
- **tolerance** (float) - Convergence tolerance for the EM algorithm.
- **w** (float) - Weighting factor for the CPD objective function.
- **scale** (bool) - Whether to estimate a uniform scale factor.

### Method
`register(callback)`

### Parameters
- **callback** (function, optional) - A function to be called after each iteration, receiving iteration number, error, target point cloud, and source point cloud.

### Returns
- **TY** (torch.Tensor) - Transformed source point cloud.
- **(s, R, t)** (tuple) - A tuple containing:
    - **s** (torch.Tensor) - Scale factor (if `scale` is True).
    - **R** (torch.Tensor) - DxD rotation matrix.
    - **t** (torch.Tensor) - 1xD translation vector.

### Method
`transform_point_cloud(points)`

### Parameters
- **points** (torch.Tensor) - New point cloud to transform.

### Returns
- **transformed_points** (torch.Tensor) - The transformed point cloud.
```

--------------------------------

### AffineRegistration

Source: https://context7.com/yuliangxiaoylx/cpd-pytorch/llms.txt

Performs affine point cloud registration, estimating a full DxD affine transformation matrix `B` and translation vector `t`. This method handles shearing, scaling, and rotation simultaneously, making it suitable for point clouds that have undergone general linear transformations.

```APIDOC
## AffineRegistration

### Description
Performs affine point cloud registration, estimating a full DxD affine transformation matrix `B` and translation vector `t`. Handles shearing, scaling, and rotation simultaneously. Well-suited for cases where the point cloud undergoes a general linear transformation.

### Method Signature
`AffineRegistration(X, Y, device)`

### Parameters
- **X** (numpy.ndarray or torch.Tensor) - Target point cloud.
- **Y** (numpy.ndarray or torch.Tensor) - Source point cloud to be aligned.
- **device** (str) - Device to perform computation on ('cuda:0' or 'cpu').

### Method
`register(callback)`

### Parameters
- **callback** (function, optional) - A function to be called after each iteration, receiving iteration number, error, target point cloud, and source point cloud.

### Returns
- **TY** (torch.Tensor) - Transformed source point cloud.
- **(B, t)** (tuple) - A tuple containing:
    - **B** (torch.Tensor) - DxD affine transformation matrix.
    - **t** (torch.Tensor) - 1xD translation vector.
```

--------------------------------

### Initialize Sigma2 for GMM

Source: https://context7.com/yuliangxiaoylx/cpd-pytorch/llms.txt

Computes the initial variance for GMM from mean squared pairwise distances. This is used by EMRegistration when sigma2 is not explicitly provided. Can be used to override default initialization.

```python
import numpy as np
import torch as th
from torchcpd.emregistration import initialize_sigma2

X = th.tensor(np.loadtxt('data/fish_target.txt'), dtype=th.float64)
Y = th.tensor(np.loadtxt('data/fish_source.txt'), dtype=th.float64)

sigma2_init = initialize_sigma2(X, Y)
print(f"Initial sigma2: {sigma2_init.item():.6f}")
# sigma2 = sum_{n,m} ||x_n - y_m||^2 / (D * M * N)
# Use this to inspect or override the default initialization:
from torchcpd import DeformableRegistration
reg = DeformableRegistration(X=X.numpy(), Y=Y.numpy(),
                             device='cpu',
                             sigma2=sigma2_init.item() * 2.0)  # start with larger variance
```

--------------------------------

### EMRegistration.register

Source: https://context7.com/yuliangxiaoylx/cpd-pytorch/llms.txt

The core EM loop for registration. It alternates between the E-step (computing soft assignment probabilities P) and the M-step (updating transform parameters and variance sigma2). The loop terminates when max_iterations is reached or the objective change falls below tolerance.

```APIDOC
## EMRegistration.register

### Description
The core EM loop shared by all registration classes. Alternates between the E-step (computing soft assignment probabilities `P`) and the M-step (updating transform parameters and variance `sigma2`). Terminates when `max_iterations` is reached or the objective change falls below `tolerance`.

### Method Signature
```python
register(callback=None)
```

### Parameters
- **callback** (function, optional): A function to be called at each iteration with `(iteration, error, X, Y)`.

### Returns
- **TY** (torch.Tensor): The aligned source point cloud.
- **params** (tuple): Transformation parameters (e.g., (s, R, t) for rigid registration).
```

--------------------------------

### Affine Point Cloud Registration

Source: https://context7.com/yuliangxiaoylx/cpd-pytorch/llms.txt

Performs affine registration on 2D point clouds, estimating an affine transformation matrix and translation vector. Useful when point clouds undergo shearing, scaling, and rotation. Requires setting up a matplotlib figure for plotting.

```python
from functools import partial
import numpy as np
import torch as th
import matplotlib.pyplot as plt
from torchcpd import AffineRegistration

device = 'cuda:0' if th.cuda.is_available() else 'cpu'

# Construct a synthetically transformed 2D fish point cloud
X = np.loadtxt('data/fish_target.txt')  # target (Nx2)
theta = np.pi / 6.0
R_mat = np.array([[np.cos(theta), -np.sin(theta)],
                  [np.sin(theta),  np.cos(theta)]])
shear = np.array([[1, 0.5], [0, 1]])
A_true = np.dot(R_mat, shear)
t_true = np.array([0.5, 1.0])
Y = np.dot(X, A_true) + t_true  # source: affinely transformed target

def visualize(iteration, error, X, Y, ax):
    plt.cla()
    X_np, Y_np = X.detach().cpu().numpy(), Y.detach().cpu().numpy()
    ax.scatter(X_np[:, 0], X_np[:, 1], color='red', label='Target')
    ax.scatter(Y_np[:, 0], Y_np[:, 1], color='blue', label='Source')
    ax.text(0.87, 0.92, f'Iter: {iteration}  Q: {error[0]:.4f}',
            transform=ax.transAxes); ax.legend(); plt.draw(); plt.pause(0.001)

fig = plt.figure()
fig.add_axes([0, 0, 1, 1])
callback = partial(visualize, ax=fig.axes[0])

reg = AffineRegistration(X=X, Y=Y, device=device)
TY, (B, t) = reg.register(callback)

# B: DxD affine matrix, t: 1xD translation
print("Affine matrix B:", B.detach().cpu().numpy())
print("Translation t:", t.detach().cpu().numpy())

```

--------------------------------

### Perform Deformable Registration

Source: https://context7.com/yuliangxiaoylx/cpd-pytorch/llms.txt

Use DeformableRegistration for non-rigid point cloud registration. Set low_rank=False for standard registration or low_rank=True for a low-rank approximation suitable for large point clouds.

```python
from functools import partial
import numpy as np
import torch as th
import matplotlib.pyplot as plt
from torchcpd import DeformableRegistration

device = 'cuda:0' if th.cuda.is_available() else 'cpu'

X = np.loadtxt('data/fish_target.txt')
Y = np.loadtxt('data/fish_source.txt')

def visualize(iteration, error, X, Y, ax):
    plt.cla()
    X_np, Y_np = X.detach().cpu().numpy(), Y.detach().cpu().numpy()
    ax.scatter(X_np[:, 0], X_np[:, 1], color='red', label='Target')
    ax.scatter(Y_np[:, 0], Y_np[:, 1], color='blue', label='Source')
    ax.text(0.87, 0.92, f'Iter: {iteration}', transform=ax.transAxes)
    ax.legend(); plt.draw(); plt.pause(0.001)

fig = plt.figure()
fig.add_axes([0, 0, 1, 1])
callback = partial(visualize, ax=fig.axes[0])

# Standard (full-rank) deformable registration
reg = DeformableRegistration(X=X, Y=Y, device=device,
                             alpha=2, beta=2,
                             low_rank=False)
TY, (G, W) = reg.register(callback)

# G: MxM Gaussian kernel matrix, W: MxD deformation weights
# Deformation field: G @ W added to original source positions
print("Kernel G shape:", G.shape)   # (M, M)
print("Weights W shape:", W.shape)  # (M, D)

# --- Low-rank variant for large point clouds ---
# Construct 3D data by stacking 2D fish slices
fish_target = np.loadtxt('data/fish_target.txt')
X3 = np.vstack([np.hstack([fish_target, np.zeros((len(fish_target), 1))]),
                np.hstack([fish_target, np.ones((len(fish_target), 1))])])
fish_source = np.loadtxt('data/fish_source.txt')
Y3 = np.vstack([np.hstack([fish_source, np.zeros((len(fish_source), 1))]),
                np.hstack([fish_source, np.ones((len(fish_source), 1))])])

reg_lr = DeformableRegistration(X=X3, Y=Y3, device=device,
                                low_rank=True, num_eig=100)
TY_lr, (G_lr, W_lr) = reg_lr.register()
plt.show()
```

--------------------------------

### Gaussian Kernel Calculation

Source: https://context7.com/yuliangxiaoylx/cpd-pytorch/llms.txt

Calculates self-kernel and cross-kernel using the Gaussian kernel function. Ensure X and Y are appropriately defined before use.

```python
G_self = gaussian_kernel(X, beta=2.0)
print("Self-kernel shape:", G_self.shape)   # (N, N)
print("Diagonal (should be 1.0):", G_self.diagonal()[:5])

# Cross-kernel (MxN): between two different point sets
G_cross = gaussian_kernel(X, beta=2.0, Y=Y)
print("Cross-kernel shape:", G_cross.shape)  # (N, M)
```

--------------------------------

### Perform Constrained Deformable Registration

Source: https://context7.com/yuliangxiaoylx/cpd-pytorch/llms.txt

Use ConstrainedDeformableRegistration to incorporate hard correspondence constraints. The e_alpha parameter controls confidence in these correspondences; smaller values indicate higher confidence.

```python
from functools import partial
import numpy as np
import torch as th
import matplotlib.pyplot as plt
from torchcpd import ConstrainedDeformableRegistration

device = 'cuda:0' if th.cuda.is_available() else 'cpu'

X = np.loadtxt('data/fish_target.txt')
Y = np.loadtxt('data/fish_source.txt')

# Simulate partial overlap — only use first 61 target points
X = X[:61]

# Known correspondence pairs (indices into source and target arrays)
IDs = [1, 10, 20, 30]
src_id = np.int32(IDs)
tgt_id = np.int32(IDs)

def visualize(iteration, error, X, Y, ax):
    plt.cla()
    all_X = np.arange(X.shape[0])
    free_X = np.delete(all_X, IDs)
    free_Y = np.delete(np.arange(Y.shape[0]), IDs)
    X_np, Y_np = X.detach().cpu().numpy(), Y.detach().cpu().numpy()
    ax.scatter(X_np[free_X, 0], X_np[free_X, 1], color='red', label='Target')
    ax.scatter(Y_np[free_Y, 0], Y_np[free_Y, 1], color='blue', label='Source')
    # Highlight constrained landmarks
    ax.scatter(X_np[IDs, 0], X_np[IDs, 1], color='blue',
               s=100, facecolors='none', label='Target Constrained')
    ax.scatter(Y_np[IDs, 0], Y_np[IDs, 1], color='green',
               s=100, marker=(5, 1), label='Source Constrained')
    ax.text(0.87, 0.92, f'Iter: {iteration}', transform=ax.transAxes)
    ax.legend(); plt.draw(); plt.pause(0.001)

fig = plt.figure()
fig.add_axes([0, 0, 1, 1])
callback = partial(visualize, ax=fig.axes[0])

# e_alpha: correspondence reliability (1e-8 = very reliable, 1 = unreliable)
reg = ConstrainedDeformableRegistration(
    X=X, Y=Y, device=device,
    e_alpha=1e-8,
    source_id=src_id,
    target_id=tgt_id,
    alpha=2, beta=2
)
TY, (G, W) = reg.register(callback)

print("Aligned source shape:", TY.shape)  # (M, D)
plt.show()
```

--------------------------------

### Compute Gaussian Kernel

Source: https://context7.com/yuliangxiaoylx/cpd-pytorch/llms.txt

The gaussian_kernel function computes a spatial coherence kernel between point sets. It is used internally by registration modules but can also be called directly for custom kernel-based computations.

```python
import numpy as np
import torch as th
from torchcpd import gaussian_kernel

X = th.tensor(np.loadtxt('data/fish_target.txt'), dtype=th.float64)
Y = th.tensor(np.loadtxt('data/fish_source.txt'), dtype=th.float64)


```
Tokens:
Raw