Volovikov_CV/color_balls/recognizer.py

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