Volovikov_CV/count_objects/main.py

import cv2
import numpy as np
import zmq
from skimage.measure import label
import math

def euclidian(x1, y1, x2, y2):
    return ((x1 - x2)**2 + (y1 - y2)**2)**0.5


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

context = zmq.Context()
socket = context.socket(zmq.SUB)
socket.setsockopt(zmq.SUBSCRIBE, b"")
port = 5055
socket.connect(f"tcp://192.168.0.113:{port}")
n = 0

while True:
    bts = socket.recv()
    n += 1
    arr = np.frombuffer(bts, np.uint8)
    imagege = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    key = cv2.waitKey(10)

    if key == ord("q"):
        break

    image = cv2.cvtColor(imagege.copy(), cv2.COLOR_BGR2HSV)
    image = image[:,:,1]

    img = image.copy()
    _, thras = cv2.threshold(img, 70, 255, cv2.THRESH_BINARY)
    image = cv2.threshold(image, 70, 255, cv2.THRESH_BINARY)[1]
    image = cv2.dilate(image, None, iterations=3)

    # sssssss

    # ssssssss


    labeled = label(image)

    cntrs = list(cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0])

    circles = 0
    cubicCount = 0

    for c in cntrs:
        (curr_x, curr_y), r = cv2.minEnclosingCircle(c)
        area = cv2.contourArea(c)
        circleArea = math.pi * r ** 2
        
        if area <= 2000:
            continue

        if area / circleArea  >= 0.9:
            circles += 1
        else:
            cubicCount += 1

    distance_map = cv2.distanceTransform(thras, cv2.DIST_L2, 5)
    ret ,dist_thres= cv2.threshold(distance_map, 0.6 * np.max(distance_map), 255, cv2.THRESH_BINARY)
    # ret, dist_thres = cv2.threshold(distance_map, 0.5, 25, cv2.THRESH_BINARY)
    # ret, markers = cv2.connectedComponents(dist_thres.astype('uint8'))
    # segments = cv2.wetershed(image, markers + 1)
    confuse = cv2.subtract(thras, dist_thres.astype('uint8'))
    ret, markers = cv2.connectedComponents(dist_thres.astype('uint8'))
    markers += 1
    markers[confuse == 255] = 0
    segments = cv2.watershed(imagege, markers)
    cnts, hierrache = cv2.findContours(segments, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)

    circlesA = 0
    cubicCountA = 0

    for i in range(len(cnts)):
        if hierrache[0][i][3] != -1: continue
        c = cnts[i]
        (curr_x, curr_y), r = cv2.minEnclosingCircle(c)
        area = cv2.contourArea(c)
        circleArea = math.pi * r ** 2
        
        if area <= 2000:
            continue

        if area / circleArea  >= 0.9:
            circlesA += 1
        else:
            cubicCountA += 1

    for i in range(len(cnts)):
        if hierrache[0][i][3] == -1:
            cv2.drawContours(imagege, cnts, i, (0, 255, 0), 10)

    cv2.putText(image, f"Circles count = {circles}, quadro count = {cubicCount}", (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)

    cv2.putText(imagege, f"Circles count = {circlesA}, quadro count = {cubicCountA}", (10, 20), cv2.FONT_HERSHEY_SIMPLEX,
                0.7, (127, 255, 255))

    cv2.imshow("Image2", imagege)

cv2.destroyAllWindows()