|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import sys
|
| sys.path.append("./")
|
| import numpy as np
|
|
|
|
|
| def meanshift_cluster(pts_in, bandwidth, weights=None, max_iter=20):
|
| """
|
| Meanshift clustering
|
| :param pts_in: input points
|
| :param bandwidth: bandwidth
|
| :param weights: weights per pts indicting its importance in the clustering
|
| :return: points after clustering
|
| """
|
| diff = 1e10
|
| num_iter = 1
|
| while diff > 1e-3 and num_iter < max_iter:
|
| Y = np.sum(((pts_in[np.newaxis, ...] - pts_in[:, np.newaxis, :]) ** 2), axis=2)
|
| K = np.maximum(bandwidth**2 - Y, np.zeros(Y.shape))
|
| if weights is not None:
|
| K = K * weights
|
| row_sums = K.sum(axis=0, keepdims=True)
|
| P = K / (row_sums + 1e-10)
|
| P = P.transpose()
|
| pts_in_prim = 0.3 * (np.matmul(P, pts_in) - pts_in) + pts_in
|
| diff = np.sqrt(np.sum((pts_in_prim - pts_in)**2))
|
| pts_in = pts_in_prim
|
| num_iter += 1
|
| return pts_in
|
|
|
|
|
| def nms_meanshift(pts_in, density, bandwidth):
|
| """
|
| NMS to extract modes after meanshift. Code refers to sci-kit-learn.
|
| :param pts_in: input points
|
| :param density: density at each point
|
| :param bandwidth: bandwidth used in meanshift. Used here as neighbor region for NMS
|
| :return: extracted clusters.
|
| """
|
| Y = np.sum(((pts_in[np.newaxis, ...] - pts_in[:, np.newaxis, :]) ** 2), axis=2)
|
| sorted_ids = np.argsort(density)[::-1]
|
| unique = np.ones(len(sorted_ids), dtype=np.bool)
|
| dist = np.sqrt(Y)
|
| for i in sorted_ids:
|
| if unique[i]:
|
| neighbor_idxs = np.argwhere(dist[:, i] <= bandwidth)
|
| unique[neighbor_idxs.squeeze()] = 0
|
| unique[i] = 1
|
| pts_in = pts_in[unique]
|
| return pts_in
|
|
|
|
|
|
|
