working

falling_ball/cv.py

@@ -2,33 +2,19 @@ import cv2
 import numpy as np
 import matplotlib.pyplot as plt
 
-cv2.namedWindow('Image', cv2.WINDOW_NORMAL)
+# 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'
+    if (len(approx) == 7 or len(approx) == 8) and sm > 60: 
+        return 'circle'
 
 
-lower = 100
-upper = 200
+def recognize(image, shape):
 
+    # =============== Находим большие синие круглишки =================
 
-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])
 
@@ -43,43 +29,58 @@ def recognize(image):
             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:
+            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])
+            cv2.drawContours(image, contours, i, (0, 255, 0), 2)
 
+    # ================= Бинаризируем полученное поле ===================
+    
     image = image[corner_right_down[1]:corner_left_up[1],corner_right_down[0]:corner_left_up[0]]
 
-    # hsv = cv2.cvtColor(image.copy(), cv2.COLOR_BGR2HSV)
-    # hsv = cv2.GaussianBlur(hsv, (1, 1), 0)
-    # _, thresh = cv2.threshold(hsv[:, :, 1], lower, upper, cv2.THRESH_BINARY)
-    # cnts, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
-    # cv2.drawContours(image, cnts, -1, (0, 255, 0), 2)
-
     image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
 
-    thresh = cv2.threshold(image, lower, upper, cv2.THRESH_BINARY)[0]
-    # thresh = cv2.adaptiveThreshold(image,upper,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY,3,2)
+    thresh = cv2.adaptiveThreshold(image,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY_INV,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)
+    binary = cv2.dilate(binary, None, iterations=2)
+
+    contours = cv2.findContours(binary.astype('uint8'), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[0]
+
+    newbinary = np.zeros(binary.shape)
+
+    for contour in contours:
+        cv2.drawContours(newbinary, [contour], -1, 1, -1)
+
+    binary = cv2.erode(newbinary, None, iterations=5)
+    binary = cv2.dilate(binary, None, iterations=2)
+
+    # ================= Трансформируем доску в экран ===================
+
+    desk_contour = np.float32([[0, binary.shape[0]], [0,0], 
+                                [binary.shape[1], 0], [binary.shape[1], binary.shape[0]]])
+    game_contour = np.float32([[0, shape[1]], [0, 0], [shape[0], 0], [shape[0], shape[1]]])
+
+    mask = cv2.getPerspectiveTransform(desk_contour, game_contour)
+    binary = cv2.warpPerspective(binary, mask, shape)
+
+    # =================== Находим объекты ===================
+
+    contours = cv2.findContours(binary.astype('uint8'), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[0]
+
+    polys = []
+    for contour in contours:
+        eps = 0.044 * cv2.arcLength(contour, True)
+        rect = cv2.minAreaRect(contour)
+        box = cv2.boxPoints(rect)
+        box = np.intp(box)
+        approx = cv2.approxPolyDP(box, eps, True)
+        approx = approx[:, 0, :].tolist()
+        polys.append(approx)
 
-    # return polys, shape
+    return polys
 
 if  __name__ == "__main__":
-    while True:
-        image = cv2.imread('falling_ball/image.png')
-        print(recognize(image))
+    image = cv2.imread('falling_ball/image.png')
+    print(recognize(image, (1200, 800)))

falling_ball/main.py

@@ -3,16 +3,20 @@ import pymunk.pygame_util
 from cv import recognize
 import cv2
 
-image = cv2.imread('project/image.png')
-polys, shape = recognize(image)
+# image = cv2.imread('falling_ball/image.png')
+cam = cv2.VideoCapture(0)
 
-print(shape)
-RES = WIDTH, HEIGHT = shape[0], shape[1]
+cap, image = cam.read()
+
+RES = WIDTH, HEIGHT = 1280, 800
+BALL_RADIUS = min(RES) / 30
 FPS = 60
 
+polys = recognize(image, RES)
+
 pygame.init()
 
-surface = pygame.display.set_mode(RES)
+surface = pygame.display.set_mode((0, 0), pygame.FULLSCREEN, display=1)
 clock = pygame.time.Clock()
 
 draw_options = pymunk.pygame_util.DrawOptions(surface)
@@ -20,7 +24,7 @@ space = pymunk.Space()
 space.gravity = 0, 2000
 
 def create_ball(space, pos: tuple[int ,int]):
-    ball_mass, ball_radius = 10, 5
+    ball_mass, ball_radius = 10, BALL_RADIUS
     ball_moment = pymunk.moment_for_circle(ball_mass, 0, ball_radius)
     ball_body = pymunk.Body(ball_mass, ball_moment)
     ball_body.position = pos
@@ -37,20 +41,20 @@ def draw_by_lines(space, lines=list[tuple[int, int]]):
 for poly in polys:
     draw_by_lines(space, poly)
 
-circle_position = (350, 20)
+circle_position = (350, BALL_RADIUS+1)
 
 while True:
     surface.fill(pygame.Color('white'))
 
-    pygame.draw.circle(surface, pygame.Color('BLUE'), circle_position, 5)
+    pygame.draw.circle(surface, pygame.Color('BLUE'), circle_position, BALL_RADIUS)
 
     for i in pygame.event.get():
         if i.type == pygame.KEYDOWN:
             if i.key == pygame.K_SPACE:
                 create_ball(space, circle_position)
-            if i.key == pygame.K_LEFT:
-                circle_position = (circle_position[0] + 10, circle_position[1])
             if i.key == pygame.K_RIGHT:
+                circle_position = (circle_position[0] + 10, circle_position[1])
+            if i.key == pygame.K_LEFT:
                 circle_position = (circle_position[0] - 10, circle_position[1])
         if i.type == pygame.QUIT:
             exit()