82 lines
2.8 KiB
Python
82 lines
2.8 KiB
Python
import numpy as np
|
||
|
||
def iou(box, clusters):
|
||
"""
|
||
Calculates the Intersection over Union (IoU) between a box and k clusters.
|
||
:param box: tuple or array, shifted to the origin (i. e. width and height)
|
||
:param clusters: numpy array of shape (k, 2) where k is the number of clusters
|
||
:return: numpy array of shape (k, 0) where k is the number of clusters
|
||
"""
|
||
x = np.minimum(clusters[:, 0], box[0])
|
||
y = np.minimum(clusters[:, 1], box[1])
|
||
if np.count_nonzero(x == 0) > 0 or np.count_nonzero(y == 0) > 0:
|
||
raise ValueError("Box has no area") # 如果报这个错，可以把这行改成pass即可
|
||
|
||
intersection = x * y
|
||
box_area = box[0] * box[1]
|
||
cluster_area = clusters[:, 0] * clusters[:, 1]
|
||
|
||
iou_ = intersection / (box_area + cluster_area - intersection)
|
||
|
||
return iou_
|
||
|
||
def avg_iou(boxes, clusters):
|
||
"""
|
||
Calculates the average Intersection over Union (IoU) between a numpy array of boxes and k clusters.
|
||
:param boxes: numpy array of shape (r, 2), where r is the number of rows
|
||
:param clusters: numpy array of shape (k, 2) where k is the number of clusters
|
||
:return: average IoU as a single float
|
||
"""
|
||
return np.mean([np.max(iou(boxes[i], clusters)) for i in range(boxes.shape[0])])
|
||
|
||
def translate_boxes(boxes):
|
||
"""
|
||
Translates all the boxes to the origin.
|
||
:param boxes: numpy array of shape (r, 4)
|
||
:return: numpy array of shape (r, 2)
|
||
"""
|
||
new_boxes = boxes.copy()
|
||
for row in range(new_boxes.shape[0]):
|
||
new_boxes[row][2] = np.abs(new_boxes[row][2] - new_boxes[row][0])
|
||
new_boxes[row][3] = np.abs(new_boxes[row][3] - new_boxes[row][1])
|
||
return np.delete(new_boxes, [0, 1], axis=1)
|
||
|
||
|
||
def kmeans(boxes, k, dist=np.median):
|
||
"""
|
||
Calculates k-means clustering with the Intersection over Union (IoU) metric.
|
||
:param boxes: numpy array of shape (r, 2), where r is the number of rows
|
||
:param k: number of clusters
|
||
:param dist: distance function
|
||
:return: numpy array of shape (k, 2)
|
||
"""
|
||
rows = boxes.shape[0]
|
||
|
||
distances = np.empty((rows, k))
|
||
last_clusters = np.zeros((rows,))
|
||
|
||
np.random.seed()
|
||
|
||
# the Forgy method will fail if the whole array contains the same rows
|
||
clusters = boxes[np.random.choice(rows, k, replace=False)]
|
||
|
||
while True:
|
||
for row in range(rows):
|
||
distances[row] = 1 - iou(boxes[row], clusters)
|
||
|
||
nearest_clusters = np.argmin(distances, axis=1)
|
||
|
||
if (last_clusters == nearest_clusters).all():
|
||
break
|
||
|
||
for cluster in range(k):
|
||
clusters[cluster] = dist(boxes[nearest_clusters == cluster], axis=0)
|
||
|
||
last_clusters = nearest_clusters
|
||
|
||
return clusters
|
||
|
||
if __name__ == '__main__':
|
||
a = np.array([[1, 2, 3, 4], [5, 7, 6, 8]])
|
||
print(translate_boxes(a))
|