Saving progress: we have color filtering, although the thresholds are...

Saving progress: we have color filtering, although the thresholds are completely wrong for the exact colors (need to fix). Will show the warped image with the corresponding filtered and warped hsv and then print to the console the colors it finds. Next steps are fixing the thresholds and taking the color_grid and turning it into FEN

board_detector.py

@@ -83,13 +83,14 @@ def filter_lines(lines, min_distance):
             
     return filtered_lines
 
-def detect_board(img):
+def find_board_and_pieces(img):
     vertical_lines, horizontal_lines = find_longest_lines(img)
     print("# of Vertical:",len(vertical_lines))
     print("# of Horizontal:",len(horizontal_lines))
 
-    height, width, _ = img.shape
-    black_img = np.zeros((height, width), dtype=np.uint8)
+    if (len(vertical_lines) != 9 or len(horizontal_lines) != 9):
+        print("Error: Grid does not match expected 9x9")
+        return
 
     # create bitmasks for vert and horiz so we can get lines and intersections
     height, width, _ = img.shape
@@ -104,9 +105,10 @@ def detect_board(img):
         x1, y1, x2, y2 = line[0]
         cv2.line(horizontal_mask, (x1, y1), (x2, y2), (255), 2)
 
+    # get lines and intersections of grid and corresponding contours
     intersection = cv2.bitwise_and(vertical_mask, horizontal_mask)
     board_lines = cv2.bitwise_or(vertical_mask, horizontal_mask)
-        
+
     contours, hierarchy = cv2.findContours(board_lines, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
     intersections, hierarchy = cv2.findContours(intersection, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
 
@@ -131,7 +133,7 @@ def detect_board(img):
             for point in points:
                 cv2.circle(board_lines_img, point, 5, (255 - (3 * i), (i % 9) * 28, 3 * i), -1)
                 i += 1
-                print(i)
+                # print(i)
         cv2.imshow('Lines of Board', board_lines_img)
         cv2.waitKey(0)
         cv2.destroyAllWindows()
@@ -178,15 +180,6 @@ def detect_board(img):
     print(corners)
     # corners.sort(key=lambda coord: (coord[0], coord[1])) # sort coords. goes from bottom left clockwise to bottom right - DIDN'T WORK
     corners_sorted = sort_square_grid_coords(corners, unpacked=True)
-    print(corners_sorted)
-
-    # corners_img = img.copy()
-    # for i, corner in enumerate(corners_sorted):
-    #     cv2.circle(corners_img, corner, 5, (60 * i, 60 * i, 60 * i), -1)
-    # if (show_cv):
-    #     cv2.imshow('Canny Filter', corners_img)
-    #     cv2.waitKey(0)
-    #     cv2.destroyAllWindows()
 
     tl = corners_sorted[0]
     tr = corners_sorted[1]
@@ -201,10 +194,26 @@ def detect_board(img):
 
     M = cv2.getPerspectiveTransform(src, dest)
     Minv = cv2.getPerspectiveTransform(dest, src)
-    warped_ip = img.copy()
-    warped_ip = cv2.drawContours(warped_ip, intersections, -1, (0, 0, 255), 2)
+    warped_ip = intersection.copy() # warped intersection points
+    # warped_ip = cv2.drawContours(warped_ip, intersections, -1, (0, 0, 255), 2)
     warped_ip = cv2.warpPerspective(np.uint8(warped_ip), M, (width, height))
 
+    # ----
+    intersections, hierarchy = cv2.findContours(warped_ip, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+
+    # find midpoints of intersection contours (this gives us exact coordinates of the corners of the grid)
+    intersection_points = []
+    for contour in intersections:
+        M = cv2.moments(contour)
+        if (M["m00"] != 0):
+            midpoint_x = int(M["m10"] / M["m00"])
+            midpoint_y = int(M["m01"] / M["m00"])
+            intersection_points.append((midpoint_x, midpoint_y))
+
+    # sort the coordinates from left to right then top to bottom
+    sorted_intersection_points = sort_square_grid_coords(intersection_points, unpacked=False)
+    # ----
+
     if (show_cv):
         contours_img = img.copy()
         # for i in range(63):
@@ -220,24 +229,128 @@ def detect_board(img):
         cv2.waitKey(0)
         cv2.destroyAllWindows()
 
+        # cv2.imshow('Warped', warped_img)
+        # cv2.waitKey(0)
+        # cv2.destroyAllWindows()
+
+    # COLOR / PIECE DETECTION ------------------------------------------------
+
+    hsv_img = cv2.cvtColor(warped_img, cv2.COLOR_BGR2HSV)
+
+    hsv_mask_sat = cv2.inRange(hsv_img[:,:,1], 100, 255) # saturation mask
+    hsv_mask_bright = cv2.inRange(hsv_img[:,:,2], 100, 255) # brightness mask
+
+    # Combine the saturation and brightness masks
+    hsv_mask = cv2.bitwise_and(hsv_mask_sat, hsv_mask_bright)
+
+    # Apply the mask to the entire HSV image
+    hsv_after = cv2.bitwise_and(hsv_img, hsv_img, mask=hsv_mask)
+
+    # contours, _ = cv2.findContours(hsv_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    # filtered_contours = [cnt for cnt in contours if cv2.contourArea(cnt) > 20]
+    # contour_mask = np.zeros_like(hsv_mask)
+    # cv2.drawContours(contour_mask, filtered_contours, -1, (255), thickness=cv2.FILLED)
+
+
+
+
+
+    # create color histogram for each square and find candidate color, if any
+    hue_thresh_dict = {'red': (170,190), 'orange':(11,30), 'yellow': (31,40), 'green': (50,70), 'blue': (110,130), 'teal': (90,109),
+                       'pink': (140,160)} 
+    # ^ note: 190 is above max hue but it should wrap around and start from the beginning again (or we'll just % by 180 ourselves)
+    count = 0
+    color_grid = []
+    for i in range(0,8):
+        color_grid.append([])
+        for j in range(0,8):
+            tl = sorted_intersection_points[i][j]
+            # cv2.circle(test_img, tl, 5, (0, 255, 255), -1)
+            # if (show_cv):
+            #     cv2.imshow('test_img', test_img)
+            #     cv2.waitKey(0)
+            #     cv2.destroyAllWindows()
+
+
+            tr = sorted_intersection_points[i][j+1]
+            # cv2.circle(test_img, tr, 5, (0, 255, 255), -1)
+            # if (show_cv):
+            #     cv2.imshow('test_img', test_img)
+            #     cv2.waitKey(0)
+            #     cv2.destroyAllWindows()
+
+            bl = sorted_intersection_points[i+1][j]
+            # cv2.circle(test_img, bl, 5, (0, 255, 255), -1)
+            # if (show_cv):
+            #     cv2.imshow('test_img', test_img)
+            #     cv2.waitKey(0)
+            #     cv2.destroyAllWindows()
+
+            br = sorted_intersection_points[i+1][j+1]
+            # cv2.circle(test_img, br, 5, (0, 255, 255), -1)
+            # if (show_cv):
+            #     cv2.imshow('hsv_img', hsv_img)
+            #     cv2.waitKey(0)
+            #     cv2.destroyAllWindows()
+
+            height, width, _ = img.shape
+            mask = np.zeros((height, width), dtype=np.uint8)
+            poly = np.array([[tl, tr, br, bl]], dtype=np.int32)
+            cv2.fillPoly(mask, poly, 255)
+            num_pixels = 1
+            hue = 0
+            for x in range(min(tl[0],bl[0]), max(tr[0],br[0])):
+                for y in range(min(tl[1],tr[1]), max(bl[1],br[1])):
+                    # print(hsv_after[y, x, 2])
+                    if mask[y,x] == 255 and hsv_after[y, x, 0] != 0:
+                        # print(hsv_img[y, x, 2])
+                        num_pixels += 1
+                        hue += hsv_after[y, x, 0]
+            avg_hue = hue / num_pixels
+            print(count, avg_hue)
+            piece_found = False
+            for color, (lower, upper) in hue_thresh_dict.items():
+                if avg_hue != 0:
+                    if lower <= avg_hue <= upper:
+                        color_grid[i].append((color, avg_hue))
+                        piece_found = True
+                        break # should only do this once
+            
+            if piece_found == False:
+                color_grid[i].append((None,0))
+
+            count += 1
+
+    for row in color_grid:
+        for tup in row:
+            if tup[0] is None:
+                print("\t\t|", end="")
+            else:
+                print(tup[0],int(tup[1]), "\t|", end="")
+        print("")
+
+    if show_cv:
         cv2.imshow('Warped', warped_img)
+
+        hsv_after_intersections = hsv_after.copy()
+        for points in sorted_intersection_points:
+                for point in points:
+                    cv2.circle(hsv_after_intersections, point, 1, (255, 255, 255), -1)
+        cv2.imshow('hsv_after_intersections', hsv_after_intersections)
         cv2.waitKey(0)
         cv2.destroyAllWindows()
 
-        # cv2.imshow('Warped', warped_ip)
-        # cv2.waitKey(0)
-        # cv2.destroyAllWindows()
-
 def sort_square_grid_coords(coordinates, unpacked):
+    # this function assumes there are a perfect square amount of coordinates
     sqrt_len = int(math.sqrt(len(coordinates)))
     sorted_coords = sorted(coordinates, key=lambda coord: coord[1]) # first sort by y values
     # then group rows of the square (for example, 9x9 grid would be 81 coordinates so split into 9 arrays of 9)
-    groups = [sorted_coords[i:i+sqrt_len] for i in range(0, len(sorted_coords), sqrt_len)]
-    for group in groups:
-        group.sort(key=lambda coord: coord[0]) # now sort each row by x
+    rows = [sorted_coords[i:i+sqrt_len] for i in range(0, len(sorted_coords), sqrt_len)]
+    for row in rows:
+        row.sort(key=lambda coord: coord[0]) # now sort each row by x
 
     if (unpacked == False):
-        return groups
+        return rows
     
-    collapsed_groups = [coord for sublist in groups for coord in sublist] # unpack/collapse groups to just be an array of tuples
-    return collapsed_groups
\ No newline at end of file
+    collapsed = [coord for row in rows for coord in row] # unpack/collapse groups to just be an array of tuples
+    return collapsed
\ No newline at end of file

game.py

 import argparse
 import chess
 import chess.pgn
-from board_detector import detect_board
+from board_detector import find_board_and_pieces
 from board_detector import init_show_cv
 import color_analyzer
 import move_translator
@@ -31,7 +31,7 @@ class ChessGame:
         h,w,c = orig_img.shape
         print(h, w)
         cropped_img = orig_img
-        # cropped_img = orig_img[0:h, int(w/2 - h/2 - 60):int(w/2 + h/2 + 100)]
+        cropped_img = orig_img[0:h, int(w/2 - h/2 - 60):int(w/2 + h/2 + 100)]
         img = cv2.resize(cropped_img, (512, 512))
         # img = cv2.resize(cropped_img, (int(w/5), int(h/5)))
         # img = orig_img
@@ -41,7 +41,7 @@ class ChessGame:
             cv2.waitKey(0)
             cv2.destroyAllWindows()
 
-        detect_board(img)
+        find_board_and_pieces(img)
 
         while(1): # game loop
             self.player_turn()