-
Notifications
You must be signed in to change notification settings - Fork 7.8k
/
yolo_kmeans.py
39 lines (31 loc) · 1.24 KB
/
yolo_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
import numpy as np
def wh_iou(wh1, wh2):
# Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2
wh1 = wh1[:, None] # [N,1,2]
wh2 = wh2[None] # [1,M,2]
inter = np.minimum(wh1, wh2).prod(2) # [N,M]
return inter / (wh1.prod(2) + wh2.prod(2) - inter) # iou = inter / (area1 + area2 - inter)
def k_means(boxes, k, dist=np.median):
"""
yolo k-means methods
refer: https://github.com/qqwweee/keras-yolo3/blob/master/kmeans.py
Args:
boxes: 需要聚类的bboxes
k: 簇数(聚成几类)
dist: 更新簇坐标的方法(默认使用中位数,比均值效果略好)
"""
box_number = boxes.shape[0]
last_nearest = np.zeros((box_number,))
# np.random.seed(0) # 固定随机数种子
# init k clusters
clusters = boxes[np.random.choice(box_number, k, replace=False)]
while True:
distances = 1 - wh_iou(boxes, clusters)
current_nearest = np.argmin(distances, axis=1)
if (last_nearest == current_nearest).all():
break # clusters won't change
for cluster in range(k):
# update clusters
clusters[cluster] = dist(boxes[current_nearest == cluster], axis=0)
last_nearest = current_nearest
return clusters