diff --git a/board_detector.py b/board_detector.py
index d63851826ff39d724aff67371d2dec994faae9a7..c5f1b46d57b4124bfa2529dee906b4334af00838 100644
--- a/board_detector.py
+++ b/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]])
@@ -327,15 +332,23 @@ def find_pieces(warped_img, sorted_warped_points, reader, img_idx):
gray_after = cv2.cvtColor(bgr_after, cv2.COLOR_BGR2GRAY)
diff_mask_contours = np.zeros_like(gray_after)
+ drawn_warped_img = deepcopy(warped_img)
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 +356,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 +367,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 +381,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 +395,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 +454,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 +462,35 @@ 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
+
+ cv2.drawContours(drawn_warped_img, diff_contours_sorted, -1, (0,0,255), thickness=cv2.FILLED)
+
+ # 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 +499,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 +540,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 +574,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)
@@ -565,7 +616,7 @@ def find_pieces(warped_img, sorted_warped_points, reader, img_idx):
prev_largest_contour = deepcopy(largest_contour)
# prev_bounding_box_array = cur_bounding_box_array.copy()
- return color_grid
+ return color_grid, drawn_warped_img
def sort_square_grid_coords(coordinates, unpacked):
# this function assumes there are a perfect square amount of coordinates
diff --git a/game.py b/game.py
index 30e72ccf1b5c8b5bfa160628cee93cc70b51a963..0b406a84fa93aa61bac27a493c459f91b72ff0dc 100644
--- a/game.py
+++ b/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
@@ -11,6 +12,7 @@ import cv2
import os
import time
import easyocr
+import sys
# was having issues installing picamera2 on my PC
# if the raspi runs this code (with picamera2 installed), skip_camera will be false
@@ -21,15 +23,17 @@ try:
except ImportError:
skip_camera = True
-example_fens = {"gunnar_enpass": "r1bq1r2/pp2n1p1/4N2k/3pPp1P/1b1n2Q1/2N5/PP3PP1/R1B1K2R b KQ - 0 14"}
+example_fens = {"gunnar_enpass": "r1bq1r2/pp2n1p1/4N2k/3pPp1P/1b1n2Q1/2N5/PP3PP1/R1B1K2R b KQ - 0 14",
+ "piece_test_castling1-": "r3k3/1b1p1p2/p3pK2/5qp1/2P5/P3P3/1P2B3/R1R1Q3 w q - 4 30"} # prins
class ChessGame:
- def __init__(self, difficulty, color_scheme, show_cv, show_cam, loop, img_idx = 1, test_img = None, save_img_as = None):
+ def __init__(self, difficulty, color_scheme, show_cv, show_board, show_cam, loop, img_idx = 1, test_img = None, save_img_as = None):
self.board = chess.Board()
self.prev_board = chess.Board()
self.difficulty = difficulty
self.color_scheme = color_scheme
self.show_cv = show_cv
+ self.show_board = show_board
self.show_cam = show_cam
self.test_img = test_img
self.loop = loop
@@ -49,6 +53,11 @@ class ChessGame:
self.img_size = 512
+ # store warping info
+ self.corners_sorted = None
+ self.intersection = None
+ self.drawn_warped_img = None
+
if not skip_camera:
self.picam2 = Picamera2()
else:
@@ -91,17 +100,26 @@ class ChessGame:
self.board = chess.Board()
# self.img_idx += 1
+ if self.show_board and self.drawn_warped_img is not None:
+ display_img([self.drawn_warped_img])
+
while(1): # game loop
# if self.board.turn == chess.WHITE:
self.player_turn()
+ if self.show_board and self.drawn_warped_img is not None:
+ display_img([self.drawn_warped_img])
+
if self.check_game_over():
break
# if self.board.turn == chess.BLACK:
self.ai_turn()
+ if self.show_board and self.drawn_warped_img is not None:
+ display_img([self.drawn_warped_img])
+
if self.check_game_over():
break
@@ -149,6 +167,11 @@ class ChessGame:
if (self.test_img):
img_txt = self.test_img + str(self.img_idx) + '.jpg'
img_path = os.path.join('ocr_test_images', img_txt)
+ if not os.path.exists(img_path):
+ print("\n-------------------")
+ print("Out of test images.")
+ print("-------------------\n")
+ sys.exit()
orig_img = cv2.imread(img_path)
self.img_idx += 1
else:
@@ -175,15 +198,20 @@ 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
- color_grid = find_pieces(warped_img, sorted_warped_points, self.reader, self.img_idx - 1)
+ color_grid, self.drawn_warped_img = find_pieces(warped_img, sorted_warped_points, self.reader, self.img_idx - 1)
# convert color_grid to board, which we can get fen string from
if self.color_scheme == 'p/y':
@@ -201,6 +229,7 @@ class ChessGame:
print(self.board)
print(self.board.fen())
+
def check_game_over(self):
if self.board.is_checkmate():
print("\n----------------------")
@@ -246,7 +275,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 +289,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
@@ -284,20 +313,18 @@ class ChessGame:
board_copy = self.prev_board.copy()
board_copy.push(move)
- # print(move)
- # print(temp_board)
- # print(temp_board.board_fen())
- # print(board_copy)
- # print(board_copy.board_fen())
if (temp_board.board_fen() == board_copy.board_fen()):
found_move = True
+ print("Found move!", move)
self.board = board_copy.copy()
break
- if found_move:
- print("Found move!")
- else:
+ if not found_move:
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)
@@ -314,6 +341,7 @@ if __name__ == "__main__":
parser.add_argument("--color_scheme", choices=["r/t", "p/y"],
default="p/y", help="Red and teal or pink and yellow for chess piece colors")
parser.add_argument("--show_cv", action="store_true", help="Show opencv images as processing occurs during game")
+ parser.add_argument("--show_board", action="store_true", help="Show only the board (will show warped version)")
parser.add_argument("--show_cam", action="store_true", help="Show persistent camera view")
parser.add_argument("--loop", action="store_true", help="Loop before cv (for taking test images)")
parser.add_argument("--img_idx", help="Where to start indexing images (for naming them; default is 1 if not specified)")
@@ -321,5 +349,5 @@ if __name__ == "__main__":
parser.add_argument("--save_img_as", help="If specified, will save image as given name in game_images")
args = parser.parse_args()
- game = ChessGame(args.difficulty, args.color_scheme, args.show_cv, args.show_cam, args.loop, args.img_idx, args.test_img, args.save_img_as)
+ game = ChessGame(args.difficulty, args.color_scheme, args.show_cv, args.show_board, args.show_cam, args.loop, args.img_idx, args.test_img, args.save_img_as)
game.start_game()
\ No newline at end of file