add tree option

fps_cluster.py

@@ -3,7 +3,7 @@ import numpy as np
 import torch
 
 
-def build_tree(all_dots):
+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)
@@ -11,10 +11,18 @@ def build_tree(all_dots):
     start_idx = np.argmin(distances_to_center)
     indices = [start_idx]
     distances = [114514,]
-    A = all_dots[:, None] - all_dots[None, :]
-    A = (A ** 2).sum(-1)
-    A = np.sqrt(A)
-    A = torch.tensor(A)
+    if dist == 'euclidean':
+        A = all_dots[:, None] - all_dots[None, :]
+        A = (A ** 2).sum(-1)
+        A = np.sqrt(A)
+        A = torch.tensor(A)
+    elif dist == 'cosine':
+        # assume all_dots is normalized
+        A = all_dots @ all_dots.T
+        A = torch.tensor(A)
+        A = 1 - A
+    else:
+        raise ValueError('dist must be euclidean or cosine')
     for i in range(num_sample - 1):
         _A = A[indices]
         min_dist = _A.min(dim=0).values