-
Notifications
You must be signed in to change notification settings - Fork 9
/
kmeans_anchor.py
executable file
·49 lines (40 loc) · 1.42 KB
/
kmeans_anchor.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
#!/usr/bin/python3
import argparse
import json
import os
import os.path as osp
import numpy as np
from sklearn.cluster import KMeans
from coco_utils import find_all_img_anns
def kmeans_anchor(coco_path, n_clusters, img_size):
wh_list = []
with open(coco_path, 'r') as f:
coco = f.read()
coco = json.loads(coco)
img_info_list, anns_list = find_all_img_anns(coco)
for i in range(len(img_info_list)):
img_info = img_info_list[i]
anns = anns_list[i]
for ann in anns:
wh_list.append([ann['bbox'][2] / img_info['width'], ann['bbox'][3] / img_info['height']])
wh_list = np.float32(wh_list)
estimator = KMeans(n_clusters=n_clusters)
estimator.fit(wh_list)
anchors = np.float32(estimator.cluster_centers_)
anchors[:, 0] *= img_size[0]
anchors[:, 1] *= img_size[1]
area = anchors[:, 0] * anchors[:, 1]
output = 'anchors: '
for i in range(n_clusters):
index = np.argmax(area)
area[index] = 0
output += '[%d,%d], ' % (int(anchors[index][0]), int(anchors[index][1]))
print(output[:-2])
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('coco', type=str)
parser.add_argument('-n', type=int, default=9)
parser.add_argument('-s', '--img-size', type=int, nargs=2, default=[416, 416])
opt = parser.parse_args()
print(opt)
kmeans_anchor(opt.coco, opt.n, opt.img_size)