diff --git a/kmeans.py b/kmeans.py new file mode 100644 index 0000000000000000000000000000000000000000..a453925962f96779917b237666ae002667824037 --- /dev/null +++ b/kmeans.py @@ -0,0 +1,85 @@ +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") + + 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(boxes, [0, 1], axis=1) + + +def kmeans(boxes, k, iterations=10): + """ + 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 iterations: number of iterations + :return: numpy array of shape (k, 2) + """ + rows = boxes.shape[0] + + distances = np.empty((rows, k)) + + result = [0.0, None] + for i in range(0, iterations): + # the Forgy method will fail if the whole array contains the same rows + clusters = boxes[np.random.choice(rows, k, replace=False)] + + tmp = [0.0, clusters] + while True: + for row in range(rows): + distances[row] = 1 - iou(boxes[row], clusters) + + nearest_clusters = np.argmin(distances, axis=1) + + for cluster in range(k): + clusters[cluster] = np.mean(boxes[nearest_clusters == cluster], axis=0) + + # improve this + avg = avg_iou(boxes, clusters) + if avg > tmp[0]: + tmp = [avg, clusters] + else: + break + + if tmp[0] > result[0]: + result = tmp + + return result[1] diff --git a/tests/test_basic.py b/tests/test_basic.py new file mode 100644 index 0000000000000000000000000000000000000000..358224e0ecca0579fb1d6913c9b64ca871d08d70 --- /dev/null +++ b/tests/test_basic.py @@ -0,0 +1,61 @@ +from unittest import TestCase + +import numpy as np + +from kmeans import iou, avg_iou, kmeans + + +class TestBasic(TestCase): + def test_iou_100(self): + self.assertEqual(iou([200, 200], np.array([[200, 200]])), 1.) + + def test_iou_50(self): + self.assertEqual(iou([200, 200], np.array([[100, 200]])), .5) + self.assertEqual(iou([200, 200], np.array([[200, 100]])), .5) + + def test_iou_75(self): + self.assertEqual(iou([200, 200], np.array([[150, 200]])), .75) + self.assertEqual(iou([200, 200], np.array([[200, 150]])), .75) + + def test_iou_20(self): + self.assertEqual(iou([183, 73], np.array([[73, 36.6]])), .2) + self.assertEqual(iou([183, 73], np.array([[36.6, 73]])), .2) + + def test_iou_multiple(self): + a = np.array([[200, 200], [100, 200], [200, 100], [150, 200], [200, 150]]) + b = np.array([1., 0.5, 0.5, 0.75, 0.75]) + self.assertTrue((iou([200, 200], a) == b).all()) + + def test_iou_0(self): + self.assertRaises(ValueError, iou, [100, 100], np.array([[0, 0]])) + self.assertRaises(ValueError, iou, [0, 0], np.array([[100, 100]])) + self.assertRaises(ValueError, iou, [0, 0], np.array([[0, 0]])) + self.assertRaises(ValueError, iou, [100, 0], np.array([[100, 100]])) + self.assertRaises(ValueError, iou, [0, 100], np.array([[100, 100]])) + + def test_avg_iou_simple(self): + self.assertEqual(avg_iou(np.array([[200, 200]]), np.array([[200, 200]])), 1.) + self.assertEqual(avg_iou(np.array([[200, 200]]), np.array([[100, 200]])), .5) + self.assertEqual(avg_iou(np.array([[200, 200]]), np.array([[200, 100]])), .5) + + def test_avg_iou_multiple(self): + a = np.array([[200, 200], [100, 200], [200, 100], [150, 200], [200, 150]]) + b = np.array([[200, 200], [100, 200], [200, 100], [150, 200], [200, 150]]) + self.assertEqual(avg_iou(a, b), 1.) + + c = np.array([[200, 200], [100, 200]]) + self.assertEqual(avg_iou(a, c), np.mean([1., 1., .5, .75, .75])) + + def test_kmeans_simple(self): + a = np.array([[200, 200]]) + b = np.array([[200, 200]]) + self.assertTrue((kmeans(a, 1) == b).all()) + + def test_kmeans_multiple(self): + a = np.array([[200, 200], [100, 200], [300, 200]]) + b = [[100, 200], [200, 200], [300, 200]] + + out = kmeans(a, 3).tolist() + out.sort() + + self.assertTrue((out == b)) diff --git a/tests/test_rectangles.py b/tests/test_rectangles.py new file mode 100644 index 0000000000000000000000000000000000000000..b3d5a121db33ca9eccd0ec30d9c426f7b6ffb6a9 --- /dev/null +++ b/tests/test_rectangles.py @@ -0,0 +1,32 @@ +from unittest import TestCase + +import numpy as np + +from kmeans import kmeans + + +class TestBasic(TestCase): + def gen_shape(self, width, height, amount): + boxes = np.empty((amount, 2)) + for i in range(0, amount): + x0 = np.random.randint(100, 1000) + y0 = np.random.randint(100, 1000) + x1 = x0 + width + y1 = y0 + height + boxes[i] = (x1 - x0, y1 - y0) + return boxes + + def test_kmeans_shift(self): + boxes = self.gen_shape(1000, 1000, 100) + self.assertTrue((kmeans(boxes, 1) == [[1000, 1000]]).all()) + + def test_kmeans_shift2(self): + boxes1 = self.gen_shape(1000, 1000, 1) + boxes2 = self.gen_shape(100, 3000, 1) + together = np.concatenate((boxes1, boxes2), axis=0) + + out = kmeans(together, 2) + + res1 = np.array([[100, 3000], [1000, 1000]]) + res2 = np.array([[1000, 1000], [100, 3000]]) + self.assertTrue(np.array_equal(out, res1) or np.array_equal(out, res2)) diff --git a/tests/test_voc2007.py b/tests/test_voc2007.py new file mode 100644 index 0000000000000000000000000000000000000000..2474051026b24d3ab1ad7ec19b4dc341c9b2af51 --- /dev/null +++ b/tests/test_voc2007.py @@ -0,0 +1,45 @@ +import glob +import xml.etree.ElementTree as ET +from unittest import TestCase + +import numpy as np + +from kmeans import kmeans, avg_iou + +ANNOTATIONS_PATH = "Annotations" + + +class TestVoc2007(TestCase): + def __load_dataset(self): + dataset = [] + for xml_file in glob.glob("{}/*xml".format(ANNOTATIONS_PATH)): + tree = ET.parse(xml_file) + + height = int(tree.findtext("./size/height")) + width = int(tree.findtext("./size/width")) + + for obj in tree.iter("object"): + xmin = int(obj.findtext("bndbox/xmin")) / width + ymin = int(obj.findtext("bndbox/ymin")) / height + xmax = int(obj.findtext("bndbox/xmax")) / width + ymax = int(obj.findtext("bndbox/ymax")) / height + + dataset.append([xmax - xmin, ymax - ymin]) + + return np.array(dataset) + + def test_kmeans_5(self): + dataset = self.__load_dataset() + out = kmeans(dataset, 5, iterations=50) + percentage = avg_iou(dataset, out) + print(percentage) + + np.testing.assert_almost_equal(percentage, 0.61, decimal=2) + + def test_kmeans_9(self): + dataset = self.__load_dataset() + out = kmeans(dataset, 9, iterations=50) + percentage = avg_iou(dataset, out) + print(percentage) + + np.testing.assert_almost_equal(percentage, 0.672, decimal=2)