Spaces:
Runtime error
Runtime error
File size: 2,234 Bytes
827b81f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import torch
import numpy as np
def pairwise_distances(x, y):
#Input: x is a Nxd matrix
# y is an optional Mxd matirx
#Output: dist is a NxM matrix where dist[i,j] is the square norm between x[i,:] and y[j,:]
# if y is not given then use 'y=x'.
#i.e. dist[i,j] = ||x[i,:]-y[j,:]||^2
x_norm = (x ** 2).sum(1).view(-1, 1)
y_t = torch.transpose(y, 0, 1)
y_norm = (y ** 2).sum(1).view(1, -1)
dist = x_norm + y_norm - 2.0 * torch.mm(x, y_t)
return torch.clamp(dist, 0.0, np.inf)
def meanshift_cluster(pts, bandwidth, weights = None, meanshift_step = 15, step_size = 0.3):
"""
meanshift written in pytorch
:param pts: input points
:param weights: weight per point during clustering
:return: clustered points
"""
pts_steps = []
for i in range(meanshift_step):
Y = pairwise_distances(pts, pts)
K = torch.nn.functional.relu(bandwidth ** 2 - Y)
if weights is not None:
K = K * weights
P = torch.nn.functional.normalize(K, p=1, dim=0, eps=1e-10)
P = P.transpose(0, 1)
pts = step_size * (torch.matmul(P, pts) - pts) + pts
pts_steps.append(pts)
return pts_steps
def distance(a,b):
return torch.sqrt(((a-b)**2).sum())
def meanshift_assign(points, bandwidth):
cluster_ids = []
cluster_idx = 0
cluster_centers = []
for i, point in enumerate(points):
if(len(cluster_ids) == 0):
cluster_ids.append(cluster_idx)
cluster_centers.append(point)
cluster_idx += 1
else:
# assign to nearest cluster
#for j,center in enumerate(cluster_centers):
# dist = distance(point, center)
# if(dist < bandwidth):
# cluster_ids.append(j)
cdist = torch.cdist(point.unsqueeze(0), torch.stack(cluster_centers), p = 2)
nearest_idx = torch.argmin(cdist, dim = 1)
if cdist[0, nearest_idx] < bandwidth:
cluster_ids.append(nearest_idx)
else:
cluster_ids.append(cluster_idx)
cluster_centers.append(point)
cluster_idx += 1
return cluster_ids, cluster_centers
|