Skip to content
Snippets Groups Projects
Commit 73599af6 authored by zalonzo2's avatar zalonzo2
Browse files

adding planning

parents a6559e7e dce59845
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
board_detector.py
+ 94
− 43
View file @ 73599af6
......@@ -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
......
game.py
+ 43
− 15
View file @ 73599af6
......@@ -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
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment