point-affinity-transformer 0.1.0__py3-none-any.whl

pat/__init__.py

@@ -0,0 +1,30 @@
+"""Point Affinity Transformer: global neuron shape reasoning on point clouds."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from . import cluster, metrics
+from .inference import EncodedCloud, Segmentation, encode, segment
+from .model import PATConfig, PointAffinityTransformer, expand_features
+from .presets import PRESETS, create_model
+from .pretrained import load_pretrained, save_pretrained
+
+try:
+    __version__ = version("point-affinity-transformer")  # defined in pyproject.toml
+except PackageNotFoundError:  # running from a source tree without installing
+    __version__ = "unknown"
+
+__all__ = [
+    "PATConfig",
+    "PointAffinityTransformer",
+    "EncodedCloud",
+    "Segmentation",
+    "encode",
+    "segment",
+    "create_model",
+    "PRESETS",
+    "load_pretrained",
+    "save_pretrained",
+    "expand_features",
+    "cluster",
+    "metrics",
+]

pat/__main__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+main()

pat/cli.py

@@ -0,0 +1,85 @@
+"""Command line interface: ``pat train | eval | convert``."""
+
+from __future__ import annotations
+
+import argparse
+
+
+def _train(args):
+    from .train import load_config, train
+
+    train(load_config(args.config, args.set))
+
+
+def _eval(args):
+    import json
+
+    import torch
+
+    from .pretrained import load_pretrained
+    from .train import MergedNeuronDataset, SkeletonStore, TensorDirDataset, evaluate, format_results, load_config
+    from .train.data import read_root_ids
+
+    cfg = load_config(args.config, args.set)
+    torch.manual_seed(cfg.seed)
+    model = load_pretrained(args.model, device=args.device or cfg.device)
+    if args.items:
+        dataset = TensorDirDataset(args.items)
+    else:
+        store = SkeletonStore(cfg.data.swc_dir, cfg.data.units, cfg.data.dust_size, cfg.data.cache_dir,
+                              features=cfg.data.features)
+        ids = read_root_ids(cfg.data.eval_neuron_ids or cfg.data.neuron_ids)
+        dataset = MergedNeuronDataset(store, ids, cfg.data, train=False)
+    thresholds = [float(t) for t in args.thresholds.split(",")]
+    results = evaluate(model, dataset, thresholds, min_cluster_size=args.min_cluster_size,
+                       symmetrize=args.symmetrize, limit=args.limit)
+    print(format_results(results))
+    if args.out:
+        with open(args.out, "w") as f:
+            json.dump(results, f, indent=2)
+
+
+def _convert(args):
+    import torch
+
+    from .pretrained import load_legacy_checkpoint, save_pretrained
+
+    model = load_legacy_checkpoint(args.checkpoint, preset=args.preset)
+    print(save_pretrained(model, args.out, dtype=torch.float16 if args.fp16 else None))
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(prog="pat", description="Point Affinity Transformer")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    p = sub.add_parser("train", help="train or fine-tune a model")
+    p.add_argument("config", help="YAML training config")
+    p.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE", help="override config entries")
+    p.set_defaults(func=_train)
+
+    p = sub.add_parser("eval", help="evaluate segmentation on merged clouds")
+    p.add_argument("config", nargs="?", help="YAML config describing the data (not needed with --items)")
+    p.add_argument("--model", required=True, help="model directory (from pat train or save_pretrained)")
+    p.add_argument("--items", help="directory of pre-generated items (pc_*.pt, labels_*.pt, mask_*.pt)")
+    p.add_argument("--thresholds", default="0.8")
+    p.add_argument("--min-cluster-size", type=int, default=30)
+    p.add_argument("--symmetrize", action="store_true", help="symmetrize affinities before clustering")
+    p.add_argument("--limit", type=int, help="evaluate only the first N items")
+    p.add_argument("--device")
+    p.add_argument("--out", help="write results as JSON")
+    p.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE")
+    p.set_defaults(func=_eval)
+
+    p = sub.add_parser("convert", help="convert a research .pth checkpoint into a model directory")
+    p.add_argument("checkpoint")
+    p.add_argument("out")
+    p.add_argument("--preset", help="model preset; overrides the one stored in the checkpoint")
+    p.add_argument("--fp16", action="store_true", help="store weights in float16")
+    p.set_defaults(func=_convert)
+
+    args = parser.parse_args(argv)
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

pat/cluster.py

@@ -0,0 +1,97 @@
+"""Clustering of affinity matrices into neuron labels.
+
+The label ``-1`` marks noise: DBSCAN outliers and, after
+[`remove_small_clusters`][pat.cluster.remove_small_clusters], points of clusters
+smaller than ``min_cluster_size``.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+from sklearn import cluster as _skcluster
+
+NOISE = -1
+
+
+def to_numpy(x) -> np.ndarray:
+    """Convert a tensor or array-like to a NumPy array."""
+    return x.detach().cpu().numpy() if torch.is_tensor(x) else np.asarray(x)
+
+
+def num_clusters(labels) -> int:
+    """Number of clusters in ``labels`` numbered ``0..K-1``, ignoring noise (``-1``)."""
+    labels = torch.as_tensor(labels)
+    return int(labels.max().item()) + 1 if (labels >= 0).any() else 0
+
+
+def agglomerative(distance, threshold: float = 0.8, linkage: str = "average") -> torch.Tensor:
+    """Agglomerative clustering of a precomputed ``(N, N)`` distance matrix.
+
+    Clusters are merged while their linkage distance is below ``threshold``.
+    """
+    d = to_numpy(distance)
+    model = _skcluster.AgglomerativeClustering(
+        n_clusters=None, metric="precomputed", distance_threshold=threshold, linkage=linkage
+    )
+    return torch.as_tensor(model.fit_predict(d), dtype=torch.int64)
+
+
+def dbscan(distance, eps: float = 0.2, min_samples: int = 5) -> torch.Tensor:
+    """DBSCAN on a precomputed distance matrix. Noise points get label ``-1``."""
+    d = to_numpy(distance)
+    labels = _skcluster.DBSCAN(eps=eps, min_samples=min_samples, metric="precomputed").fit_predict(d)
+    return torch.as_tensor(labels, dtype=torch.int64)
+
+
+def remove_small_clusters(labels: torch.Tensor, min_size: int) -> torch.Tensor:
+    """Mark clusters with fewer than ``min_size`` points as noise and relabel.
+
+    The remaining clusters are numbered ``0..K-1``, largest first. Existing
+    ``-1`` labels stay noise.
+    """
+    labels = torch.as_tensor(labels, dtype=torch.int64)
+    out = torch.full_like(labels, NOISE)
+    valid = labels >= 0
+    if not valid.any():
+        return out
+    ids, inverse, counts = torch.unique(labels[valid], return_inverse=True, return_counts=True)
+    keep = counts >= min_size
+    order = torch.argsort(counts.masked_fill(~keep, -1), descending=True, stable=True)
+    new_id = torch.full_like(ids, NOISE)
+    new_id[order[: int(keep.sum())]] = torch.arange(int(keep.sum()))
+    out[valid] = new_id[inverse]
+    return out
+
+
+def cluster_affinity(
+    affinity,
+    threshold: float = 0.8,
+    min_cluster_size: int = 30,
+    method: str = "agglomerative",
+    symmetrize: bool = True,
+) -> torch.Tensor:
+    """Turn an ``(N, N)`` affinity matrix into one label per point.
+
+    Args:
+        affinity: Pairwise same-neuron probabilities.
+        threshold: Distance threshold on ``1 - affinity``. For average linkage,
+            two groups merge while their mean affinity exceeds ``1 - threshold``.
+            For DBSCAN this is ``eps``.
+        min_cluster_size: Smaller clusters become noise (``-1``). Use ``0`` to keep all.
+        method: ``"agglomerative"`` or ``"dbscan"``.
+        symmetrize: Average ``A`` and ``A.T`` first. The model is not exactly
+            symmetric in its two inputs; the paper evaluation used ``False``.
+    """
+    a = to_numpy(affinity).astype(np.float32)
+    if symmetrize:
+        a = 0.5 * (a + a.T)
+    dist = np.clip(1.0 - a, 0.0, None)
+    np.fill_diagonal(dist, 0.0)
+    if method == "agglomerative":
+        labels = agglomerative(dist, threshold=threshold)
+    elif method == "dbscan":
+        labels = dbscan(dist, eps=threshold)
+    else:
+        raise ValueError(f"Unknown method {method!r}")
+    return remove_small_clusters(labels, min_cluster_size) if min_cluster_size > 0 else labels

pat/inference.py

@@ -0,0 +1,212 @@
+"""High-level inference: encode a cloud once, then query affinities anywhere.
+
+Coordinates are used exactly as given. Normalize them yourself, the same way
+for training and inference.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from .cluster import NOISE, cluster_affinity, num_clusters
+from .model import PointAffinityTransformer
+from .preprocess import as_points, subsample
+
+DEFAULT_CHUNK_SIZE = 16384
+
+
+@dataclass
+class EncodedCloud:
+    """A point cloud encoded by PAT. Query it with arbitrary points.
+
+    Create it with [`pat.encode`][pat.encode]. Queries are given like the encoded
+    cloud: ``(Q, 3 + F)``, coordinates in the same frame followed by the same
+    ``F`` feature columns.
+
+    Attributes:
+        model: The model that produced the latents.
+        latents: Latent vectors, shape ``(1, num_latents, dim)``.
+        points: The (possibly subsampled) points that were encoded, including
+            feature columns.
+        index: Indices of `points` in the cloud passed to `encode`.
+    """
+
+    model: PointAffinityTransformer
+    latents: torch.Tensor
+    points: torch.Tensor
+    index: torch.Tensor
+
+    @property
+    def device(self) -> torch.device:
+        return self.latents.device
+
+    def _embed(self, queries) -> torch.Tensor:
+        return self.model.point_embed(as_points(queries, device=self.device, dim=self.model.in_dim))
+
+    @torch.no_grad()
+    def pair_affinity(self, qa, qb=None, *, chunk_size: int = DEFAULT_CHUNK_SIZE) -> torch.Tensor:
+        """Affinity of explicit point pairs.
+
+        Args:
+            qa: First points ``(P, 3 + F)``, or pairs ``(P, 2, 3 + F)`` if ``qb`` is None.
+            qb: Second points ``(P, 3 + F)``.
+
+        Returns:
+            Affinities of shape ``(P,)``.
+        """
+        if qb is None:
+            pairs = torch.as_tensor(qa, dtype=torch.float32)
+            qa, qb = pairs[:, 0], pairs[:, 1]
+        ea, eb = self._embed(qa), self._embed(qb)
+        if ea.shape != eb.shape:
+            raise ValueError("qa and qb must contain the same number of points")
+        out = [
+            self.model.decode(self.latents, a[None], b[None], embedded=True)[0]
+            for a, b in zip(ea.split(chunk_size), eb.split(chunk_size))
+        ]
+        return torch.cat(out)
+
+    @torch.no_grad()
+    def affinity(self, queries=None, others=None, *, chunk_size: int = DEFAULT_CHUNK_SIZE) -> torch.Tensor:
+        """Affinity matrix between two sets of query points.
+
+        Args:
+            queries: Points ``(Q, 3 + F)``. Defaults to the encoded points.
+            others: Points ``(K, 3 + F)``. Defaults to ``queries``.
+
+        Returns:
+            ``A[i, j]`` is the affinity of ``(queries[i], others[j])``, shape ``(Q, K)``.
+        """
+        queries = self.points if queries is None else queries
+        ea = self._embed(queries)
+        eb = ea if others is None else self._embed(others)
+        Q, K = ea.shape[0], eb.shape[0]
+
+        out = torch.empty(Q * K, device=self.device)
+        for start in range(0, Q * K, chunk_size):
+            flat = torch.arange(start, min(start + chunk_size, Q * K), device=self.device)
+            i, j = flat // K, flat % K
+            out[start : start + flat.numel()] = self.model.decode(
+                self.latents, ea[i][None], eb[j][None], embedded=True
+            )[0]
+        return out.view(Q, K)
+
+
+@torch.no_grad()
+def encode(
+    model: PointAffinityTransformer,
+    points,
+    *,
+    max_points: int | None = None,
+    generator: torch.Generator | None = None,
+) -> EncodedCloud:
+    """Encode one point cloud.
+
+    Args:
+        model: A PAT model, e.g. from [`load_pretrained`][pat.load_pretrained].
+        points: Cloud ``(N, 3 + F)``: coordinates in the same frame as the training
+            data, then the model's ``F = model.config.feature_dim`` feature columns.
+        max_points: Clouds larger than this are randomly subsampled before
+            encoding. Defaults to ``model.config.num_points``.
+        generator: Random generator for subsampling.
+    """
+    device = next(model.parameters()).device
+    points = as_points(points, dim=model.in_dim)
+    max_points = max_points or model.config.num_points
+    encoded, index = subsample(points, max_points, generator=generator)
+
+    encoded = encoded.to(device)
+    latents = model.encode(encoded[None])
+    return EncodedCloud(model, latents, encoded, index.to(device))
+
+
+@dataclass
+class Segmentation:
+    """Result of [`pat.segment`][pat.segment].
+
+    Attributes:
+        labels: Neuron label per query point, ``-1`` for noise. Shape ``(Q,)``.
+        queries: The query points, including feature columns.
+        anchors: Indices of the queries that were clustered directly.
+        affinity: Affinity matrix between the anchors.
+    """
+
+    labels: torch.Tensor
+    queries: torch.Tensor
+    anchors: torch.Tensor
+    affinity: torch.Tensor
+
+    @property
+    def num_neurons(self) -> int:
+        return num_clusters(self.labels)
+
+
+@torch.no_grad()
+def segment(
+    model: PointAffinityTransformer,
+    points,
+    queries=None,
+    *,
+    threshold: float = 0.8,
+    min_cluster_size: int = 30,
+    method: str = "agglomerative",
+    max_anchors: int = 4096,
+    symmetrize: bool = True,
+    max_points: int | None = None,
+    chunk_size: int = DEFAULT_CHUNK_SIZE,
+    generator: torch.Generator | None = None,
+) -> Segmentation:
+    """Split a point cloud into neurons.
+
+    The cloud ``points`` is encoded once. Labels are computed for ``queries``,
+    which can be any points (default: ``points`` itself).
+
+    If there are more than ``max_anchors`` queries, a random subset (the anchors)
+    is clustered on its full affinity matrix, and each remaining query joins
+    the cluster with the highest mean affinity to it, provided that mean is
+    above ``1 - threshold``. Otherwise it becomes noise. This keeps the cost
+    linear in the number of queries.
+
+    Args:
+        model: A PAT model.
+        points: Cloud to encode, ``(N, 3 + F)``: coordinates, then features.
+        queries: Points to label, ``(Q, 3 + F)``. Defaults to ``points``.
+        threshold: Distance threshold, see [`cluster_affinity`][pat.cluster.cluster_affinity].
+        min_cluster_size: Minimum anchors per neuron.
+        method: Clustering method for the anchors.
+        max_anchors: Maximum number of queries clustered directly.
+        symmetrize: Symmetrize the anchor affinity matrix.
+        max_points: Subsample ``points`` to at most this many before encoding,
+            see [`encode`][pat.encode].
+        chunk_size: Pairs per decoder pass; lower it if you run out of memory.
+        generator: Random generator for subsampling the cloud and the anchors.
+    """
+    enc = encode(model, points, max_points=max_points, generator=generator)
+    queries = as_points(points if queries is None else queries, device=enc.device, dim=model.in_dim)
+    Q = queries.shape[0]
+
+    _, anchors = subsample(queries, max_anchors, generator=generator)
+    A = enc.affinity(queries[anchors], chunk_size=chunk_size)
+    anchor_labels = cluster_affinity(
+        A, threshold=threshold, min_cluster_size=min_cluster_size, method=method, symmetrize=symmetrize
+    ).to(enc.device)
+
+    labels = torch.full((Q,), NOISE, dtype=torch.int64, device=enc.device)
+    labels[anchors] = anchor_labels
+    rest = torch.ones(Q, dtype=torch.bool, device=enc.device)
+    rest[anchors] = False
+    n_clusters = num_clusters(anchor_labels)
+
+    if rest.any() and n_clusters > 0:
+        valid = anchor_labels >= 0
+        onehot = torch.nn.functional.one_hot(anchor_labels[valid], n_clusters).float()
+        rest_idx = rest.nonzero().squeeze(1)
+        for chunk in rest_idx.split(max(1, chunk_size // max(1, anchors.numel()))):
+            a = enc.affinity(queries[chunk], queries[anchors][valid], chunk_size=chunk_size)
+            mean_aff = (a @ onehot) / onehot.sum(0)
+            best, arg = mean_aff.max(dim=1)
+            labels[chunk] = torch.where(best > 1.0 - threshold, arg, torch.full_like(arg, NOISE))
+
+    return Segmentation(labels=labels.cpu(), queries=queries.cpu(), anchors=anchors.cpu(), affinity=A.cpu())

pat/io.py

@@ -0,0 +1,49 @@
+"""Reading neuron skeletons."""
+
+from __future__ import annotations
+
+import navis
+import numpy as np
+
+
+def load_neuron(path, units: str = "um"):
+    """Read an SWC file as a ``navis.TreeNeuron``.
+
+    Coordinates are converted to ``units`` if the file declares its units
+    (as navis-written files do); otherwise they are used as they are.
+    """
+    neuron = navis.read_swc(str(path))
+    if neuron.units is not None and not neuron.units.dimensionless:
+        neuron = neuron.convert_units(units)
+    return neuron
+
+
+def node_array(neuron, features=()) -> np.ndarray:
+    """Node table of a ``navis.TreeNeuron`` as ``(N, 3 + F)``: xyz, then the ``features`` columns."""
+    return neuron.nodes[["x", "y", "z", *features]].to_numpy(dtype=np.float32)
+
+
+def read_swc(path, units: str = "um", features=()) -> np.ndarray:
+    """Read the nodes of an SWC file (see `load_neuron`).
+
+    Args:
+        path: SWC file.
+        units: Coordinate units to convert to.
+        features: Node-table columns appended after xyz, e.g. ``["radius"]``.
+
+    Returns:
+        Array of shape ``(N, 3 + len(features))``.
+    """
+    return node_array(load_neuron(path, units), features)
+
+
+def read_swcs(paths, units: str = "um", features=()) -> tuple[np.ndarray, np.ndarray]:
+    """Read several SWC files into one cloud.
+
+    Returns:
+        ``(points, labels)`` where ``labels[i]`` is the index of the file that
+        point ``i`` came from.
+    """
+    clouds = [read_swc(p, units, features) for p in paths]
+    labels = np.concatenate([np.full(len(c), i) for i, c in enumerate(clouds)])
+    return np.concatenate(clouds), labels

pat/metrics.py

@@ -0,0 +1,41 @@
+"""Segmentation metrics comparing predicted to ground-truth labels."""
+
+from __future__ import annotations
+
+import numpy as np
+from sklearn.metrics import adjusted_rand_score as _ari
+from sklearn.metrics.cluster import contingency_matrix
+
+from .cluster import to_numpy
+
+
+def _labels(x) -> np.ndarray:
+    return to_numpy(x).astype(np.int64).ravel()
+
+
+def _entropy(counts: np.ndarray, n: int) -> float:
+    p = counts[counts > 0] / n
+    return float(-(p * np.log(p)).sum())
+
+
+def voi(pred, gt) -> tuple[float, float, float]:
+    """Variation of information (natural log).
+
+    Returns:
+        ``(voi, split, merge)`` with ``split = H(pred | gt)`` (over-segmentation)
+        and ``merge = H(gt | pred)`` (under-segmentation).
+        Noise (``-1``) is treated as one regular label.
+    """
+    pred, gt = _labels(pred), _labels(gt)
+    n = pred.size
+    joint = contingency_matrix(gt, pred, sparse=True).data
+    h_joint = _entropy(joint, n)
+    h_pred = _entropy(np.unique(pred, return_counts=True)[1], n)
+    h_gt = _entropy(np.unique(gt, return_counts=True)[1], n)
+    split, merge = h_joint - h_gt, h_joint - h_pred
+    return split + merge, split, merge
+
+
+def adjusted_rand_score(pred, gt) -> float:
+    """Adjusted Rand index (1 = identical partitions)."""
+    return float(_ari(_labels(gt), _labels(pred)))