import cv2

cv2.namedWindow("Image", cv2.WINDOW_GUI_NORMAL)

def chose(contour):
    sm = cv2.arcLength(contour, True)
    approx = cv2.approxPolyDP(contour, sm*0.025, True)
    if len(approx) == 4: return 'square'
    elif len(approx) == 8: return 'circle'



def recognize(image, n, video=False):
    image_hsv = cv2.cvtColor(image.copy(), cv2.COLOR_BGR2HSV)
    image_hsv = image_hsv[:,:,1]

    _, thres = cv2.threshold(image_hsv.copy(), 50, 255, cv2.THRESH_BINARY)
    image_hsv = cv2.dilate(thres, None, iterations=3)

    distance_map = cv2.distanceTransform(thres, cv2.DIST_L2, 5)
    ret ,dist_thres= cv2.threshold(distance_map, 0.6 * distance_map.max(), 255, cv2.THRESH_BINARY)

    confuse = cv2.subtract(thres, dist_thres.astype('uint8'))
    ret, markers = cv2.connectedComponents(dist_thres.astype('uint8'))
    markers += 1
    markers[confuse == 255] = 0
    segments = cv2.watershed(image, markers)
    cnts, hierrache = cv2.findContours(segments, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)

    squares = 0
    circles = 0

    color = 100
    for i,contour in enumerate(cnts[:-2]):
        x,y = cv2.boundingRect(contour)[0:2]
        x2,y2 = cv2.boundingRect(cnts[i+1])[0:2]
        if abs(x-x2) < 25 and (y-y2) < 25: continue
        cv2.drawContours(image, cnts[:-2], i, (0, color, 0), 10)
        clase = chose(contour)
        if clase == 'square': squares += 1
        elif clase == 'circle': circles += 1
        color+=10

    cv2.putText(image, f"Circles count = {circles}, quadro count = {squares}", (10, 20), cv2.FONT_HERSHEY_SIMPLEX,
                0.7, (127, 255, 255))
    
    cv2.putText(image, f"Image = {n}", (10, 60), cv2.FONT_HERSHEY_SIMPLEX,
                0.7, (127, 255, 255))
    
    cv2.imshow("Image", image)
    while not video and cv2.waitKey(1) != ord('q'): continue


if __name__ == "__main__":
    image = cv2.imread('/home/jezvcp/Projects/Volovikov_CV/count_objects/out-fotor-20240514183622.png')
    recognize(image, 0)
    cv2.destroyAllWindows()