Skip to content
Snippets Groups Projects
Commit 3ea9715a authored by zalonzo2's avatar zalonzo2
Browse files

Absolutely garbage code for difference checking between images that is... 

Absolutely garbage code for difference checking between images that is actually working and faster than before cuz it skips some ocr
parent ec7eb073
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
board_detector.py
+ 152
− 61
View file @ 3ea9715a
...@@ -8,6 +8,7 @@ import psutil ...@@ -8,6 +8,7 @@ import psutil
import time import time
import chess import chess
import easyocr import easyocr
import os
# global show_cv because I didn't want to have show_cv as an input to every function # global show_cv because I didn't want to have show_cv as an input to every function
# also need skip_camera so I can use cv2.imshow if I'm using my PC # also need skip_camera so I can use cv2.imshow if I'm using my PC
...@@ -267,10 +268,30 @@ def find_board(img): ...@@ -267,10 +268,30 @@ def find_board(img):
return warped_img, sorted_warped_points return warped_img, sorted_warped_points
def find_pieces(warped_img, sorted_warped_points): def find_pieces(warped_img, sorted_warped_points, img_idx):
# save warped img
img_path = os.path.join('game_images', 'warped' + str(img_idx) + '.jpg')
cv2.imwrite(img_path, warped_img)
compare_prev_warped = img_idx > 1
prev_warped_img = None
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(warped_img, prev_warped_img)
diff_mask = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
# display_img([diff_mask])
diff_thresh = 20
diff_mask[diff_mask > diff_thresh] = 255
diff_mask[diff_mask <= diff_thresh] = 0
# display_img([diff_mask])
hsv_img = cv2.cvtColor(warped_img, cv2.COLOR_BGR2HSV) hsv_img = cv2.cvtColor(warped_img, cv2.COLOR_BGR2HSV)
gray_img = cv2.cvtColor(warped_img, cv2.COLOR_BGR2GRAY) gray_img = cv2.cvtColor(warped_img, cv2.COLOR_BGR2GRAY)
if compare_prev_warped:
prev_hsv_img = cv2.cvtColor(prev_warped_img, cv2.COLOR_BGR2HSV)
# threshold to find strongest colors in image # threshold to find strongest colors in image
saturation_thresh = 60 saturation_thresh = 60
brightness_thresh = 80 brightness_thresh = 80
...@@ -278,10 +299,20 @@ def find_pieces(warped_img, sorted_warped_points): ...@@ -278,10 +299,20 @@ def find_pieces(warped_img, sorted_warped_points):
hsv_mask_bright = cv2.inRange(hsv_img[:,:,2], brightness_thresh, 255) # brightness mask hsv_mask_bright = cv2.inRange(hsv_img[:,:,2], brightness_thresh, 255) # brightness mask
hsv_mask = cv2.bitwise_and(hsv_mask_sat, hsv_mask_bright) hsv_mask = cv2.bitwise_and(hsv_mask_sat, hsv_mask_bright)
if compare_prev_warped:
prev_hsv_mask_sat = cv2.inRange(prev_hsv_img[:,:,1], saturation_thresh, 255) # saturation mask
prev_hsv_mask_bright = cv2.inRange(prev_hsv_img[:,:,2], brightness_thresh, 255) # brightness mask
prev_hsv_mask = cv2.bitwise_and(prev_hsv_mask_sat, prev_hsv_mask_bright)
hsv_after = cv2.bitwise_and(hsv_img, hsv_img, mask=hsv_mask) hsv_after = cv2.bitwise_and(hsv_img, hsv_img, mask=hsv_mask)
bgr_after = cv2.cvtColor(hsv_after, cv2.COLOR_HSV2BGR) bgr_after = cv2.cvtColor(hsv_after, cv2.COLOR_HSV2BGR)
gray_after = cv2.cvtColor(bgr_after, cv2.COLOR_BGR2GRAY) gray_after = cv2.cvtColor(bgr_after, cv2.COLOR_BGR2GRAY)
if compare_prev_warped:
prev_hsv_after = cv2.bitwise_and(prev_hsv_img, prev_hsv_img, mask=prev_hsv_mask)
prev_bgr_after = cv2.cvtColor(prev_hsv_after, cv2.COLOR_HSV2BGR)
prev_gray_after = cv2.cvtColor(prev_bgr_after, cv2.COLOR_BGR2GRAY)
# define color thresholds to use to classify colors later on # 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), 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 'teal': (80,105), 'pink': (140,170)} # CHANGE
...@@ -293,12 +324,17 @@ def find_pieces(warped_img, sorted_warped_points): ...@@ -293,12 +324,17 @@ def find_pieces(warped_img, sorted_warped_points):
reader = easyocr.Reader(['en'], gpu=False, verbose=False) reader = easyocr.Reader(['en'], gpu=False, verbose=False)
filled_contour_mask = np.zeros_like(hsv_after) filled_contour_mask = np.zeros_like(hsv_after)
diff_contour_mask = np.zeros_like(hsv_after)
if compare_prev_warped:
prev_filled_contour_mask = np.zeros_like(prev_hsv_after)
pixel_thresh = 40 pixel_thresh = 40
color_grid = [] color_grid = []
# loop through each square of chess board # loop through each square of chess board
for i in range(0,8): for i in range(0,8):
color_grid.append([]) color_grid.append([])
for j in range(0,8): for j in range(0,8):
do_ocr = False
# establish corners of current square # establish corners of current square
tl = sorted_warped_points[i][j] tl = sorted_warped_points[i][j]
...@@ -314,71 +350,123 @@ def find_pieces(warped_img, sorted_warped_points): ...@@ -314,71 +350,123 @@ def find_pieces(warped_img, sorted_warped_points):
# get contours inside square mask # get contours inside square mask
masked_hsv_after = cv2.bitwise_and(gray_after, gray_after, mask=rect_mask) masked_hsv_after = cv2.bitwise_and(gray_after, gray_after, mask=rect_mask)
if compare_prev_warped:
prev_masked_hsv_after = cv2.bitwise_and(prev_gray_after, prev_gray_after, mask=rect_mask)
# display_img([masked_hsv_after,prev_masked_hsv_after])
contours, hierarchy = cv2.findContours(masked_hsv_after, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE) contours, hierarchy = cv2.findContours(masked_hsv_after, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
prev_contours = None
if compare_prev_warped:
prev_contours, hierarchy = cv2.findContours(prev_masked_hsv_after, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
num_pixels = 1 num_pixels = 1
hue = 0 hue = 0
avg_hue = 0 avg_hue = 0
cur_bounding_box = None cur_bounding_box = None
if contours is not None: # if contours are found in the square, find the largest one (which will be the letter if present) if contours is not None or prev_contours is not None: # if contours are found in the square, find the largest one (which will be the letter if present)
try: try:
largest_contour = max(contours, key=cv2.contourArea) largest_contour = max(contours, key=cv2.contourArea)
except ValueError: except ValueError:
largest_contour = None largest_contour = None
prev_largest_contour = None
if prev_contours is not None:
try:
prev_largest_contour = max(prev_contours, key=cv2.contourArea)
except ValueError:
prev_largest_contour = None
# add contour to contour mask # add contour to contour mask
if largest_contour is not None: if largest_contour is not None or prev_largest_contour is not None:
cv2.drawContours(filled_contour_mask, [largest_contour], -1, (255, 0, 0), thickness=cv2.FILLED) if prev_largest_contour is not None:
cur_bounding_box = cv2.boundingRect(largest_contour) cv2.drawContours(prev_filled_contour_mask, [prev_largest_contour], -1, (255, 0, 0), thickness=cv2.FILLED)
num_pixels = 1 cur_bounding_box = cv2.boundingRect(prev_largest_contour)
if largest_contour is not None:
cv2.drawContours(filled_contour_mask, [largest_contour], -1, (255, 0, 0), thickness=cv2.FILLED)
cur_bounding_box = cv2.boundingRect(largest_contour)
hue = 0 hue = 0
if compare_prev_warped:
# loop through all pixels in square and find average hue # display_img([filled_contour_mask])
for x in range(min(tl[0],bl[0]), max(tr[0],br[0])): # filled_contour_mask = cv2.bitwise_and(filled_contour_mask, filled_contour_mask, mask=diff_mask)
for y in range(min(tl[1],tr[1]), max(bl[1],br[1])): # prev_filled_contour_mask = cv2.bitwise_and(prev_filled_contour_mask, prev_filled_contour_mask, mask=diff_mask)
if filled_contour_mask[y, x, 0] > 0 and hsv_after[y, x, 0] != 0: diff_contour_mask = cv2.bitwise_or(filled_contour_mask, prev_filled_contour_mask, mask=diff_mask)
num_pixels += 1 # diff_contour_mask = cv2.bitwise_and(diff_contour_mask, diff_contour_mask, mask=rect_mask)
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])):
# only save the contour if it has enough pixels, otherwise erase it for y in range(min(tl[1],tr[1]), max(bl[1],br[1])):
if num_pixels < pixel_thresh: if diff_contour_mask[y, x, 0] > 0:
cv2.drawContours(filled_contour_mask, [largest_contour], -1, (0, 0, 0), thickness=cv2.FILLED) num_pixels += 1
largest_contour = None
if num_pixels > pixel_thresh:
do_ocr = True
num_pixels = 1
# loop through all pixels in square and find average hue
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
# 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)
largest_contour = None
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
# 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)
largest_contour = None
else:
do_ocr = True
# if there is a color in square, then label based on custom hue thresholds, run ocr, and add info to color_grid # if there is a color in square, then label based on custom hue thresholds, run ocr, and add info to color_grid
piece_found = False piece_found = False
if largest_contour is not None: # loop through all large contours (which should be spots with letters) if do_ocr:
x, y, w, h = cur_bounding_box
border = 10 # widen bounding box by _ pixels for each letter
bl = (max(y-border,0),max(x-border,0))
tr = (min(y+h+border,img_size),min(x+w+border,img_size))
img_to_read = masked_hsv_after[bl[0]:tr[0], bl[1]:tr[1]]
img_to_read[img_to_read != 0] = 255 # bitmap of img of current letter
# display_img([img_to_read])
# detect letter in current square
result = reader.readtext(
image=img_to_read,
allowlist="PKQRBN", # only want these letters
rotation_info=[180], # either rightside up or upside down
text_threshold=0.1,
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"
if show_cv:
display_img([img_to_read])
for color, (lower, upper) in hue_thresh_dict.items(): # for each color threshold for color, (lower, upper) in hue_thresh_dict.items(): # for each color threshold
if lower <= avg_hue < upper: if lower <= avg_hue < upper:
# create img of letter
x, y, w, h = cur_bounding_box
border = 10 # widen bounding box by _ pixels for each letter
bl = (max(y-border,0),max(x-border,0))
tr = (min(y+h+border,img_size),min(x+w+border,img_size))
img_to_read = masked_hsv_after[bl[0]:tr[0], bl[1]:tr[1]]
img_to_read[img_to_read != 0] = 255 # bitmap of img of current letter
# detect letter in current square
result = reader.readtext(
image=img_to_read,
allowlist="PKQRBN", # only want these letters
rotation_info=[180], # either rightside up or upside down
text_threshold=0.1,
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 = None
if show_cv:
display_img([img_to_read])
# update color_grid # update color_grid
color_grid[i].append([color, avg_hue, num_pixels, letter]) color_grid[i].append([color, avg_hue, num_pixels, letter])
piece_found = True piece_found = True
...@@ -387,10 +475,10 @@ def find_pieces(warped_img, sorted_warped_points): ...@@ -387,10 +475,10 @@ def find_pieces(warped_img, sorted_warped_points):
if show_cv: if show_cv:
y,x = tl[0] + 5, tl[1] + 5 y,x = tl[0] + 5, tl[1] + 5
warped_img_draw.text((y,x), color, fill=(255, 0, 0)) warped_img_draw.text((y,x), color, fill=(255, 0, 0))
break break
if not piece_found:
if piece_found == False: color_grid[i].append([None, avg_hue, num_pixels, letter])
else:
color_grid[i].append([None, avg_hue, num_pixels, None]) color_grid[i].append([None, avg_hue, num_pixels, None])
if show_cv: if show_cv:
...@@ -402,18 +490,21 @@ def find_pieces(warped_img, sorted_warped_points): ...@@ -402,18 +490,21 @@ def find_pieces(warped_img, sorted_warped_points):
cv2.circle(bgr_after_intersections, point, 1, (255, 255, 255), -1) 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])
# 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) # 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|") # print("|avg_hue, num_pixels, letter|")
# for row in color_grid: # for row in color_grid:
# print("||", end="") # print("||", end="")
# for color, avg_hue, num_pixels, letter in row: # for color, avg_hue, num_pixels, letter in row:
# if num_pixels > pixel_thresh: # if num_pixels > pixel_thresh:
# print(f"{int(avg_hue)}, {letter}\t|", end="") # print(f"{color}, {letter}\t|", end="")
# # print(f"{color}, {letter}\t|", end="") # # print(f"{color}, {letter}\t|", end="")
# else: # else:
# print("\t\t|", end="") # print("\t\t|", end="")
# print("|") # print("|")
return color_grid return color_grid
......
game.py
+ 15
− 8
View file @ 3ea9715a
...@@ -78,22 +78,23 @@ class ChessGame: ...@@ -78,22 +78,23 @@ class ChessGame:
# initial setup of board # initial setup of board
print("\n--- STARTING BOARD ---") print("\n--- STARTING BOARD ---")
self.do_cv()
# use example fen if not starting from blank game # use example fen if not starting from blank game
if self.test_img in example_fens: if self.test_img in example_fens:
self.board = chess.Board(example_fens[self.test_img]) self.board = chess.Board(example_fens[self.test_img])
print(self.board) print(self.board)
print(self.board.fen()) print(self.board.fen())
self.img_idx += 1 # self.img_idx += 1
else:
self.do_cv()
while(1): # game loop while(1): # game loop
# if self.board.turn == chess.WHITE:
self.player_turn() self.player_turn()
if self.check_game_over(): if self.check_game_over():
break break
# if self.board.turn == chess.BLACK:
self.ai_turn() self.ai_turn()
if self.check_game_over(): if self.check_game_over():
...@@ -176,7 +177,7 @@ class ChessGame: ...@@ -176,7 +177,7 @@ class ChessGame:
return return
# get the pieces based on color thresholding and easyocr # get the pieces based on color thresholding and easyocr
color_grid = find_pieces(warped_img, sorted_warped_points) color_grid = find_pieces(warped_img, sorted_warped_points, self.img_idx - 1)
# convert color_grid to board, which we can get fen string from # convert color_grid to board, which we can get fen string from
if self.color_scheme == 'p/y': if self.color_scheme == 'p/y':
...@@ -189,9 +190,10 @@ class ChessGame: ...@@ -189,9 +190,10 @@ class ChessGame:
print("Bad FEN.") print("Bad FEN.")
return return
# print the board to terminal # print the board to terminal:
print(self.board) if self.test_img not in example_fens or self.img_idx != 2:
print(self.board.fen()) print(self.board)
print(self.board.fen())
def check_game_over(self): def check_game_over(self):
if self.board.is_checkmate(): if self.board.is_checkmate():
...@@ -251,6 +253,9 @@ class ChessGame: ...@@ -251,6 +253,9 @@ class ChessGame:
piece_type = chess.QUEEN piece_type = chess.QUEEN
elif letter == "K": elif letter == "K":
piece_type = chess.KING piece_type = chess.KING
elif letter == "X":
temp_board.remove_piece_at(chess.square(j, 7 - i))
continue
piece_color = None piece_color = None
if color == color1: if color == color1:
...@@ -260,7 +265,9 @@ class ChessGame: ...@@ -260,7 +265,9 @@ class ChessGame:
if piece_color is not None and piece_type is not None: if piece_color is not None and piece_type is not None:
temp_board.set_piece_at(chess.square(j, 7 - i), chess.Piece(piece_type, piece_color)) temp_board.set_piece_at(chess.square(j, 7 - i), chess.Piece(piece_type, piece_color))
continue continue
temp_board.remove_piece_at(chess.square(j, 7 - i))
if self.img_idx == 1:
temp_board.remove_piece_at(chess.square(j, 7 - i))
# find move with correct fen string (with castling, en passant, half move counter, etc.) # find move with correct fen string (with castling, en passant, half move counter, etc.)
# Would be more efficient to update the extra fen string parameters myself. Might want to change # Would be more efficient to update the extra fen string parameters myself. Might want to change
......
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