This package extends pytorch-cluster
library with custom graph topology-based clustering algorithms for the use in PyTorch.
The package consists of the following clustering algorithms:
- Vector Radius clustering based on vector: Rygiel et al. Eigenvector Grouping for Point Cloud Vessel Labeling (GeoMedIA 2022)
- Centerline Grouping: Rygiel et al. CenterlinePointNet++: A new point cloud based architecture for coronary artery pressure drop and vFFR estimation (MICCAI 2023)
All included operations work on varying data types and are implemented for GPU - there is no version for CPU yet (opened for external contribution).
Clone repository locally, navigate to the root directory and run command below to install package in current python environment:
pip install -e .
A point clustering algorithm around the line segment (vector) specified by the pair of
import torch
from torch_cluster_topology import vector_radius
x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]]).cuda()
y = torch.Tensor([[-1, 0, 0, 0], [1, 0, 0, 0]]).cuda()
row, col = vector_radius(x, y, 1.2)
print(row, col)
tensor([0, 0, 1, 1]) tensor([0, 1, 2, 3])
A point clustering algorithm that groups surface points based on the underlying centerline.
The surface points are provided as the relative mapping ids to the centerline, and centerline is provided
as the thresholded all pair graph distance matrix (1
are considered to be neighbours for the grouping and O
not).
import torch
from torch_cluster_topology import centerline_group
x = torch.Tensor([[0], [0], [1], [1], [2], [3], [4]]).cuda()
y = torch.Tensor([[1, 0, 1, 1, 0]]).cuda()
row, col = centerline_group(x, y)
print(row, col)
tensor([0, 0, 0, 0]) tensor([0, 1, 4, 5])