point-affinity-transformer 0.2.0

Point Affinity Transformers

Installation

pip install point-affinity-transformer

The package is imported as pat. Install a PyTorch build matching your CUDA version first (pytorch.org).

Getting Started

import pat
from pat.io import read_swcs

points, gt = read_swcs(["a.swc", "b.swc"])          # (N, 3) coordinates in µm
points = points / 659.88                            # your normalization: the same as for the training data
model = pat.load_pretrained("my_model", device="cuda")  # directory from pat train or pat.save_pretrained

seg = pat.segment(model, points, threshold=0.8)     # one label per point, -1 = noise
print(seg.num_neurons, pat.metrics.voi(seg.labels, gt))

enc = pat.encode(model, points)                     # encode once ...
A = enc.affinity(queries, others)                   # ... then query affinities between any points

Point Clouds and Features

Point clouds of shape (B, N, 3) with D extra features are passed to the model as (B, N, 3 + D). Training only needs a neuron label per point.

import torch
import torch.nn.functional as F
import pat

B, N, D = 8, 1024, 4                           # random toy data: replace with your own
xyz = torch.rand(B, N, 3) * 2 - 1              # coordinates, already normalized to roughly [-1, 1]
features = torch.randn(B, N, D)                # D features per point
labels = torch.randint(0, 3, (B, N))           # neuron id of every point (training targets)
pc = torch.cat([xyz, features], dim=-1)        # (B, N, 3 + D)

model = pat.create_model("ae_d256_m64", depth=4, feature_dim=D)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

# Training: sample random point pairs, target = 1 if both lie on the same neuron
def gather(x, idx):                            # x: (B, N, C), idx: (B, Q) -> (B, Q, C)
    return x.gather(1, idx[..., None].expand(-1, -1, x.shape[-1]))

model.train()
for step in range(100):
    i, j = torch.randint(0, N, (2, B, 2048))   # 2048 query pairs per cloud
    target = (labels.gather(1, i) == labels.gather(1, j)).float()
    logits = model(pc, None, gather(pc, i), gather(pc, j), logits=True)  # mask=None: no padding
    loss = F.binary_cross_entropy_with_logits(logits, target)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Inference
model.eval()
with torch.no_grad():
    latents = model.encode(pc)                             # (B, num_latents, dim), once per cloud
    affinity = model.decode(latents, pc[:, :5], pc[:, 5:10])  # (B, 5): any query pairs, with features

seg = pat.segment(model, pc[0])                # cluster one cloud: labels (N,), -1 = noise
pat.save_pretrained(model, "my_model")         # reload with pat.load_pretrained("my_model")

Coordinates should lie within roughly [-1, 1]. Clouds of different sizes can be padded with pat.preprocess.pad_batch, which returns the mask to pass instead of None. A trained 3D model can be extended to D features with pat.expand_features(model, D) and then fine-tuned the same way. To train from SWC files, run pat train config.yaml (examples in research/configs), optionally with data.features: [radius] or, from Python, pat.train.train(cfg, feature_fn=...).

See the quickstart notebook for an end-to-end example, and the documentation (pip install -e ".[docs]" && mkdocs serve) for the API reference.

Citation

@techreport{troidlgnsr2024,
  title = {Global Neuron Shape Reasoning with Point Affinity Transformers},
  author = {Troidl, Jakob and Knittel, Johannes and Li, Wanhua and Zhan, Fengnang and Pfister*, Hanspeter and Turaga*, Srinivas},
  journal = {bioRxiv},
  year = {2024},
  publisher = {Cold Spring Harbor Laboratory},
  keywords = {preprint}
}

Acknowledgements

We acknowledge NSF grants CRCNS-2309041, NCS-FO-2124179, and NIH grant R01HD104969. We also thank the HHMI Janelia Visiting Scientist Program and the Harvard Data Science Initiative Postdoctoral Fellowship for their support. The code is partially based on 3DShape2VecSet by Zhang et. al.