see recognizer

count_objects/main.py

@@ -1,8 +1,7 @@
 import cv2
 import numpy as np
 import zmq
-from skimage.measure import label
-import math
+from recognizer import recognize
 
 def euclidian(x1, y1, x2, y2):
     return ((x1 - x2)**2 + (y1 - y2)**2)**0.5
@@ -23,89 +22,7 @@ while True:
     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)
+    recognize(imagege, n, True)
 
 cv2.destroyAllWindows()

count_objects/recognizer.py

@@ -0,0 +1,57 @@
+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()