mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-22 23:24:31 +08:00
81 lines
2.6 KiB
Python
81 lines
2.6 KiB
Python
# replaced torch.ops.torch_cluster.fps with a manual implementation
|
|
# to avoid having torch_cluster downloaded as dependency
|
|
# also the dependency takes a long time to install
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
import math
|
|
|
|
def fps(src: Tensor, batch: Tensor, sampling_ratio: float, start_random: bool = True):
|
|
|
|
# manually create the pointer vector
|
|
assert src.size(0) == batch.numel()
|
|
|
|
batch_size = int(batch.max()) + 1
|
|
deg = src.new_zeros(batch_size, dtype = torch.long)
|
|
|
|
deg.scatter_add_(0, batch, torch.ones_like(batch))
|
|
|
|
ptr_vec = deg.new_zeros(batch_size + 1)
|
|
torch.cumsum(deg, 0, out=ptr_vec[1:])
|
|
|
|
#return fps_sampling(src, ptr_vec, ratio)
|
|
sampled_indicies = []
|
|
|
|
for b in range(batch_size):
|
|
# start and the end of each batch
|
|
start, end = ptr_vec[b].item(), ptr_vec[b + 1].item()
|
|
# points from the point cloud
|
|
points = src[start:end]
|
|
|
|
num_points = points.size(0)
|
|
num_samples = max(1, math.ceil(num_points * sampling_ratio))
|
|
|
|
selected = torch.zeros(num_samples, device = src.device, dtype = torch.long)
|
|
distances = torch.full((num_points,), float("inf"), device = src.device)
|
|
|
|
# select a random start point
|
|
if start_random:
|
|
farthest = torch.randint(0, num_points, (1,), device = src.device)
|
|
else: farthest = torch.tensor([0], device = src.device, dtype = torch.long)
|
|
|
|
for i in range(num_samples):
|
|
selected[i] = farthest
|
|
centroid = points[farthest].squeeze(0)
|
|
dist = torch.norm(points - centroid, dim = 1) # compute euclidean distance
|
|
distances = torch.minimum(distances, dist)
|
|
farthest = torch.argmax(distances)
|
|
|
|
sampled_indicies.append(torch.arange(start, end)[selected])
|
|
|
|
return torch.cat(sampled_indicies, dim = 0)
|
|
|
|
|
|
def test_fps():
|
|
|
|
torch.manual_seed(2025)
|
|
|
|
# 2 batches with different numbers of points
|
|
points = torch.tensor([
|
|
[0.0, 0.0, 0.0], # batch 0
|
|
[1.0, 0.0, 0.0], # batch 0
|
|
[2.0, 0.0, 0.0], # batch 0
|
|
[0.0, 1.0, 0.0], # batch 1
|
|
[0.0, 2.0, 0.0], # batch 1
|
|
[0.0, 3.0, 0.0] # batch 1
|
|
], dtype=torch.float)
|
|
|
|
batch = torch.tensor([0, 0, 0, 1, 1, 1]) # batch IDs
|
|
|
|
ratio = 0.5 # sample 50% of points per batch
|
|
# jit compilation for speedups
|
|
#optimized_fps = torch.compile(fps)
|
|
|
|
outputs = fps(points, batch, ratio, start_random = True)
|
|
#outputs2 = torch.ops.torch_cluster.fps(points, batch, ratio, True) # shouldn't work
|
|
|
|
print(outputs)
|
|
|
|
if __name__ == "__main__":
|
|
test_fps()
|
|
|