Volovikov_CV/falling_ball/cv.py

import cv2
import numpy as np
import matplotlib.pyplot as plt

# cv2.namedWindow('Image', cv2.WINDOW_NORMAL)

def chose(contour):
    sm = cv2.arcLength(contour, True)
    approx = cv2.approxPolyDP(contour, sm*0.025, True)
    if (len(approx) == 7 or len(approx) == 8) and sm > 60: 
        return 'circle'


def recognize(image, shape):

    # =============== Находим большие синие круглишки =================

    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    binary = np.zeros(hsv.shape[:-1])

    binary[np.any(hsv > 190, axis=2)] = 1

    contours = cv2.findContours(binary.astype('uint8'), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[0]
    corner_left_up, corner_right_down = None, None
    for i,contour in enumerate(contours):
        if chose(contour) == 'circle':
            x = contour[:, 0, 0]
            y = contour[:, 0, 1]
            x_max, x_min, y_max, y_min = x.max(), x.min(), y.max(), y.min()
            if not corner_left_up or corner_left_up[0] < x_min or corner_left_up[1] < y_min:
                corner_left_up = (x_max, y_max)
            elif not corner_right_down or corner_right_down[0] > x_min or corner_right_down[1] > y_min:
                corner_right_down = (x_min, y_min)

            cv2.drawContours(image, contours, i, (0, 255, 0), 2)

    # ================= Бинаризируем полученное поле ===================
    
    image = image[corner_right_down[1]:corner_left_up[1],corner_right_down[0]:corner_left_up[0]]

    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    thresh = cv2.adaptiveThreshold(image,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY_INV,3,2)
    binary = np.zeros(image.shape)
    binary[image < thresh] = 1

    binary = cv2.dilate(binary, None, iterations=2)

    contours = cv2.findContours(binary.astype('uint8'), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[0]

    newbinary = np.zeros(binary.shape)

    for contour in contours:
        cv2.drawContours(newbinary, [contour], -1, 1, -1)

    binary = cv2.erode(newbinary, None, iterations=5)
    binary = cv2.dilate(binary, None, iterations=2)

    # ================= Трансформируем доску в экран ===================

    desk_contour = np.float32([[0, binary.shape[0]], [0,0], 
                                [binary.shape[1], 0], [binary.shape[1], binary.shape[0]]])
    game_contour = np.float32([[0, shape[1]], [0, 0], [shape[0], 0], [shape[0], shape[1]]])

    mask = cv2.getPerspectiveTransform(desk_contour, game_contour)
    binary = cv2.warpPerspective(binary, mask, shape)

    # =================== Находим объекты ===================

    contours = cv2.findContours(binary.astype('uint8'), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[0]

    polys = []
    for contour in contours:
        eps = 0.044 * cv2.arcLength(contour, True)
        rect = cv2.minAreaRect(contour)
        box = cv2.boxPoints(rect)
        box = np.intp(box)
        approx = cv2.approxPolyDP(box, eps, True)
        approx = approx[:, 0, :].tolist()
        polys.append(approx)

    return polys

if  __name__ == "__main__":
    image = cv2.imread('falling_ball/image.png')
    print(recognize(image, (1200, 800)))