-
Notifications
You must be signed in to change notification settings - Fork 126
/
kmeans.py
executable file
·317 lines (266 loc) · 13.5 KB
/
kmeans.py
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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
'''
2019.10.8 ming71
功能: 对box进行anchor的kmeans聚类
注意:
- 停止条件是最小值索引不变而不是最小值不变，会造成早停，可以改
- 暂时仅支持voc标注,如需改动再重写get_all_boxes函数即可
评价方法：
- anchor聚类采用iou评价 / 可视化(method1情况下)
- area和ratio聚类采用可视化散点图
'''
import numpy as np
import glob
import os
import cv2
from decimal import Decimal
from tqdm import tqdm
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
## a sample for kmeans via sklearn:
# import numpy as np
# from sklearn.cluster import KMeans
# data = np.random.rand(100, 3) #生成一个随机数据，样本大小为100, 特征数为3
# import ipdb; ipdb.set_trace()
# #假如我要构造一个聚类数为3的聚类器
# estimator = KMeans(n_clusters=3)#构造聚类器
# estimator.fit(data)#聚类
# label_pred = estimator.labels_ #获取聚类标签
# centroids = estimator.cluster_centers_ #获取聚类中心
# inertia = estimator.inertia_ # 获取聚类准则的总和
class Kmeans:
def __init__(self, cluster_number, all_boxes, save_path=None):
self.cluster_number = cluster_number
self.all_boxes = all_boxes
self.save_path = save_path
# 输入两个二维数组:所有box和种子点box
# 输出[num_boxes, k]的结果
def iou(self, boxes, clusters): # 1 box -> k clusters
n = boxes.shape[0]
k = self.cluster_number #类别
box_area = boxes[:, 0] * boxes[:, 1] #列表切片操作：取所有行0列和1列相乘 ，得到gt的面积的行向量
box_area = box_area.repeat(k) #行向量进行重复
box_area = np.reshape(box_area, (n, k))
cluster_area = clusters[:, 0] * clusters[:, 1] #种子点的面积行向量
cluster_area = np.tile(cluster_area, [1, n])
cluster_area = np.reshape(cluster_area, (n, k))
box_w_matrix = np.reshape(boxes[:, 0].repeat(k), (n, k))
cluster_w_matrix = np.reshape(np.tile(clusters[:, 0], (1, n)), (n, k))
min_w_matrix = np.minimum(cluster_w_matrix, box_w_matrix)
box_h_matrix = np.reshape(boxes[:, 1].repeat(k), (n, k))
cluster_h_matrix = np.reshape(np.tile(clusters[:, 1], (1, n)), (n, k))
min_h_matrix = np.minimum(cluster_h_matrix, box_h_matrix)
inter_area = np.multiply(min_w_matrix, min_h_matrix)
result = inter_area / (box_area + cluster_area - inter_area + 1e-16)
assert (result>0).all() == True , 'negtive anchors present , cluster again!'
return result
def avg_iou(self, boxes, clusters):
accuracy = np.mean([np.max(self.iou(boxes, clusters), axis=1)])
return accuracy
def result2txt(self, data):
f = open(self.save_path, 'w')
row = np.shape(data)[0]
for i in range(row):
if i == 0:
x_y = "%d,%d" % (data[i][0], data[i][1])
else:
x_y = ", %d,%d" % (data[i][0], data[i][1])
f.write(x_y)
f.close() #最终输出的是w1,h1,w2,h2,w3,h3,...
def anchor_clusters(self):
boxes = np.array(self.all_boxes) #返回全部gt的宽高二维数组
k=self.cluster_number
############ K-means聚类计算 ######
##### Method 1 : sklearn implemention
estimator = KMeans(n_clusters=k)
estimator.fit(boxes) #聚类
label_pred = estimator.labels_ #获取聚类标签
centroids = estimator.cluster_centers_ #获取聚类中心
centroids = np.array(centroids)
result = centroids[np.lexsort(centroids.T[0, None])] #将得到的三个anchor按照宽进行从小到大，重新排序
print("K anchors:\n {}\n".format(result))
print("Accuracy: {:.2f}%\n".format(self.avg_iou(boxes, result) * 100))
plt.figure()
plt.scatter(boxes[:,0], boxes[:,1], marker='.',c=label_pred)
plt.xlabel('anchor_w')
plt.ylabel('anchor_h')
plt.title('anchor_clusters')
for c in centroids:
plt.annotate(s='cluster' ,xy=c ,xytext=c-20,arrowprops=dict(facecolor='red',width=3,headwidth = 6))
plt.scatter(c[0], c[1], marker='*',c='red',s=100)
##### Method 2 : 自己写一边kmeans，这个的iou更高，推荐使用
# #注意：这里代码选择的停止聚类的条件是最小值的索引不变，而不是种子点的数值不变。这样的误差理论会大一点，其实关系不大。
# box_number = boxes.shape[0] # box个数
# distances = np.empty((box_number, k)) # 初始化[box_number , k]二维数组，存放自定义iou距离（obj*anchor）
# last_nearest = np.zeros((box_number,)) # [box_number , ]的标量
# np.random.seed()
# clusters = boxes[np.random.choice(
# box_number, k, replace=False)] # 种子点随机初始化
# # 种子点一旦重复会有计算错误,避免!
# while True :
# uniques_clusters = np.unique(clusters,axis=0)
# if len(uniques_clusters)==len(clusters) :
# break
# clusters = boxes[np.random.choice(box_number, k, replace=False)]
# # k-means
# while True:
# # 每轮循环，计算种子点外所有点各自到k个种子点的自定义距离，并且按照距离各个点找离自己最近的种子点进行归类；计算新的各类中心；然后下一轮循环
# distances = 1 - self.iou(boxes, clusters) # iou越大,距离越小
# current_nearest = np.argmin(distances, axis=1) # 展开为box_number长度向量,代表每个box当前属于哪个种子点类别(0,k-1)
# if (last_nearest == current_nearest).all(): # 每个box的当前类别所属和上一次相同,不再移动聚类
# break
# #计算新的k个种子点坐标
# for cluster in range(k):
# clusters[cluster] = np.median(boxes[current_nearest == cluster], axis=0) # 只对还需要聚类的种子点进行位移
# last_nearest = current_nearest
# result = clusters[np.lexsort(clusters.T[0, None])] #将得到的三个anchor按照宽进行从小到大，重新排序
# print('\n-----anchor_cluster-----\n')
# print("K anchors:\n {}\n".format(result))
# print("Accuracy: {:.2f}%\n".format(self.avg_iou(boxes, result) * 100))
# if self.save_path:
# self.result2txt(result)
# ## 聚类结果分析
# with open(self.save_path,'r') as f:
# contents = f.read()
# w = list(map(int, contents.split(',')[::2]))
# h = list(map(int, contents.split(',')[1::2]))
# anchors = [anchor for anchor in zip(w,h)]
# ratio = [Decimal(anchor[0]/anchor[1]).quantize(Decimal('0.00')) for anchor in anchors]
# ratio.sort()
# area = [Decimal(anchor[0]*anchor[1]).quantize(Decimal('0.00')) for anchor in anchors]
# area.sort()
# ##### 自定义需要分析的数据 ###
# squre=[float(s)**0.5 for s in area]
# print('ratio:\n{}\n\narea:\n{}\n'.format(ratio,area))
# print('sqrt(area):\n{}'.format(squre))
## 懒得重写，直接用sklearn
def area_cluster(self,vis=False):
# 面积聚类
boxes = np.array(self.all_boxes)
areas = boxes[:,0]*boxes[:,1]
estimator = KMeans(n_clusters=self.cluster_number)
estimator.fit(areas.reshape(-1,1)) #聚类
label_pred = estimator.labels_ #获取聚类标签
centroids = estimator.cluster_centers_ #获取聚类中心
centroids = centroids[np.lexsort(centroids.T)] # 排个序
centroids = np.array([int(i) for i in centroids]).reshape(-1,1) # 取个整
print('\n-----area_cluster-----\n')
print(centroids)
if vis:
plt.figure()
plt.scatter(range(len(areas)), areas.squeeze(), marker='.',c=label_pred)
plt.xlabel('gt_num')
plt.ylabel('area')
plt.title('area_cluster')
for c in centroids:
xy = np.array([int(0.5*len(boxes)),c.item()])
plt.scatter(int(0.5*len(boxes)),c.item(), marker='*',c='red',s=100)
def ratio_cluster(self,vis=False):
# 宽高比聚类
boxes = np.array(self.all_boxes)
ratios = boxes[:,0]/boxes[:,1]
estimator = KMeans(n_clusters=self.cluster_number)
estimator.fit(ratios.reshape(-1,1)) #聚类
label_pred = estimator.labels_ #获取聚类标签
centroids = estimator.cluster_centers_ #获取聚类中心
centroids = centroids[np.lexsort(centroids.T)] # 排个序(从小到大)
# 表示为分子或分母1便于直观观察
print('\n-----ratio_cluster-----\n')
for i,c in enumerate(centroids):
num, den = c.item().as_integer_ratio()
if c > 1 : num /= den ; den = 1 ; num = Decimal(num).quantize(Decimal('0.00'))
if c < 1 : den /= num ; num = 1 ; den = Decimal(den).quantize(Decimal('0.00'))
ratio = str(num) + '/' +str(den)
print(ratio)
if vis:
plt.figure()
plt.scatter(range(len(ratios)), ratios.squeeze(), marker='.',c=label_pred)
plt.xlabel('gt_num')
plt.ylabel('ratio')
plt.title('ratio_cluster')
for c in centroids:
xy = np.array([int(0.5*len(boxes)),c.item()])
plt.scatter(int(0.5*len(boxes)),c.item(), marker='*',c='red',s=100)
# 返回所有label的box,形式为[[w1,h1],[w2,h2],...]
def get_all_boxes(path,mode=None):
assert not mode is None,'Input correct label mode,such as : voc, hrsc, yolo'
boxes = []
if mode == 'voc':
labels = sorted(glob.glob(os.path.join(path, '*.*')))
for label in labels:
with open(label,'r') as f:
contents = f.read()
objects = contents.split('<object>')
objects.pop(0)
if len(objects) == 0: pass
for object in objects:
xmin = int(float(object[object.find('<xmin>')+6 : object.find('</xmin>')]))
xmax = int(float(object[object.find('<xmax>')+6 : object.find('</xmax>')]))
ymin = int(float(object[object.find('<ymin>')+6 : object.find('</ymin>')]))
ymax = int(float(object[object.find('<ymax>')+6 : object.find('</ymax>')]))
box_w = xmax - xmin
box_h = ymax - ymin
boxes.append((box_w,box_h))
elif mode == 'hrsc': # xml格式
rotate = True
labels = sorted(glob.glob(os.path.join(path, '*.*')))
for label in labels:
with open(label,'r') as f:
contents = f.read()
objects = contents.split('<HRSC_Object>')
objects.pop(0)
if len(objects) == 0: pass
for object in objects:
if not rotate:
xmin = int(object[object.find('<box_xmin>')+10 : object.find('</box_xmin>')])
ymin = int(object[object.find('<box_ymin>')+10 : object.find('</box_ymin>')])
xmax = int(object[object.find('<box_xmax>')+10 : object.find('</box_xmax>')])
ymax = int(object[object.find('<box_ymax>')+10 : object.find('</box_ymax>')])
box_w = xmax - xmin; box_h = ymax - ymin
else: # 旋转框
box_w = int(float(object[object.find('<mbox_w>')+8 : object.find('</mbox_w>')]))
box_h = int(float(object[object.find('<mbox_h>')+8 : object.find('</mbox_h>')]))
boxes.append((box_w,box_h))
elif mode == 'yolo':
labels = sorted(glob.glob(os.path.join(path, '*.txt*')))
for label in tqdm(labels,desc='Loading labels'):
img_path = os.path.join(os.path.split(label)[0], os.path.split(label)[1][:-4]+'.jpg')
height,width,_ = cv2.imread(img_path).shape
with open(label,'r') as f:
contents=f.read()
lines=contents.split('\n')
lines = [x for x in contents.split('\n') if x] # 移除空格
for object in lines:
coors = object.split(' ')
box_w = int(float(coors[3])*width)
box_h = int(float(coors[4])*height)
boxes.append((box_w,box_h))
else:
print('Unrecognized label mode!!')
return boxes
if __name__ == "__main__":
cluster_number = 3 # 种子点个数
label_path = '/py/datasets/HRSC2016/yolo-dataset/train'
save_path = 'anchor-cluster.txt'
all_boxes = get_all_boxes(label_path,'yolo')
kmeans = Kmeans(cluster_number, all_boxes, save_path=save_path)
vis = True
kmeans.anchor_clusters()
kmeans.area_cluster(vis=vis)
kmeans.ratio_cluster(vis=vis)
if vis:
plt.show()
'''
K anchors:
[[153.87419355 31.30967742]
[343.26412214 52.12366412]
[574.57024793 109.7231405 ]]
Accuracy: 69.19%
-----area_cluster-----
[[ 12130]
[ 42951]
[113378]]
-----ratio_cluster-----
4.18/1
6.50/1
8.75/1
'''