fix fps_cluster

fps_cluster.py

@@ -5,12 +5,7 @@ import torch
 
 def build_tree(all_dots, dist='euclidean'):
     num_sample = all_dots.shape[0]
-    # center = all_dots.mean(axis=0)
-    center = np.median(all_dots, axis=0)
-    distances_to_center = np.linalg.norm(all_dots - center, axis=1)
-    start_idx = np.argmin(distances_to_center)
-    indices = [start_idx]
-    distances = [114514,]
+    
     if dist == 'euclidean':
         A = all_dots[:, None] - all_dots[None, :]
         A = (A ** 2).sum(-1)
@@ -23,6 +18,12 @@ def build_tree(all_dots, dist='euclidean'):
         A = 1 - A
     else:
         raise ValueError('dist must be euclidean or cosine')
+
+    d_sum = A.mean(dim=1)
+    start_idx = torch.argmin(d_sum).item()
+    indices = [start_idx]
+    distances = [114514,]
+    
     for i in range(num_sample - 1):
         _A = A[indices]
         min_dist = _A.min(dim=0).values
@@ -55,7 +56,7 @@ def build_tree(all_dots, dist='euclidean'):
     pi_indices = np.array(pi_indices)
 
     edges = np.stack([indices, pi_indices], axis=1)
-    return edges
+    return edges, levels
 
 
 def find_connected_component(edges, start_node):