当前位置:   article > 正文

Pytorch机器学习(十)—— 目标检测中k-means聚类方法生成锚框anchor_锚框聚类

锚框聚类

Pytorch机器学习(十)—— YOLO中k-means聚类方法生成锚框anchor




前言

前面文章说过有关锚框的一些知识,但有个坑一直没填,就是在YOLO中锚框的大小是如何确定出来的。其实在YOLOV3中就有采用k-means聚类方法计算锚框的方法,而在YOLOV5中作者在基于k-means聚类方法的结果之后,采用了遗传算法,进一步得到效果更好的锚框。

如果对锚框概念不理解的,可以看一下这篇文章

Pytorch机器学习(九)—— YOLO中对于锚框,预测框,产生候选区域及对候选区域进行标注详解


一、K-means聚类

在YOLOV3中,锚框大小的计算就是采用的k-means聚类的方法形成的。

从直观的理解,我们知道所有已经标注的bbox的长宽大小,而锚框则是对于预测这些bbox的潜在候选框,所以锚框的长宽形状应该越接近真实bbox越好。而又由于YOLO网络的预测层是包含3种尺度的信息的(分别对应3种感受野),每种尺度的anchor又是三种,所以我们就需要9种尺度的anchor,也即我们需要对所有的bbox的尺寸聚类成9种类别!!

聚类方法比较常用的是使用k-means聚类方法,其算法流程如下。

  • 从数据集中随机选取 K 个点作为初始聚类的中心,中心点为C=\left \{ c_{1},c_{2},....,c_{k} \right \} 
  • 针对数据集中每个样本 xi,计算它们到各个聚类中心点的距离,到哪个聚类中心点的距离最小,就将其划分到对应聚类中心的类中
  • 针对每个类别 i ,重新计算该类别的聚类中心  {c_{i}}=\frac{1}{\left |i \right |}\sum x(其中 | ||i| 表示的是该类别数据的总个数)
  • 重复第二步和第三步,直到聚类中心的位置不再发生变化(我们也可以设置迭代次数)

 k-means代码

  1. # 计算中心点和其他点直接的距离
  2. def calc_distance(obs, guess_central_points):
  3. """
  4. :param obs: 所有的观测点
  5. :param guess_central_points: 中心点
  6. :return:每个点对应中心点的距离
  7. """
  8. distances = []
  9. for x, y in obs:
  10. distance = []
  11. for xc, yc in guess_central_points:
  12. distance.append(math.dist((x, y), (xc, yc)))
  13. distances.append(distance)
  14. return distances
  15. def k_means(obs, k, dist=np.median):
  16. """
  17. :param obs: 待观测点
  18. :param k: 聚类数k
  19. :param dist: 表征聚类中心函数
  20. :return: guess_central_points中心点
  21. current_cluster 分类结果
  22. """
  23. obs_num = obs.shape[0]
  24. if k < 1:
  25. raise ValueError("Asked for %d clusters." % k)
  26. # 随机取中心点
  27. guess_central_points = obs[np.random.choice(obs_num, size=k, replace=False)] # 初始化最大距离
  28. last_cluster = np.zeros((obs_num, ))
  29. # 当小于一定值时聚类完成
  30. while True:
  31. # 关键是下面的calc_distance,来计算需要的距离
  32. distances = calc_distance(obs, guess_central_points)
  33. # 获得对应距离最小值的索引
  34. current_cluster = np.argmin(distances, axis=1)
  35. # 如果聚类类别没有改变, 则直接退出
  36. if (last_cluster == current_cluster).all():
  37. break
  38. # 计算新的中心
  39. for i in range(k):
  40. guess_central_points[i] = dist(obs[current_cluster == i], axis=0)
  41. last_cluster = current_cluster
  42. return guess_central_points, current_cluster

 聚类效果如下

k-means++算法

还有一种k-means++算法,是属于k-means算法的衍生吧,其主要解决的是k-means算法第一步,随机选择中心点的问题。

整个代码也十分简单,只需要把最先随机选取中心点用下面代码计算出来就可以。

  1. # k_means++计算中心坐标
  2. def calc_center(boxes):
  3. box_number = boxes.shape[0]
  4. # 随机选取第一个中心点
  5. first_index = np.random.choice(box_number, size=1)
  6. clusters = boxes[first_index]
  7. # 计算每个样本距中心点的距离
  8. dist_note = np.zeros(box_number)
  9. dist_note += np.inf
  10. for i in range(k):
  11. # 如果已经找够了聚类中心,则退出
  12. if i+1 == k:
  13. break
  14. # 计算当前中心点和其他点的距离
  15. for j in range(box_number):
  16. j_dist = single_distance(boxes[j], clusters[i])
  17. if j_dist < dist_note[j]:
  18. dist_note[j] = j_dist
  19. # 转换为概率
  20. dist_p = dist_note / dist_note.sum()
  21. # 使用赌轮盘法选择下一个点
  22. next_index = np.random.choice(box_number, 1, p=dist_p)
  23. next_center = boxes[next_index]
  24. clusters = np.vstack([clusters, next_center])
  25. return clusters

但我自己在使用过程中,对于提升不大。主要因为其实bbox的尺度差异一般不会太大,所以这个中心点的选取,对于最后影响不大。


二、YOLO中使用k-means聚类生成anchor

下面重点说一下如何使用这个k-means算法来生成anchor,辅助我们训练,下面的代码和上面的有一点不一样,因为我们上面的代码是基于点的(x,y),而我们聚类中,是bbox的(w,h),下面代码都以VOC格式的训练集为例,如果是coco格式的,得麻烦你自己转一下格式了。

如果不想用我下面的代码但也想用k-means聚类,请读取自己数据集时,读取bbox和图片的(w,h)以列表的形式保存,确保自己n*2或者m*2的列表.

读取VOC格式数据集

我下面的代码,不仅读取了voc格式的数据集,还做了一些数据的统计,如果不想要,自己注释点就好,代码比较简单,也写了注释。

大家可以不用太纠结代码实现,记得改一下自己的图片路径即可。

  1. from xml.dom.minidom import parse
  2. import matplotlib.pyplot as plt
  3. import cv2 as cv
  4. import os
  5. train_annotation_path = '/home/aistudio/data/train/Annotations' # 训练集annotation的路径
  6. train_image_path = '/home/aistudio/data/train/JPEGImages' # 训练集图片的路径
  7. # 展示图片的数目
  8. show_num = 12
  9. #打开xml文档
  10. def parase_xml(xml_path):
  11. """
  12. 输入:xml路径
  13. 返回:image_name, width, height, bboxes
  14. """
  15. domTree = parse(xml_path)
  16. rootNode = domTree.documentElement
  17. # 得到object,sizem,图片名称属性
  18. object_node = rootNode.getElementsByTagName("object")
  19. shape_node = rootNode.getElementsByTagName("size")
  20. image_node = rootNode.getElementsByTagName("filename")
  21. image_name = image_node[0].childNodes[0].data
  22. bboxes = []
  23. # 解析图片的长宽
  24. for size in shape_node:
  25. width = int(size.getElementsByTagName('width')[0].childNodes[0].data)
  26. height = int(size.getElementsByTagName('height')[0].childNodes[0].data)
  27. # 解析图片object属性
  28. for obj in object_node:
  29. # 解析name属性,并统计类别数
  30. class_name = obj.getElementsByTagName("name")[0].childNodes[0].data
  31. # 解析bbox属性,并统计bbox的大小
  32. bndbox = obj.getElementsByTagName("bndbox")
  33. for bbox in bndbox:
  34. x1 = int(bbox.getElementsByTagName('xmin')[0].childNodes[0].data)
  35. y1 = int(bbox.getElementsByTagName('ymin')[0].childNodes[0].data)
  36. x2 = int(bbox.getElementsByTagName('xmax')[0].childNodes[0].data)
  37. y2 = int(bbox.getElementsByTagName('ymax')[0].childNodes[0].data)
  38. bboxes.append([class_name, x1, y1, x2, y2])
  39. return image_name, width, height, bboxes
  40. def read_voc(train_annotation_path, train_image_path, show_num):
  41. """
  42. train_annotation_path:训练集annotation的路径
  43. train_image_path:训练集图片的路径
  44. show_num:展示图片的大小
  45. """
  46. # 用于统计图片的长宽
  47. total_width, total_height = 0, 0
  48. # 用于统计图片bbox长宽
  49. bbox_total_width, bbox_total_height, bbox_num = 0, 0, 0
  50. min_bbox_size = 40000
  51. max_bbox_size = 0
  52. # 用于统计聚类所用的图片长宽,bbox长宽
  53. img_wh = []
  54. bbox_wh = []
  55. # 用于统计标签
  56. total_size = []
  57. class_static = {'crazing': 0, 'inclusion': 0, 'patches': 0, 'pitted_surface': 0, 'rolled-in_scale': 0, 'scratches': 0}
  58. num_index = 0
  59. for root, dirs, files in os.walk(train_annotation_path):
  60. for file in files:
  61. num_index += 1
  62. xml_path = os.path.join(root, file)
  63. image_name, width, height, bboxes = parase_xml(xml_path)
  64. image_path = os.path.join(train_image_path, image_name)
  65. img_wh.append([width, height])
  66. total_width += width
  67. total_height += height
  68. # 如果需要展示,则读取图片
  69. if num_index < show_num:
  70. image_path = os.path.join(train_image_path, image_name)
  71. image = cv.imread(image_path)
  72. # 统计有关bbox的信息
  73. wh = []
  74. for bbox in bboxes:
  75. class_name = bbox[0]
  76. class_static[class_name] += 1
  77. x1, y1, x2, y2 = bbox[1], bbox[2], bbox[3], bbox[4]
  78. bbox_width = x2 - x1
  79. bbox_height = y2 - y1
  80. bbox_size = bbox_width*bbox_height
  81. # 统计bbox的最大最小尺寸
  82. if min_bbox_size > bbox_size:
  83. min_bbox_size = bbox_size
  84. if max_bbox_size < bbox_size:
  85. max_bbox_size = bbox_size
  86. total_size.append(bbox_size)
  87. # 统计bbox平均尺寸
  88. bbox_total_width += bbox_width
  89. bbox_total_height += bbox_height
  90. # 用于聚类使用
  91. wh.append([bbox_width / width, bbox_height / height]) # 相对坐标
  92. bbox_num += 1
  93. # 如果需要展示,绘制方框
  94. if num_index < show_num:
  95. cv.rectangle(image, (x1, y1), (x2, y2), color=(255, 0, 0), thickness=2)
  96. cv.putText(image, class_name, (x1, y1+10), cv.FONT_HERSHEY_SIMPLEX, fontScale=0.2, color=(0, 255, 0), thickness=1)
  97. bbox_wh.append(wh)
  98. # 如果需要展示
  99. if num_index < show_num:
  100. plt.figure()
  101. plt.imshow(image)
  102. plt.show()
  103. # 去除2个检查文件
  104. # num_index -= 2
  105. print("total train num is: {}".format(num_index))
  106. print("avg total_width is {}, avg total_height is {}".format((total_width / num_index), (total_height / num_index)))
  107. print("avg bbox width is {}, avg bbox height is {} ".format((bbox_total_width / bbox_num), (bbox_total_height / bbox_num)))
  108. print("min bbox size is {}, max bbox size is {}".format(min_bbox_size, max_bbox_size))
  109. print("class_static show below:", class_static)
  110. return img_wh, bbox_wh
  111. img_wh, bbox_wh = read_voc(train_annotation_path, train_image_path, show_num)

k-means聚类生成anchor

我这里的k-means代码集合了k-means++的实现,也集合了 太阳花的小绿豆这位博主提出用IOU作为评价指标来计算k-means而不是用欧拉距离的方法可以测试发现,使用IOU确实效果要比使用欧拉距离做为评价指标要好)

  1. import numpy as np
  2. # 这里IOU的概念更像是只是考虑anchor的长宽
  3. def wh_iou(wh1, wh2):
  4. # Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2
  5. wh1 = wh1[:, None] # [N,1,2]
  6. wh2 = wh2[None] # [1,M,2]
  7. inter = np.minimum(wh1, wh2).prod(2) # [N,M]
  8. return inter / (wh1.prod(2) + wh2.prod(2) - inter) # iou = inter / (area1 + area2 - inter)
  9. # 计算单独一个点和一个中心的距离
  10. def single_distance(center, point):
  11. center_x, center_y = center[0]/2 , center[1]/2
  12. point_x, point_y = point[0]/2, point[1]/2
  13. return np.sqrt((center_x - point_x)**2 + (center_y - point_y)**2)
  14. # 计算中心点和其他点直接的距离
  15. def calc_distance(boxes, clusters):
  16. """
  17. :param obs: 所有的观测点
  18. :param clusters: 中心点
  19. :return:每个点对应中心点的距离
  20. """
  21. distances = []
  22. for box in boxes:
  23. # center_x, center_y = x/2, y/2
  24. distance = []
  25. for center in clusters:
  26. # center_xc, cneter_yc = xc/2, yc/2
  27. distance.append(single_distance(box, center))
  28. distances.append(distance)
  29. return distances
  30. # k_means++计算中心坐标
  31. def calc_center(boxes, k):
  32. box_number = boxes.shape[0]
  33. # 随机选取第一个中心点
  34. first_index = np.random.choice(box_number, size=1)
  35. clusters = boxes[first_index]
  36. # 计算每个样本距中心点的距离
  37. dist_note = np.zeros(box_number)
  38. dist_note += np.inf
  39. for i in range(k):
  40. # 如果已经找够了聚类中心,则退出
  41. if i+1 == k:
  42. break
  43. # 计算当前中心点和其他点的距离
  44. for j in range(box_number):
  45. j_dist = single_distance(boxes[j], clusters[i])
  46. if j_dist < dist_note[j]:
  47. dist_note[j] = j_dist
  48. # 转换为概率
  49. dist_p = dist_note / dist_note.sum()
  50. # 使用赌轮盘法选择下一个点
  51. next_index = np.random.choice(box_number, 1, p=dist_p)
  52. next_center = boxes[next_index]
  53. clusters = np.vstack([clusters, next_center])
  54. return clusters
  55. # k-means聚类,且评价指标采用IOU
  56. def k_means(boxes, k, dist=np.median, use_iou=True, use_pp=False):
  57. """
  58. yolo k-means methods
  59. Args:
  60. boxes: 需要聚类的bboxes,bboxes为n*2包含w,h
  61. k: 簇数(聚成几类)
  62. dist: 更新簇坐标的方法(默认使用中位数,比均值效果略好)
  63. use_iou:是否使用IOU做为计算
  64. use_pp:是否是同k-means++算法
  65. """
  66. box_number = boxes.shape[0]
  67. last_nearest = np.zeros((box_number,))
  68. # 在所有的bboxes中随机挑选k个作为簇的中心
  69. if not use_pp:
  70. clusters = boxes[np.random.choice(box_number, k, replace=False)]
  71. # k_means++计算初始值
  72. else:
  73. clusters = calc_center(boxes, k)
  74. # print(clusters)
  75. while True:
  76. # 计算每个bboxes离每个簇的距离 1-IOU(bboxes, anchors)
  77. if use_iou:
  78. distances = 1 - wh_iou(boxes, clusters)
  79. else:
  80. distances = calc_distance(boxes, clusters)
  81. # 计算每个bboxes距离最近的簇中心
  82. current_nearest = np.argmin(distances, axis=1)
  83. # 每个簇中元素不在发生变化说明以及聚类完毕
  84. if (last_nearest == current_nearest).all():
  85. break # clusters won't change
  86. for cluster in range(k):
  87. # 根据每个簇中的bboxes重新计算簇中心
  88. clusters[cluster] = dist(boxes[current_nearest == cluster], axis=0)
  89. last_nearest = current_nearest
  90. return clusters

使用我下面的auot_anchor代码注意!!(这里代码也是借鉴的太阳花的小绿豆博主的,他把里面的torch函数改为np函数后,使得代码移植性变强了!)

传入的参数中img_wh和bbox_wh即读取voc数据集中图片的长宽和bbox的长宽,为n*2和m*2的列表 !!

这里我还加入了YOLOV5中的遗传算法,具体细节就不展开了。

  1. from tqdm import tqdm
  2. import random
  3. # 计算聚类和遗传算法出来的anchor和真实bbox之间的重合程度
  4. def anchor_fitness(k: np.ndarray, wh: np.ndarray, thr: float): # mutation fitness
  5. """
  6. 输入:k:聚类完后的结果,且排列为升序
  7. wh:包含bbox中w,h的集合,且转换为绝对坐标
  8. thr:bbox中和k聚类的框重合阈值
  9. """
  10. r = wh[:, None] / k[None]
  11. x = np.minimum(r, 1. / r).min(2) # ratio metric
  12. best = x.max(1)
  13. f = (best * (best > thr).astype(np.float32)).mean() # fitness
  14. bpr = (best > thr).astype(np.float32).mean() # best possible recall
  15. return f, bpr
  16. def auto_anchor(img_size, n, thr, gen, img_wh, bbox_wh):
  17. """
  18. 输入:img_size:图片缩放的大小
  19. n:聚类数
  20. thr:fitness的阈值
  21. gen:遗传算法迭代次数
  22. img_wh:图片的长宽集合
  23. bbox_wh:bbox的长框集合
  24. """
  25. # 最大边缩放到img_size
  26. img_wh = np.array(img_wh, dtype=np.float32)
  27. shapes = (img_size * img_wh / img_wh).max(1, keepdims=True)
  28. wh0 = np.concatenate([l * s for s, l in zip(shapes, bbox_wh)]) # wh
  29. i = (wh0 < 3.0).any(1).sum()
  30. if i:
  31. print(f'WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.')
  32. wh = wh0[(wh0 >= 2.0).any(1)] # 只保留wh都大于等于2个像素的box
  33. # k_means 聚类计算anchor
  34. k = k_means(wh, n, use_iou=True, use_pp=False)
  35. k = k[np.argsort(k.prod(1))] # sort small to large
  36. f, bpr = anchor_fitness(k, wh, thr)
  37. print("kmeans: " + " ".join([f"[{int(i[0])}, {int(i[1])}]" for i in k]))
  38. print(f"fitness: {f:.5f}, best possible recall: {bpr:.5f}")
  39. # YOLOV5改进遗传算法
  40. npr = np.random
  41. f, sh, mp, s = anchor_fitness(k, wh, thr)[0], k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma
  42. pbar = tqdm(range(gen), desc=f'Evolving anchors with Genetic Algorithm:') # progress bar
  43. for _ in pbar:
  44. v = np.ones(sh)
  45. while (v == 1).all(): # mutate until a change occurs (prevent duplicates)
  46. v = ((npr.random(sh) < mp) * random.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0)
  47. kg = (k.copy() * v).clip(min=2.0)
  48. fg, bpr = anchor_fitness(kg, wh, thr)
  49. if fg > f:
  50. f, k = fg, kg.copy()
  51. pbar.desc = f'Evolving anchors with Genetic Algorithm: fitness = {f:.4f}'
  52. # 按面积排序
  53. k = k[np.argsort(k.prod(1))] # sort small to large
  54. print("genetic: " + " ".join([f"[{int(i[0])}, {int(i[1])}]" for i in k]))
  55. print(f"fitness: {f:.5f}, best possible recall: {bpr:.5f}")
  56. auto_anchor(img_size=416, n=9, thr=0.25, gen=1000, img_wh=img_wh, bbox_wh=bbox_wh)

如果有兴趣代码细节的,可以看里面的注释,如果还有不懂的,可以私信我交流。

最后计算出来的结果如下,可以看到计算出来的anchor的是长方型的,这是因为我的bbox中长方型的anchor居多,符合我的预期。我们只需要把下面的anchor,替换掉默认的anchor即可!

 

最后说明一下,用聚类算法算出来的anchor并不一定比初始值即coco上的anchor要好,原因是目标检测大部分基于迁移学习,backbone网络的训练参数是基于coco上的anchor学习的,所以其实大部分情况用这个聚类效果并没有直接使用coco上的好!!,而且聚类效果跟数据集的数量有很大关系,一两千张图片,聚类出来效果可能不会很好



总结

整个算法思路其实不难,但代码有一些冗余和长,主要也是结合自己在学习和使用过程中,发现很多博主没有说明白如何使用这些代码。

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/从前慢现在也慢/article/detail/812963
推荐阅读
相关标签
  

闽ICP备14008679号