Volovikov_CV/falling_ball/cv.py

import cv2
import numpy as np

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


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

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

    cv2.imshow('Image', binary)

    while cv2.waitKey(1) != ord('q'): continue

    labels, labeledScreen = cv2.connectedComponents(binary)
    cnts, hierrache = cv2.findContours(labels, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
    for label in labels:
        print(chose(label))

    # image = image[t_min_loc1[1] + 20:t_min_loc2[1], t_min_loc1[0] + 20:t_min_loc2[0]]

    # hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    # hsv = cv2.GaussianBlur(hsv, (9, 9), 0)
    # _, thresh = cv2.threshold(hsv[:, :, 1], 27, 255, cv2.THRESH_BINARY)
    # cnts, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    # 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)

    # return polys, (t_min_loc2[0] + 20, t_min_loc2[1] + 20)

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