add color_balls

color_balls/main.py

@@ -0,0 +1,82 @@
+from recognizer import recognize, triple, quadriple
+import cv2
+
+red = (52, 28, 216)
+green = (71, 143, 31)
+blue = (143, 62, 1)
+yellow = (70, 223, 232)
+colors = [red, green, blue, yellow]
+
+# Test 1
+image = cv2.imread('color_balls/pictures/test1.jpg')
+
+result, image_result = recognize(image, triple, [yellow, blue])
+print("Oreder right?: ", result)
+
+cv2.namedWindow('Image', cv2.WINDOW_NORMAL)
+cv2.imshow('Image', image_result)
+
+while cv2.waitKey(1) != ord('q'):
+    continue
+
+cv2.destroyAllWindows()
+
+# Test 2
+
+image = cv2.imread('color_balls/pictures/test2.jpg')
+
+result, image_result = recognize(image, triple, [blue, yellow, green])
+print("Oreder right?: ", result)
+
+cv2.namedWindow('Image', cv2.WINDOW_NORMAL)
+cv2.imshow('Image', image_result)
+
+while cv2.waitKey(1) != ord('q'):
+    continue
+
+cv2.destroyAllWindows()
+
+# Test 3
+
+image = cv2.imread('color_balls/pictures/test3.jpg')
+
+result, image_result = recognize(image, triple, [blue, yellow, green])
+print("Oreder right?: ", result)
+
+cv2.namedWindow('Image', cv2.WINDOW_NORMAL)
+cv2.imshow('Image', image_result)
+
+while cv2.waitKey(1) != ord('q'):
+    continue
+
+cv2.destroyAllWindows()
+
+# Test 4
+
+image = cv2.imread('color_balls/pictures/test4.jpg')
+
+result, image_result = recognize(image, quadriple, [red, blue, yellow, green])
+print("Oreder right?: ", result)
+
+cv2.namedWindow('Image', cv2.WINDOW_NORMAL)
+cv2.imshow('Image', image_result)
+
+while cv2.waitKey(1) != ord('q'):
+    continue
+
+cv2.destroyAllWindows()
+
+# Test 3
+
+image = cv2.imread('color_balls/pictures/test5.jpg')
+
+result, image_result = recognize(image, quadriple, [blue, yellow, red, green])
+print("Oreder right?: ", result)
+
+cv2.namedWindow('Image', cv2.WINDOW_NORMAL)
+cv2.imshow('Image', image_result)
+
+while cv2.waitKey(1) != ord('q'):
+    continue
+
+cv2.destroyAllWindows()

color_balls/recognizer.py

@@ -0,0 +1,34 @@
+import cv2
+
+
+def triple(centers):
+    return len(centers) == 3 and abs(centers[0][1] - centers[1][1]) < 100 and abs(centers[2][1] - centers[1][1]) < 100
+
+
+def quadriple(centers):
+    return len(centers) == 4 and abs(centers[0][1] - centers[1][1]) < 100 and abs(centers[2][1] - centers[3][1]) < 100 and abs(centers[0][0] - centers[2][0]) < 100
+
+
+def color_distance(color1, color2):
+    return ((color1[0] - color2[0])**2 + (color1[1] - color2[1])**2 + (color1[2] - color2[2])**2)**0.5
+
+
+def recognize(image, rtype, ncolors):
+    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
+    _, hsv = cv2.threshold(hsv[:,:,1], 100, 255, cv2.THRESH_BINARY)
+    hsv = cv2.erode(hsv, None, iterations=5)
+    hsv = cv2.dilate(hsv, None, iterations=10)
+    contours = cv2.findContours(hsv, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
+
+    centers = []
+    colors = []
+    for i,contour in enumerate(contours):
+        rot_react = cv2.minAreaRect(contour)
+        centers.append(rot_react[0])
+        color = tuple(map(int, image[int(centers[-1][1]), int(centers[-1][0]+10)]))
+        colors.append(color)
+        cv2.circle(image, tuple(map(int, rot_react[0])), radius=2, color=color, thickness=10)
+        cv2.drawContours(image, contours, i, color, 10)
+
+    colors_approve = all([color_distance(ncolors[i], colors[i]) < 100 for i in range(len(ncolors))])
+    return colors_approve and rtype(centers), image