nystrom-ncut 0.0.6__py3-none-any.whl → 0.0.8__py3-none-any.whl
Sign up to get free protection for your applications and to get access to all the features.
- nystrom_ncut/__init__.py +2 -3
- nystrom_ncut/common.py +5 -1
- nystrom_ncut/ncut_pytorch.py +19 -29
- nystrom_ncut/nystrom.py +2 -2
- nystrom_ncut/propagation_utils.py +63 -120
- nystrom_ncut/visualize_utils.py +21 -50
- {nystrom_ncut-0.0.6.dist-info → nystrom_ncut-0.0.8.dist-info}/METADATA +1 -1
- nystrom_ncut-0.0.8.dist-info/RECORD +11 -0
- nystrom_ncut-0.0.6.dist-info/RECORD +0 -11
- {nystrom_ncut-0.0.6.dist-info → nystrom_ncut-0.0.8.dist-info}/LICENSE +0 -0
- {nystrom_ncut-0.0.6.dist-info → nystrom_ncut-0.0.8.dist-info}/WHEEL +0 -0
- {nystrom_ncut-0.0.6.dist-info → nystrom_ncut-0.0.8.dist-info}/top_level.txt +0 -0
nystrom_ncut/__init__.py
CHANGED
@@ -4,8 +4,8 @@ from .ncut_pytorch import (
|
|
4
4
|
)
|
5
5
|
from .propagation_utils import (
|
6
6
|
affinity_from_features,
|
7
|
-
|
8
|
-
|
7
|
+
extrapolate_knn_with_subsampling,
|
8
|
+
extrapolate_knn,
|
9
9
|
quantile_normalize,
|
10
10
|
)
|
11
11
|
from .visualize_utils import (
|
@@ -17,6 +17,5 @@ from .visualize_utils import (
|
|
17
17
|
rgb_from_cosine_tsne_3d,
|
18
18
|
rotate_rgb_cube,
|
19
19
|
convert_to_lab_color,
|
20
|
-
propagate_rgb_color,
|
21
20
|
get_mask,
|
22
21
|
)
|
nystrom_ncut/common.py
CHANGED
@@ -1,10 +1,14 @@
|
|
1
|
-
from typing import Any
|
1
|
+
from typing import Any, Literal
|
2
2
|
|
3
3
|
import numpy as np
|
4
4
|
import torch
|
5
5
|
import torch.nn.functional as Fn
|
6
6
|
|
7
7
|
|
8
|
+
DistanceOptions = Literal["cosine", "euclidean", "rbf"]
|
9
|
+
SampleOptions = Literal["farthest", "random"]
|
10
|
+
|
11
|
+
|
8
12
|
def ceildiv(a: int, b: int) -> int:
|
9
13
|
return -(-a // b)
|
10
14
|
|
nystrom_ncut/ncut_pytorch.py
CHANGED
@@ -1,9 +1,13 @@
|
|
1
1
|
import logging
|
2
|
-
from typing import
|
2
|
+
from typing import Tuple
|
3
3
|
|
4
4
|
import torch
|
5
5
|
import torch.nn.functional as Fn
|
6
6
|
|
7
|
+
from .common import (
|
8
|
+
DistanceOptions,
|
9
|
+
SampleOptions,
|
10
|
+
)
|
7
11
|
from .nystrom import (
|
8
12
|
EigSolverOptions,
|
9
13
|
OnlineKernel,
|
@@ -16,9 +20,6 @@ from .propagation_utils import (
|
|
16
20
|
)
|
17
21
|
|
18
22
|
|
19
|
-
DistanceOptions = Literal["cosine", "euclidean", "rbf"]
|
20
|
-
|
21
|
-
|
22
23
|
class LaplacianKernel(OnlineKernel):
|
23
24
|
def __init__(
|
24
25
|
self,
|
@@ -46,9 +47,10 @@ class LaplacianKernel(OnlineKernel):
|
|
46
47
|
affinity_focal_gamma=self.affinity_focal_gamma,
|
47
48
|
distance=self.distance,
|
48
49
|
) # [n x n]
|
50
|
+
d = features.shape[-1]
|
49
51
|
U, L = solve_eig(
|
50
52
|
self.A,
|
51
|
-
num_eig=
|
53
|
+
num_eig=d + 1, # d * (d + 3) // 2 + 1,
|
52
54
|
eig_solver=self.eig_solver,
|
53
55
|
) # [n x (d + 1)], [d + 1]
|
54
56
|
self.Ainv = U @ torch.diag(1 / L) @ U.mT # [n x n]
|
@@ -97,11 +99,10 @@ class NCUT(OnlineNystrom):
|
|
97
99
|
n_components: int = 100,
|
98
100
|
affinity_focal_gamma: float = 1.0,
|
99
101
|
num_sample: int = 10000,
|
100
|
-
sample_method:
|
102
|
+
sample_method: SampleOptions = "farthest",
|
101
103
|
distance: DistanceOptions = "cosine",
|
102
104
|
eig_solver: EigSolverOptions = "svd_lowrank",
|
103
105
|
normalize_features: bool = None,
|
104
|
-
move_output_to_cpu: bool = False,
|
105
106
|
chunk_size: int = 8192,
|
106
107
|
):
|
107
108
|
"""
|
@@ -117,7 +118,6 @@ class NCUT(OnlineNystrom):
|
|
117
118
|
eig_solver (str): eigen decompose solver, ['svd_lowrank', 'lobpcg', 'svd', 'eigh'].
|
118
119
|
normalize_features (bool): normalize input features before computing affinity matrix,
|
119
120
|
default 'None' is True for cosine distance, False for euclidean distance and rbf
|
120
|
-
move_output_to_cpu (bool): move output to CPU, set to True if you have memory issue
|
121
121
|
chunk_size (int): chunk size for large-scale matrix multiplication
|
122
122
|
"""
|
123
123
|
OnlineNystrom.__init__(
|
@@ -127,18 +127,18 @@ class NCUT(OnlineNystrom):
|
|
127
127
|
eig_solver=eig_solver,
|
128
128
|
chunk_size=chunk_size,
|
129
129
|
)
|
130
|
-
self.num_sample = num_sample
|
131
|
-
self.sample_method = sample_method
|
132
|
-
self.
|
133
|
-
self.
|
130
|
+
self.num_sample: int = num_sample
|
131
|
+
self.sample_method: SampleOptions = sample_method
|
132
|
+
self.anchor_indices: torch.Tensor = None
|
133
|
+
self.distance: DistanceOptions = distance
|
134
|
+
self.normalize_features: bool = normalize_features
|
134
135
|
if self.normalize_features is None:
|
135
136
|
if distance in ["cosine"]:
|
136
137
|
self.normalize_features = True
|
137
138
|
if distance in ["euclidean", "rbf"]:
|
138
139
|
self.normalize_features = False
|
139
140
|
|
140
|
-
self.
|
141
|
-
self.chunk_size = chunk_size
|
141
|
+
self.chunk_size: int = chunk_size
|
142
142
|
|
143
143
|
def _fit_helper(
|
144
144
|
self,
|
@@ -152,16 +152,6 @@ class NCUT(OnlineNystrom):
|
|
152
152
|
)
|
153
153
|
self.num_sample = _n
|
154
154
|
|
155
|
-
# check if features dimension greater than num_eig
|
156
|
-
if self.eig_solver in ["svd_lowrank", "lobpcg"]:
|
157
|
-
assert (
|
158
|
-
_n >= self.n_components * 2
|
159
|
-
), "number of nodes should be greater than 2*num_eig"
|
160
|
-
elif self.eig_solver in ["svd", "eigh"]:
|
161
|
-
assert (
|
162
|
-
_n >= self.n_components
|
163
|
-
), "number of nodes should be greater than num_eig"
|
164
|
-
|
165
155
|
assert self.distance in ["cosine", "euclidean", "rbf"], "distance should be 'cosine', 'euclidean', 'rbf'"
|
166
156
|
|
167
157
|
if self.normalize_features:
|
@@ -169,20 +159,20 @@ class NCUT(OnlineNystrom):
|
|
169
159
|
features = torch.nn.functional.normalize(features, dim=-1)
|
170
160
|
|
171
161
|
if precomputed_sampled_indices is not None:
|
172
|
-
|
162
|
+
_sampled_indices = precomputed_sampled_indices
|
173
163
|
else:
|
174
|
-
|
164
|
+
_sampled_indices = run_subgraph_sampling(
|
175
165
|
features,
|
176
166
|
self.num_sample,
|
177
167
|
sample_method=self.sample_method,
|
178
168
|
)
|
179
|
-
|
180
|
-
sampled_features = features[
|
169
|
+
self.anchor_indices = torch.sort(_sampled_indices).values
|
170
|
+
sampled_features = features[self.anchor_indices]
|
181
171
|
OnlineNystrom.fit(self, sampled_features)
|
182
172
|
|
183
173
|
_n_not_sampled = _n - len(sampled_features)
|
184
174
|
if _n_not_sampled > 0:
|
185
|
-
unsampled_indices = torch.full((_n,), True, device=features.device).scatter_(0,
|
175
|
+
unsampled_indices = torch.full((_n,), True, device=features.device).scatter_(0, self.anchor_indices, False)
|
186
176
|
unsampled_features = features[unsampled_indices]
|
187
177
|
V_unsampled, _ = OnlineNystrom.update(self, unsampled_features)
|
188
178
|
else:
|
nystrom_ncut/nystrom.py
CHANGED
@@ -72,7 +72,7 @@ class OnlineNystrom:
|
|
72
72
|
self.anchor_features = features
|
73
73
|
|
74
74
|
self.kernel.fit(self.anchor_features)
|
75
|
-
self.inverse_approximation_dim = max(self.n_components, features.shape[-1]
|
75
|
+
self.inverse_approximation_dim = max(self.n_components, features.shape[-1] + 1)
|
76
76
|
U, L = self._update_to_kernel() # [n x (? + 1)], [? + 1]
|
77
77
|
|
78
78
|
self.transform_matrix = (U / L)[:, :self.n_components] # [n x n_components]
|
@@ -135,7 +135,7 @@ class OnlineNystrom:
|
|
135
135
|
def solve_eig(
|
136
136
|
A: torch.Tensor,
|
137
137
|
num_eig: int,
|
138
|
-
eig_solver:
|
138
|
+
eig_solver: EigSolverOptions,
|
139
139
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
140
140
|
"""PyTorch implementation of Eigensolver cut without Nystrom-like approximation.
|
141
141
|
|
@@ -3,9 +3,14 @@ from typing import Literal
|
|
3
3
|
|
4
4
|
import numpy as np
|
5
5
|
import torch
|
6
|
-
import torch.nn.functional as
|
6
|
+
import torch.nn.functional as Fn
|
7
7
|
|
8
|
-
from .common import
|
8
|
+
from .common import (
|
9
|
+
DistanceOptions,
|
10
|
+
SampleOptions,
|
11
|
+
ceildiv,
|
12
|
+
lazy_normalize,
|
13
|
+
)
|
9
14
|
|
10
15
|
|
11
16
|
@torch.no_grad()
|
@@ -13,7 +18,7 @@ def run_subgraph_sampling(
|
|
13
18
|
features: torch.Tensor,
|
14
19
|
num_sample: int,
|
15
20
|
max_draw: int = 1000000,
|
16
|
-
sample_method:
|
21
|
+
sample_method: SampleOptions = "farthest",
|
17
22
|
):
|
18
23
|
if num_sample >= features.shape[0]:
|
19
24
|
# if too many samples, use all samples and bypass Nystrom-like approximation
|
@@ -74,7 +79,7 @@ def farthest_point_sampling(
|
|
74
79
|
def distance_from_features(
|
75
80
|
features: torch.Tensor,
|
76
81
|
features_B: torch.Tensor,
|
77
|
-
distance:
|
82
|
+
distance: DistanceOptions,
|
78
83
|
):
|
79
84
|
"""Compute affinity matrix from input features.
|
80
85
|
Args:
|
@@ -93,7 +98,11 @@ def distance_from_features(
|
|
93
98
|
D = torch.cdist(features, features_B, p=2)
|
94
99
|
elif distance == "rbf":
|
95
100
|
D = torch.cdist(features, features_B, p=2) ** 2
|
96
|
-
|
101
|
+
|
102
|
+
# Outlier-robust scale invariance using quantiles to estimate standard deviation
|
103
|
+
stds = torch.quantile(features, q=torch.tensor((0.158655, 0.841345), device=features.device), dim=0)
|
104
|
+
stds = (stds[1] - stds[0]) / 2
|
105
|
+
D = D / (2 * torch.linalg.norm(stds) ** 2)
|
97
106
|
else:
|
98
107
|
raise ValueError("distance should be 'cosine' or 'euclidean', 'rbf'")
|
99
108
|
return D
|
@@ -103,7 +112,7 @@ def affinity_from_features(
|
|
103
112
|
features: torch.Tensor,
|
104
113
|
features_B: torch.Tensor = None,
|
105
114
|
affinity_focal_gamma: float = 1.0,
|
106
|
-
distance:
|
115
|
+
distance: DistanceOptions = "cosine",
|
107
116
|
):
|
108
117
|
"""Compute affinity matrix from input features.
|
109
118
|
|
@@ -131,23 +140,23 @@ def affinity_from_features(
|
|
131
140
|
return A
|
132
141
|
|
133
142
|
|
134
|
-
def
|
135
|
-
|
136
|
-
|
137
|
-
|
138
|
-
knn: int = 10,
|
139
|
-
distance:
|
143
|
+
def extrapolate_knn(
|
144
|
+
anchor_features: torch.Tensor, # [n x d]
|
145
|
+
anchor_output: torch.Tensor, # [n x d']
|
146
|
+
extrapolation_features: torch.Tensor, # [m x d]
|
147
|
+
knn: int = 10, # k
|
148
|
+
distance: DistanceOptions = "cosine",
|
140
149
|
affinity_focal_gamma: float = 1.0,
|
141
150
|
chunk_size: int = 8192,
|
142
151
|
device: str = None,
|
143
|
-
move_output_to_cpu: bool = False
|
144
|
-
):
|
152
|
+
move_output_to_cpu: bool = False
|
153
|
+
) -> torch.Tensor: # [m x d']
|
145
154
|
"""A generic function to propagate new nodes using KNN.
|
146
155
|
|
147
156
|
Args:
|
148
|
-
|
149
|
-
|
150
|
-
|
157
|
+
anchor_features (torch.Tensor): features from subgraph, shape (num_sample, n_features)
|
158
|
+
anchor_output (torch.Tensor): output from subgraph, shape (num_sample, D)
|
159
|
+
extrapolation_features (torch.Tensor): features from existing nodes, shape (new_num_samples, n_features)
|
151
160
|
knn (int): number of KNN to propagate eige nvectors
|
152
161
|
distance (str): distance metric, 'cosine' (default) or 'euclidean', 'rbf'
|
153
162
|
chunk_size (int): chunk size for matrix multiplication
|
@@ -159,121 +168,55 @@ def propagate_knn(
|
|
159
168
|
>>> old_eigenvectors = torch.randn(3000, 20)
|
160
169
|
>>> old_features = torch.randn(3000, 100)
|
161
170
|
>>> new_features = torch.randn(200, 100)
|
162
|
-
>>> new_eigenvectors =
|
171
|
+
>>> new_eigenvectors = extrapolate_knn(old_features,old_eigenvectors,new_features,knn=3)
|
163
172
|
>>> # new_eigenvectors.shape = (200, 20)
|
164
173
|
|
165
174
|
"""
|
166
|
-
device =
|
167
|
-
|
168
|
-
if knn == 1:
|
169
|
-
return propagate_nearest(
|
170
|
-
subgraph_output,
|
171
|
-
inp_features,
|
172
|
-
subgraph_features,
|
173
|
-
chunk_size=chunk_size,
|
174
|
-
device=device,
|
175
|
-
move_output_to_cpu=move_output_to_cpu,
|
176
|
-
)
|
175
|
+
device = anchor_output.device if device is None else device
|
177
176
|
|
178
177
|
# used in nystrom_ncut
|
179
178
|
# propagate eigen_vector from subgraph to full graph
|
180
|
-
|
179
|
+
anchor_output = anchor_output.to(device)
|
181
180
|
|
182
|
-
n_chunks = ceildiv(
|
181
|
+
n_chunks = ceildiv(extrapolation_features.shape[0], chunk_size)
|
183
182
|
V_list = []
|
184
|
-
for _v in torch.chunk(
|
185
|
-
_v = _v.to(device)
|
186
|
-
|
187
|
-
# _A = affinity_from_features(subgraph_features, _v, affinity_focal_gamma, distance).mT
|
188
|
-
# if knn is not None:
|
189
|
-
# mask = torch.full_like(_A, True, dtype=torch.bool)
|
190
|
-
# mask[torch.arange(len(_v))[:, None], _A.topk(knn, dim=-1, largest=True).indices] = False
|
191
|
-
# _A[mask] = 0.0
|
192
|
-
# _A = F.normalize(_A, p=1, dim=-1)
|
193
|
-
|
194
|
-
if distance == 'cosine':
|
195
|
-
_A = _v @ subgraph_features.T
|
196
|
-
elif distance == 'euclidean':
|
197
|
-
_A = - torch.cdist(_v, subgraph_features, p=2)
|
198
|
-
elif distance == 'rbf':
|
199
|
-
_A = - torch.cdist(_v, subgraph_features, p=2) ** 2
|
200
|
-
else:
|
201
|
-
raise ValueError("distance should be 'cosine' or 'euclidean', 'rbf'")
|
202
|
-
|
203
|
-
# keep topk KNN for each row
|
204
|
-
topk_sim, topk_idx = _A.topk(knn, dim=-1, largest=True)
|
205
|
-
row_id = torch.arange(topk_idx.shape[0], device=_A.device)[:, None].expand(
|
206
|
-
-1, topk_idx.shape[1]
|
207
|
-
)
|
208
|
-
_A = torch.sparse_coo_tensor(
|
209
|
-
torch.stack([row_id, topk_idx], dim=-1).reshape(-1, 2).T,
|
210
|
-
topk_sim.reshape(-1),
|
211
|
-
size=(_A.shape[0], _A.shape[1]),
|
212
|
-
device=_A.device,
|
213
|
-
)
|
214
|
-
_A = _A.to_dense().to(dtype=subgraph_output.dtype)
|
215
|
-
_D = _A.sum(-1)
|
216
|
-
_A /= _D[:, None]
|
217
|
-
|
218
|
-
_V = _A @ subgraph_output
|
219
|
-
if move_output_to_cpu:
|
220
|
-
_V = _V.cpu()
|
221
|
-
V_list.append(_V)
|
183
|
+
for _v in torch.chunk(extrapolation_features, n_chunks, dim=0):
|
184
|
+
_v = _v.to(device) # [_m x d]
|
222
185
|
|
223
|
-
|
224
|
-
|
186
|
+
_A = affinity_from_features(anchor_features, _v, affinity_focal_gamma, distance).mT # [_m x n]
|
187
|
+
if knn is not None:
|
188
|
+
_A, indices = _A.topk(k=knn, dim=-1, largest=True) # [_m x k], [_m x k]
|
189
|
+
_anchor_output = anchor_output[indices] # [_m x k x d]
|
190
|
+
else:
|
191
|
+
_anchor_output = anchor_output[None] # [1 x n x d]
|
225
192
|
|
193
|
+
_A = Fn.normalize(_A, p=1, dim=-1) # [_m x k]
|
194
|
+
_V = (_A[:, None, :] @ _anchor_output).squeeze(1) # [_m x d]
|
226
195
|
|
227
|
-
def propagate_nearest(
|
228
|
-
subgraph_output: torch.Tensor,
|
229
|
-
inp_features: torch.Tensor,
|
230
|
-
subgraph_features: torch.Tensor,
|
231
|
-
distance: Literal["cosine", "euclidean", "rbf"] = "cosine",
|
232
|
-
chunk_size: int = 8192,
|
233
|
-
device: str = None,
|
234
|
-
move_output_to_cpu: bool = False,
|
235
|
-
):
|
236
|
-
device = subgraph_output.device if device is None else device
|
237
|
-
if distance == 'cosine':
|
238
|
-
inp_features = lazy_normalize(inp_features, dim=-1)
|
239
|
-
subgraph_features = lazy_normalize(subgraph_features, dim=-1)
|
240
|
-
|
241
|
-
# used in nystrom_tsne, equivalent to propagate_by_knn with knn=1
|
242
|
-
# propagate tSNE from subgraph to full graph
|
243
|
-
V_list = []
|
244
|
-
subgraph_features = subgraph_features.to(device)
|
245
|
-
for i in range(0, inp_features.shape[0], chunk_size):
|
246
|
-
end = min(i + chunk_size, inp_features.shape[0])
|
247
|
-
_v = inp_features[i:end].to(device)
|
248
|
-
_A = -distance_from_features(subgraph_features, _v, distance).mT
|
249
|
-
|
250
|
-
# keep top1 for each row
|
251
|
-
top_idx = _A.argmax(dim=-1).cpu()
|
252
|
-
_V = subgraph_output[top_idx]
|
253
196
|
if move_output_to_cpu:
|
254
197
|
_V = _V.cpu()
|
255
198
|
V_list.append(_V)
|
256
199
|
|
257
|
-
|
258
|
-
return
|
200
|
+
anchor_output = torch.cat(V_list, dim=0)
|
201
|
+
return anchor_output
|
259
202
|
|
260
203
|
|
261
204
|
# wrapper functions for adding new nodes to existing graph
|
262
|
-
def
|
263
|
-
|
264
|
-
|
265
|
-
|
205
|
+
def extrapolate_knn_with_subsampling(
|
206
|
+
full_features: torch.Tensor,
|
207
|
+
full_output: torch.Tensor,
|
208
|
+
extrapolation_features: torch.Tensor,
|
266
209
|
knn: int,
|
267
210
|
num_sample: int,
|
268
|
-
sample_method:
|
211
|
+
sample_method: SampleOptions,
|
269
212
|
chunk_size: int,
|
270
|
-
device: str
|
213
|
+
device: str
|
271
214
|
):
|
272
215
|
"""Propagate eigenvectors to new nodes using KNN. Note: this is equivalent to the class API `NCUT.tranform(new_features)`, expect for the sampling is re-done in this function.
|
273
216
|
Args:
|
274
|
-
|
275
|
-
|
276
|
-
|
217
|
+
full_output (torch.Tensor): eigenvectors from existing nodes, shape (num_sample, num_eig)
|
218
|
+
full_features (torch.Tensor): features from existing nodes, shape (n_samples, n_features)
|
219
|
+
extrapolation_features (torch.Tensor): features from new nodes, shape (n_new_samples, n_features)
|
277
220
|
knn (int): number of KNN to propagate eigenvectors, default 3
|
278
221
|
num_sample (int): number of samples for subgraph sampling, default 50000
|
279
222
|
sample_method (str): sample method, 'farthest' (default) or 'random'
|
@@ -286,31 +229,31 @@ def propagate_eigenvectors(
|
|
286
229
|
>>> old_eigenvectors = torch.randn(3000, 20)
|
287
230
|
>>> old_features = torch.randn(3000, 100)
|
288
231
|
>>> new_features = torch.randn(200, 100)
|
289
|
-
>>> new_eigenvectors =
|
232
|
+
>>> new_eigenvectors = extrapolate_knn_with_subsampling(extrapolation_features,old_eigenvectors,old_features,knn=3,num_sample=,sample_method=,chunk_size=,device=)
|
290
233
|
>>> # new_eigenvectors.shape = (200, 20)
|
291
234
|
"""
|
292
235
|
|
293
|
-
device =
|
236
|
+
device = full_output.device if device is None else device
|
294
237
|
|
295
238
|
# sample subgraph
|
296
|
-
|
297
|
-
|
239
|
+
anchor_indices = run_subgraph_sampling(
|
240
|
+
full_features,
|
298
241
|
num_sample,
|
299
242
|
sample_method=sample_method,
|
300
243
|
)
|
301
244
|
|
302
|
-
|
303
|
-
|
304
|
-
|
245
|
+
anchor_output = full_output[anchor_indices].to(device)
|
246
|
+
anchor_features = full_features[anchor_indices].to(device)
|
247
|
+
extrapolation_features = extrapolation_features.to(device)
|
305
248
|
|
306
249
|
# propagate eigenvectors from subgraph to new nodes
|
307
|
-
new_eigenvectors =
|
308
|
-
|
309
|
-
|
310
|
-
|
250
|
+
new_eigenvectors = extrapolate_knn(
|
251
|
+
anchor_features,
|
252
|
+
anchor_output,
|
253
|
+
extrapolation_features,
|
311
254
|
knn=knn,
|
312
255
|
chunk_size=chunk_size,
|
313
|
-
device=device
|
256
|
+
device=device
|
314
257
|
)
|
315
258
|
return new_eigenvectors
|
316
259
|
|
nystrom_ncut/visualize_utils.py
CHANGED
@@ -6,11 +6,14 @@ import torch
|
|
6
6
|
import torch.nn.functional as F
|
7
7
|
from sklearn.base import BaseEstimator
|
8
8
|
|
9
|
-
from .common import
|
9
|
+
from .common import (
|
10
|
+
DistanceOptions,
|
11
|
+
lazy_normalize,
|
12
|
+
)
|
10
13
|
from .propagation_utils import (
|
11
14
|
run_subgraph_sampling,
|
12
|
-
|
13
|
-
|
15
|
+
extrapolate_knn,
|
16
|
+
extrapolate_knn_with_subsampling,
|
14
17
|
quantile_min_max,
|
15
18
|
quantile_normalize
|
16
19
|
)
|
@@ -31,14 +34,24 @@ def _rgb_with_dimensionality_reduction(
|
|
31
34
|
reduction_dim: int,
|
32
35
|
reduction_kwargs: Dict[str, Any],
|
33
36
|
transform_func: Callable[[torch.Tensor], torch.Tensor] = _identity,
|
37
|
+
pre_smooth: bool = True,
|
34
38
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
39
|
+
|
40
|
+
if pre_smooth:
|
41
|
+
features = extrapolate_knn(
|
42
|
+
features,
|
43
|
+
features,
|
44
|
+
features,
|
45
|
+
distance="cosine",
|
46
|
+
)
|
47
|
+
|
35
48
|
subgraph_indices = run_subgraph_sampling(
|
36
49
|
features,
|
37
50
|
num_sample,
|
38
51
|
sample_method="farthest",
|
39
52
|
)
|
40
53
|
|
41
|
-
_inp = features[subgraph_indices].
|
54
|
+
_inp = features[subgraph_indices].numpy(force=True)
|
42
55
|
_subgraph_embed = reduction(
|
43
56
|
n_components=reduction_dim,
|
44
57
|
metric=metric,
|
@@ -47,14 +60,14 @@ def _rgb_with_dimensionality_reduction(
|
|
47
60
|
).fit_transform(_inp)
|
48
61
|
|
49
62
|
_subgraph_embed = torch.tensor(_subgraph_embed, dtype=torch.float32)
|
50
|
-
X_nd = transform_func(
|
63
|
+
X_nd = transform_func(extrapolate_knn(
|
64
|
+
features[subgraph_indices],
|
51
65
|
_subgraph_embed,
|
52
66
|
features,
|
53
|
-
features[subgraph_indices],
|
54
|
-
distance=metric,
|
55
67
|
knn=knn,
|
68
|
+
distance=metric,
|
56
69
|
device=device,
|
57
|
-
move_output_to_cpu=True
|
70
|
+
move_output_to_cpu=True
|
58
71
|
))
|
59
72
|
rgb = rgb_func(X_nd, q)
|
60
73
|
return X_nd, rgb
|
@@ -413,48 +426,6 @@ def rgb_from_2d_colormap(X_2d, q=0.95):
|
|
413
426
|
return rgb
|
414
427
|
|
415
428
|
|
416
|
-
def propagate_rgb_color(
|
417
|
-
rgb: torch.Tensor,
|
418
|
-
eigenvectors: torch.Tensor,
|
419
|
-
new_eigenvectors: torch.Tensor,
|
420
|
-
knn: int = 10,
|
421
|
-
num_sample: int = 1000,
|
422
|
-
sample_method: Literal["farthest", "random"] = "farthest",
|
423
|
-
chunk_size: int = 8192,
|
424
|
-
device: str = None,
|
425
|
-
):
|
426
|
-
"""Propagate RGB color to new nodes using KNN.
|
427
|
-
Args:
|
428
|
-
rgb (torch.Tensor): RGB color for each data sample, shape (n_samples, 3)
|
429
|
-
features (torch.Tensor): features from existing nodes, shape (n_samples, n_features)
|
430
|
-
new_features (torch.Tensor): features from new nodes, shape (n_new_samples, n_features)
|
431
|
-
knn (int): number of KNN to propagate RGB color, default 1
|
432
|
-
num_sample (int): number of samples for subgraph sampling, default 50000
|
433
|
-
sample_method (str): sample method, 'farthest' (default) or 'random'
|
434
|
-
chunk_size (int): chunk size for matrix multiplication, default 8192
|
435
|
-
device (str): device to use for computation, if None, will not change device
|
436
|
-
Returns:
|
437
|
-
torch.Tensor: propagated RGB color for each data sample, shape (n_new_samples, 3)
|
438
|
-
|
439
|
-
Examples:
|
440
|
-
>>> old_rgb = torch.randn(3000, 3)
|
441
|
-
>>> old_eigenvectors = torch.randn(3000, 20)
|
442
|
-
>>> new_eigenvectors = torch.randn(200, 20)
|
443
|
-
>>> new_rgb = propagate_rgb_color(old_rgb, new_eigenvectors, old_eigenvectors)
|
444
|
-
>>> # new_eigenvectors.shape = (200, 3)
|
445
|
-
"""
|
446
|
-
return propagate_eigenvectors(
|
447
|
-
eigenvectors=rgb,
|
448
|
-
features=eigenvectors,
|
449
|
-
new_features=new_eigenvectors,
|
450
|
-
knn=knn,
|
451
|
-
num_sample=num_sample,
|
452
|
-
sample_method=sample_method,
|
453
|
-
chunk_size=chunk_size,
|
454
|
-
device=device,
|
455
|
-
)
|
456
|
-
|
457
|
-
|
458
429
|
# application: get segmentation mask fron a reference eigenvector (point prompt)
|
459
430
|
def _transform_heatmap(heatmap, gamma=1.0):
|
460
431
|
"""Transform the heatmap using gamma, normalize and min-max normalization.
|
@@ -1,6 +1,6 @@
|
|
1
1
|
Metadata-Version: 2.2
|
2
2
|
Name: nystrom_ncut
|
3
|
-
Version: 0.0.
|
3
|
+
Version: 0.0.8
|
4
4
|
Summary: Normalized Cut and Nyström Approximation
|
5
5
|
Author-email: Huzheng Yang <huze.yann@gmail.com>, Wentinn Liao <wentinn.liao@gmail.com>
|
6
6
|
Project-URL: Documentation, https://github.com/JophiArcana/Nystrom-NCUT/
|
@@ -0,0 +1,11 @@
|
|
1
|
+
nystrom_ncut/__init__.py,sha256=Vlc_iAlfvTNUiJXpZLWUOaL2Q-YqZqgr7WoG6cVnD0g,439
|
2
|
+
nystrom_ncut/common.py,sha256=G6w_8_BfBUMc6r8WFgA0NH4K6am7AzZCSdrQEVjra7U,671
|
3
|
+
nystrom_ncut/ncut_pytorch.py,sha256=-SKs9AdkafJSGkeYt4LwhbKZr8oq9JA5caAqjiVDAzU,11220
|
4
|
+
nystrom_ncut/nystrom.py,sha256=-l26oiJ0oPReSGlMlYV3gftszgFdAAHAi7OFtGPZ4Ic,8802
|
5
|
+
nystrom_ncut/propagation_utils.py,sha256=0d2VhT0JrLRurd44hZbnxBvBh-QscPKxtV7VrwYtTdo,11569
|
6
|
+
nystrom_ncut/visualize_utils.py,sha256=jDjuyZ9rdd25jqrPObJgK8zCLHc3Oms0fQnaIetHk-U,17112
|
7
|
+
nystrom_ncut-0.0.8.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
|
8
|
+
nystrom_ncut-0.0.8.dist-info/METADATA,sha256=zQpx3REOOckpJSuc7N6UNpXZoqgsM5UoFWV6__DuaRQ,6058
|
9
|
+
nystrom_ncut-0.0.8.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
|
10
|
+
nystrom_ncut-0.0.8.dist-info/top_level.txt,sha256=j7g_j0S048EvguFFnGgD5Ewd3r2H6klsxd5A4dd-wHw,13
|
11
|
+
nystrom_ncut-0.0.8.dist-info/RECORD,,
|
@@ -1,11 +0,0 @@
|
|
1
|
-
nystrom_ncut/__init__.py,sha256=Cww-_OsyQHLKpgw_Wh28_tUOvIMMr7Ey8w-tH7v99xQ,452
|
2
|
-
nystrom_ncut/common.py,sha256=qdR_JwknT9H1Cv5LopwdwZfORFx-O8MLiRI6ZF1Qohc,558
|
3
|
-
nystrom_ncut/ncut_pytorch.py,sha256=wRQXUPBOW2_vutocKf0J19HrFVkBYQePAYUEfotLfx4,11701
|
4
|
-
nystrom_ncut/nystrom.py,sha256=HbwON9pLW3gEZvOmbDJwkQNHolOo1EBvwBPeh2p2uJE,8833
|
5
|
-
nystrom_ncut/propagation_utils.py,sha256=OCqnv7P9kzDlwqeJzNWpJ3TdTEpk7AD0rJhX8MazZYs,13061
|
6
|
-
nystrom_ncut/visualize_utils.py,sha256=QmBatlX7Q-ZWF_iJ1zFDnPHFuofz3tCmtoNeeoMPw3U,18558
|
7
|
-
nystrom_ncut-0.0.6.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
|
8
|
-
nystrom_ncut-0.0.6.dist-info/METADATA,sha256=FD53Ov3g9u4tbBP_Sxxd2hf1yUdg_Hy3ShWq_xGOZFA,6058
|
9
|
-
nystrom_ncut-0.0.6.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
|
10
|
-
nystrom_ncut-0.0.6.dist-info/top_level.txt,sha256=j7g_j0S048EvguFFnGgD5Ewd3r2H6klsxd5A4dd-wHw,13
|
11
|
-
nystrom_ncut-0.0.6.dist-info/RECORD,,
|
File without changes
|
File without changes
|
File without changes
|