Fixed diff issue by only doing hough lines once and saving that (assuming...

Fixed diff issue by only doing hough lines once and saving that (assuming board will not be moved). Also for loops are much more efficient now

board_detector.py

@@ -230,6 +230,11 @@ def find_board(img):
     corners = [tuple(corner) for corner in corners] # convert to tuples
     corners_sorted = sort_square_grid_coords(corners, unpacked=True)
 
+    return corners_sorted, intersection
+
+def warp_board(img, corners_sorted, intersection):
+    height, width, _ = img.shape
+
     tl, tr, bl, br = corners_sorted
     src = np.float32([list(tl), list(tr), list(bl), list(br)])
     dest = np.float32([[0,0], [width, 0], [0, height], [width, height]])
@@ -330,12 +335,19 @@ def find_pieces(warped_img, sorted_warped_points, reader, img_idx):
     if compare_prev_warped:
         # img_path = os.path.join('game_images', 'warped' + str(img_idx - 1) + '.jpg')
         # prev_warped_img = cv2.imread(img_path)
-        diff = cv2.absdiff(bgr_after, prev_bgr_after)
-        diff_mask = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
+        # diff = cv2.absdiff(bgr_after, prev_bgr_after)
+        # diff_mask = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
+
+        gray_after_mask = np.zeros_like(gray_after)
+        prev_gray_after_mask = np.zeros_like(prev_gray_after)
+        gray_after_mask[gray_after != 0] = 255
+        prev_gray_after_mask[prev_gray_after != 0] = 255
+        # prev_gray_after[prev_gray_after != 0] = 255
+        diff_mask = cv2.absdiff(gray_after_mask, prev_gray_after_mask)
         # display_img([diff_mask])
-        diff_thresh = 70
-        diff_mask[diff_mask > diff_thresh] = 255
-        diff_mask[diff_mask <= diff_thresh] = 0
+        # diff_thresh = 60
+        # diff_mask[diff_mask > diff_thresh] = 255
+        # diff_mask[diff_mask <= diff_thresh] = 0
 
         diff_contours, hierarchy = cv2.findContours(diff_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
         diff_contours_sorted = sorted(diff_contours, reverse=True, key=cv2.contourArea)
@@ -343,7 +355,7 @@ def find_pieces(warped_img, sorted_warped_points, reader, img_idx):
         cv2.drawContours(diff_mask_contours, diff_contours_sorted, -1, (255), thickness=cv2.FILLED)
 
         if show_cv:
-            display_img([bgr_after, prev_bgr_after])
+            display_img([prev_bgr_after, bgr_after])
             display_img([diff_mask, diff_mask_contours])
 
     # if compare_prev_warped:
@@ -354,6 +366,8 @@ def find_pieces(warped_img, sorted_warped_points, reader, img_idx):
     # define color thresholds to use to classify colors later on
     hue_thresh_dict = {'red': (170,190), 'orange':(8,18), 'yellow': (18,44), 'green': (50,70), 'purple': (120,140), 
                        'teal': (80,105), 'pink': (140,170)} # CHANGE
+    
+    letter_dict = {0: 'A', 1: 'B', 2: 'C', 3: 'D', 4: 'E', 5: 'F', 6: 'G', 7: 'H'}
 
     if (show_cv):
         warped_img_pil = cv2_to_pil(warped_img)
@@ -366,7 +380,7 @@ def find_pieces(warped_img, sorted_warped_points, reader, img_idx):
     
     # cur_bounding_box_array = np.empty((8,8), dtype=object)
     # cur_bounding_box_array.fill([])
-    pixel_thresh = 40
+    pixel_thresh = 150 # was 40
     color_grid = []
     # loop through each square of chess board
     for i in range(0,8):
@@ -380,6 +394,12 @@ def find_pieces(warped_img, sorted_warped_points, reader, img_idx):
             bl = sorted_warped_points[i+1][j]
             br = sorted_warped_points[i+1][j+1]
 
+            # straight edges around square
+            x_min = min(tl[0], bl[0])
+            x_max = max(tr[0], br[0])
+            y_min = min(tl[1], tr[1])
+            y_max = max(bl[1], br[1])
+
             # create mask of the square
             height, width, _ = warped_img.shape
             rect_mask = np.zeros((height, width), dtype=np.uint8)
@@ -433,7 +453,7 @@ def find_pieces(warped_img, sorted_warped_points, reader, img_idx):
                         cur_bounding_box = cv2.boundingRect(largest_contour)
                         # cur_bounding_box_array[j][i] = cur_bounding_box
                         
-                    hue = 0
+                    hue_sum = 0
                     if compare_prev_warped:
                         # display_img([filled_contour_mask])
                         # filled_contour_mask = cv2.bitwise_and(filled_contour_mask, filled_contour_mask, mask=diff_mask)
@@ -441,19 +461,33 @@ def find_pieces(warped_img, sorted_warped_points, reader, img_idx):
                         diff_contour_mask = cv2.bitwise_or(filled_contour_mask, prev_filled_contour_mask, mask=diff_mask_contours)
                         # diff_contour_mask = cv2.bitwise_and(diff_contour_mask, diff_contour_mask, mask=rect_mask)
 
-                        num_pixels = 1
-                        num_diff_pixels = 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])):
-                                if diff_contour_mask[y, x, 0] > 0:
-                                    num_diff_pixels += 1
-                                if filled_contour_mask[y, x, 0] > 0 and hsv_after[y, x, 0] != 0:
-                                    num_pixels += 1
-                                    hue += hsv_after[y, x, 0]
-                        avg_hue = hue / num_pixels
+                        # num_pixels = 1
+                        # num_diff_pixels = 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])):
+                        #         if diff_contour_mask[y, x, 0] > 0:
+                        #             num_diff_pixels += 1
+                        #             if filled_contour_mask[y, x, 0] > 0 and hsv_after[y, x, 0] != 0:
+                        #                 num_pixels += 1
+                        #                 hue += hsv_after[y, x, 0]
+                        # avg_hue = hue / num_pixels
+
+                        # diff calculation
+                        diff_rect = diff_contour_mask[y_min:y_max, x_min:x_max, 0]
+                        num_diff_pixels = np.count_nonzero(diff_rect)
                         
                         if num_diff_pixels > pixel_thresh:
                             do_ocr = True
+
+                            # get avg hue
+                            filled_rect = filled_contour_mask[y_min:y_max, x_min:x_max, 0]
+                            hsv_rect = hsv_after[y_min:y_max, x_min:x_max, 0]
+                            both_nonzero_mask = (filled_rect > 0) & (hsv_rect != 0)
+                            num_pixels = np.count_nonzero(both_nonzero_mask)
+                            hue_sum = np.sum(hsv_rect[both_nonzero_mask])
+
+                            avg_hue = hue_sum / num_pixels if num_pixels != 0 else 0
+
                             # only save the contour if it has enough pixels, otherwise erase it
                             if largest_contour is not None and num_pixels < pixel_thresh:
                                 cv2.drawContours(filled_contour_mask, [largest_contour], -1, (0, 0, 0), thickness=cv2.FILLED)
@@ -462,13 +496,22 @@ def find_pieces(warped_img, sorted_warped_points, reader, img_idx):
                             avg_hue = 0
                     else:
                         # loop through all pixels in square and find average hue
-                        num_pixels = 1
-                        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])):
-                                if filled_contour_mask[y, x, 0] > 0 and hsv_after[y, x, 0] != 0:
-                                    num_pixels += 1
-                                    hue += hsv_after[y, x, 0]
-                        avg_hue = hue / num_pixels
+                        # num_pixels = 1
+                        # 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])):
+                        #         if filled_contour_mask[y, x, 0] > 0 and hsv_after[y, x, 0] != 0:
+                        #             num_pixels += 1
+                        #             hue += hsv_after[y, x, 0]
+                        # avg_hue = hue / num_pixels
+
+                        # get avg hue
+                        filled_rect = filled_contour_mask[y_min:y_max, x_min:x_max, 0]
+                        hsv_rect = hsv_after[y_min:y_max, x_min:x_max, 0]
+                        both_nonzero_mask = (filled_rect > 0) & (hsv_rect != 0)
+                        num_pixels = np.count_nonzero(both_nonzero_mask)
+                        hue_sum = np.sum(hsv_rect[both_nonzero_mask])
+
+                        avg_hue = hue_sum / num_pixels if num_pixels != 0 else 0
 
                         # only save the contour if it has enough pixels, otherwise erase it
                         if largest_contour is not None and num_pixels < pixel_thresh:
@@ -494,19 +537,21 @@ def find_pieces(warped_img, sorted_warped_points, reader, img_idx):
                     image=img_to_read,
                     allowlist="PKQRBN", # only want these letters
                     rotation_info=[180], # either rightside up or upside down
-                    text_threshold=0.7,
-                    low_text = 0.7,
+                    text_threshold=0.45,
+                    low_text = 0.1,
                     min_size = 5
                 )
 
                 # get letter if found
                 if len(result) != 0:
                     bound_box, letter, confidence = result[0]
-                    if show_cv:
-                        print(letter, confidence)
                 else:
-                    letter = "X"
+                    letter = "X - remove letter"
+                    confidence = ""
                 if show_cv:
+
+                    coord = letter_dict[j] + str(8 - i)
+                    print(letter, confidence, coord)
                     display_img([img_to_read])
 
                 for color, (lower, upper) in hue_thresh_dict.items(): # for each color threshold
@@ -526,30 +571,33 @@ def find_pieces(warped_img, sorted_warped_points, reader, img_idx):
             else:
                 color_grid[i].append([None, avg_hue, num_pixels, None])
 
-    if show_cv:
+    if show_cv and False:
         warped_img_draw = pil_to_cv2(warped_img_draw._image)
 
         bgr_after_intersections = bgr_after.copy()
         for points in sorted_warped_points:
                 for point in points:
                     cv2.circle(bgr_after_intersections, point, 1, (255, 255, 255), -1)
-        display_img([warped_img_draw, bgr_after_intersections, filled_contour_mask])
+        # display_img([warped_img_draw, bgr_after_intersections, filled_contour_mask])
+        display_img([bgr_after_intersections, filled_contour_mask])
 
     # if compare_prev_warped:
     #     display_img([diff_contour_mask])
 
     # print color_grid. only print when the color is found a lot in the square (> pixel_thresh times)
-    # if show_cv:
+    if show_cv:
         # print("|avg_hue, num_pixels, letter|")
-    # for row in color_grid:
-    #     print("||", end="")
-    #     for color, avg_hue, num_pixels, letter in row:
-    #         if num_pixels > pixel_thresh:
-    #             print(f"{color},   {letter}\t|", end="")
-    #             # print(f"{color},   {letter}\t|", end="")
-    #         else:
-    #             print("\t\t|", end="")
-    #     print("|")
+        for row in color_grid:
+            print("||", end="")
+            for color, avg_hue, num_pixels, letter in row:
+                if letter is not None and letter[0] == "X":
+                    print(f"{int(num_pixels)},   {letter}\t|", end="")
+                elif num_pixels > pixel_thresh:
+                    # print(f"{int(avg_hue)}, {color}, {letter}\t|", end="") # THE GOOD ONE
+                    print(f"{int(num_pixels)},   {letter}\t|", end="")
+                else:
+                    print("\t\t|", end="")
+            print("|")
 
     prev_warped_img = deepcopy(warped_img)
     prev_hsv_img = deepcopy(hsv_img)

game.py

@@ -2,6 +2,7 @@ import argparse
 import chess
 import chess.pgn
 from board_detector import find_board
+from board_detector import warp_board
 from board_detector import find_pieces
 from board_detector import init_global
 from board_detector import display_img
@@ -49,6 +50,10 @@ class ChessGame:
 
         self.img_size = 512
 
+        # store warping info
+        self.corners_sorted = None
+        self.intersection = None
+
         if not skip_camera:
             self.picam2 = Picamera2()
         else:
@@ -175,11 +180,16 @@ class ChessGame:
 
             if (not self.loop):
                 break
+
+        if self.img_idx == 2:
+            self.corners_sorted, self.intersection = find_board(img)
         
         # warp the image based on the lines of the board
-        warped_img, sorted_warped_points = find_board(img)
+        # if self.img_idx == 1:
+        warped_img, sorted_warped_points = warp_board(img, self.corners_sorted, self.intersection)
 
         if (warped_img is None):
+            print("warped_img is None")
             return
         
         # get the pieces based on color thresholding and easyocr
@@ -246,7 +256,7 @@ class ChessGame:
             for j, (color, _, _, letter) in enumerate(row):
                 piece_type = None
                 if letter is not None:
-                    print("letter:", letter)
+                    # print("letter:", letter)
                     letter = letter[0]
                     if letter == "P":
                         piece_type = chess.PAWN
@@ -260,7 +270,7 @@ class ChessGame:
                         piece_type = chess.QUEEN
                     elif letter == "K":
                         piece_type = chess.KING
-                    elif letter == "X":
+                    elif letter == "X": # 'X - remove letter'
                         temp_board.remove_piece_at(chess.square(j, 7 - i))
                         continue
 
@@ -298,6 +308,10 @@ class ChessGame:
             print("Found move!")
         else:
             print("Did NOT find move.")
+            if self.img_idx != 2:
+                print("Detected Board: ")
+                print(board_copy)
+                print("\n")
 
         # print("PREVIOUS BOARD:")
         # print(self.prev_board)