import numpy as np
image_target_shape = np.array([512,720])
image_original_shape=np.array([416,416])
num_k_means = 9
path_txt_data = 'C:/Users/10107472/Desktop/my_yolo/data/box/train.txt'
anchor_save_file = 'C:/Users/10107472/Desktop/my_yolo/data/train_anchors.txt'
def input_data(path_txt_data): # 读取文件，其中文件格式为2.bmp 69,172,270,330,2 150,141,229,284,2 285,201,327,331,3 258,198,297,329,1
input_X_all = []
with open(path_txt_data, 'r') as f:
box_txt = f.readlines()
read_data_txt = [line.strip() for line in box_txt if len(line.strip().split()[1:]) != 0]
box_len = len(read_data_txt)
for j in range(box_len):
box_txt_one = read_data_txt[j].split()
box_txt_one = box_txt_one[1:]
box_txt_one_len = len(box_txt_one)
for i in range(box_txt_one_len):
box_txt_one_temp = np.array(list(map(int, box_txt_one[i].split(','))))[:4]
input_X_all.append(box_txt_one_temp)
input_X_all = np.array(input_X_all) # 将所有box变成二维矩阵，number * [x1，y1，x2，y2]
number_X, _ = input_X_all.shape
input_X_wh = np.zeros((number_X, 2))
input_X_wh[:, 0] = abs(input_X_all[:,0]-input_X_all[:, 2])
input_X_wh[:, 1] = abs(input_X_all[:, 1] - input_X_all[:, 3])
return input_X_wh
def IOU(input_X_wh, centroids): # 一个真实框与K个centroids 进行nms，会得到K个iou值,这里不需要x与y坐标
ious = []
X_w, X_h = input_X_wh # ground truth的w,h
for centroid in centroids:
c_w, c_h = centroid # anchor的w,h
if c_w >= X_w and c_h >= X_h: # anchor包围ground truth
iou = X_w * X_h / (c_w * c_h)
elif c_w >= X_w and c_h <= X_h: # anchor宽矮
iou = X_w * c_h / (X_w * X_h + (c_w - X_w) * c_h)
elif c_w <= X_w and c_h >= X_h: # anchor瘦长
iou = c_w * X_h / (X_w * X_h + c_w * (c_h - X_h))
else: # ground truth包围anchor means both w,h are bigger than c_w and c_h respectively
iou = (c_w * c_h) / (X_w * X_h)
ious.append(iou) # will become (k,) shape
return np.array(ious)
def avg_IOU(input_X_all, centroids):
'''
:param input_X_all: ground truth的w,h的集合[(w,h),(),...]
:param centroids: anchor的w,h的集合[(w,h),(),...]，共k个
'''
n, d = input_X_all.shape
sum = 0.
for i in range(input_X_all.shape[0]):
sum += max(IOU(input_X_all[i], centroids)) # 每次返回一个ground truth与所有anchor的IoU中的最大值，将其累加
return sum / n # 对所有ground truth求平均
def write_anchors_to_file(centroids, input_X_all, anchor_save_file, image_target_shape, image_original_shape):
f = open(anchor_save_file, 'w')
anchors = centroids.copy() # k*2 的w，h
anchors[:, 0] = image_target_shape[0] / image_original_shape[0]*anchors[:,0] # 将其得到的K对anchors变成目标尺寸对应的anchors
anchors[:, 1] = image_target_shape[1] / image_original_shape[1] * anchors[:,1]
widths = anchors[:, 0]
sorted_indices = np.argsort(widths) # withs 从小到大排序，给出widths排列的下标（原因在于最小的宽度就会对应最小的高度）
for i in sorted_indices[:-1]: # 只循环了 n-1，即保留了K-1个anchors
f.write('%0.2f,%0.2f, ' % (anchors[i, 0], anchors[i, 1])) # 将其保留在文件中
f.write('%0.2f,%0.2f\n' % (anchors[sorted_indices[-1:], 0], anchors[sorted_indices[-1:], 1])) # 将其最后一个K的anchors写入文件中，主要有了换行符号
f.write('%f\n' % (avg_IOU(input_X_all, centroids))) # 在下一行保存avg_IOU
def kmeans(input_X_all, num_k_means, anchor_save_path, image_target_shape, image_original_shape):
input_X_all_Num = input_X_all.shape[0] # ground truth的个数
centroids=np.zeros((num_k_means,2))
# print("centroids.shape", centroids)
cen_k, cen_dim_wh = centroids.shape # cen_k represent number of means and cen_dim_wh represent w ,h all of centroids # anchor的个数k以及w,h两维，dim默认等于2
prev_assignments = np.ones(input_X_all_Num) * (-1) # 1 维，共有N个-1 对每个ground truth分配初始标签
K_cen_old = input_X_all[np.random.choice(input_X_all_Num, num_k_means, replace=False)] # initialize the cluster centers to be k items [k,2]
while True:
N_K = [] # 最后保存的shape为[n,k]
for i in range(input_X_all_Num): # 真实框循环
one_X_K_value = 1 - IOU(input_X_all[i], K_cen_old) # 给出一个真实框对应k个anchors的IOU值，用1-IOU值将最大变为最小
N_K.append(one_X_K_value) # 在D中保存一个真实框对应K个anchors的值，如[array1([anchor1,anchor2,anchor3...]),array2([anchor1,anchor2,anchor3...]),...]
N_K = np.array(N_K) # 将其转换N,K的数组,其shape = (N,k)，得到每个ground truth在行中对应每个anchor的IoU
# print(N_K)
assignments = np.argmin(N_K, axis=1) # axis表示在每一行中找出最小值的序号（编号），该序号对应相应某个k中心的一个，即assignments的shape=（N,），表示相应真实框对应的K类型中心
# if (list(assignments) == list(prev_assignments)): # 如果前一次分配的结果和这次的结果相同，就输出anchor以及平均IoU，然后执行下面代码，将其保存在文件中
# write_anchors_to_file(centroids, input_X_all, anchor_save_path, image_target_shape, image_original_shape)
# break # 这句话说明如果if成立就会执行该语句，则相当于结束此函数
centroid_sums = np.zeros((cen_k, cen_dim_wh), np.float) # 初始化以便对每个簇的w,h求和
centroid_nums = np.zeros((cen_k), np.float) # 初始化以便对每个簇的w,h求和
for i in range(input_X_all_Num):
centroid_sums[assignments[i]] = centroid_sums[assignments[i]]+input_X_all[i] # 将每个簇中的ground truth的w和h分别累加
centroid_nums[assignments[i]] = centroid_nums[assignments[i]]+1
for j in range(cen_k): # 对簇中的w,h求平均
centroids[j] = centroid_sums[j] / centroid_nums[j]
prev_assignments = assignments.copy()
K_cen_old = centroids.copy()
# print(K_cen_old)
return centroids
input_X_all = input_data(path_txt_data)
aa = kmeans(input_X_all,num_k_means, anchor_save_file, image_target_shape, image_original_shape)
print(aa)
