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: return 'circle'


lower = 100
upper = 200


def lower_update(value):
    global lower
    lower = value


def upper_update(value):
    global upper
    upper = value


cv2.createTrackbar("Lower", "Image", lower, 255, lower_update)
cv2.createTrackbar("Upper", "Image", upper, 255, upper_update)


def recognize(image):
    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)

    shape = (corner_left_up[0]-corner_right_down[0], corner_left_up[1]-corner_right_down[1])

    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,upper,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY,3,2)
    binary = np.zeros(image.shape)
    binary[image < thresh] = 1

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

    cv2.imshow('Image', binary)
    
    cv2.waitKey(1)

    # return polys, shape

if  __name__ == "__main__":
    while True:
        image = cv2.imread('falling_ball/image3.png')
        print(recognize(image))