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
......@@ -8,6 +8,7 @@ import psutil
import time
import chess
import easyocr
import os
# 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
......@@ -267,10 +268,30 @@ def find_board(img):
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)
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
saturation_thresh = 60
brightness_thresh = 80
......@@ -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 = 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)
bgr_after = cv2.cvtColor(hsv_after, cv2.COLOR_HSV2BGR)
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
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
......@@ -293,12 +324,17 @@ def find_pieces(warped_img, sorted_warped_points):
reader = easyocr.Reader(['en'], gpu=False, verbose=False)
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
color_grid = []
# loop through each square of chess board
for i in range(0,8):
color_grid.append([])
for j in range(0,8):
do_ocr = False
# establish corners of current square
tl = sorted_warped_points[i][j]
......@@ -314,71 +350,123 @@ def find_pieces(warped_img, sorted_warped_points):
# get contours inside square 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)
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
hue = 0
avg_hue = 0
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:
largest_contour = max(contours, key=cv2.contourArea)
except ValueError:
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
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)
num_pixels = 1
if largest_contour is not None or prev_largest_contour is not None:
if prev_largest_contour is not None:
cv2.drawContours(prev_filled_contour_mask, [prev_largest_contour], -1, (255, 0, 0), thickness=cv2.FILLED)
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
# 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 num_pixels < pixel_thresh:
cv2.drawContours(filled_contour_mask, [largest_contour], -1, (0, 0, 0), thickness=cv2.FILLED)
largest_contour = None
if compare_prev_warped:
# display_img([filled_contour_mask])
# filled_contour_mask = cv2.bitwise_and(filled_contour_mask, filled_contour_mask, mask=diff_mask)
# prev_filled_contour_mask = cv2.bitwise_and(prev_filled_contour_mask, prev_filled_contour_mask, mask=diff_mask)
diff_contour_mask = cv2.bitwise_or(filled_contour_mask, prev_filled_contour_mask, mask=diff_mask)
# diff_contour_mask = cv2.bitwise_and(diff_contour_mask, diff_contour_mask, mask=rect_mask)
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 diff_contour_mask[y, x, 0] > 0:
num_pixels += 1
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
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
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
color_grid[i].append([color, avg_hue, num_pixels, letter])
piece_found = True
......@@ -387,10 +475,10 @@ def find_pieces(warped_img, sorted_warped_points):
if show_cv:
y,x = tl[0] + 5, tl[1] + 5
warped_img_draw.text((y,x), color, fill=(255, 0, 0))
break
if piece_found == False:
if not piece_found:
color_grid[i].append([None, avg_hue, num_pixels, letter])
else:
color_grid[i].append([None, avg_hue, num_pixels, None])
if show_cv:
......@@ -402,18 +490,21 @@ def find_pieces(warped_img, sorted_warped_points):
cv2.circle(bgr_after_intersections, point, 1, (255, 255, 255), -1)
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)
# 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"{int(avg_hue)}, {letter}\t|", end="")
# # print(f"{color}, {letter}\t|", end="")
# else:
# print("\t\t|", end="")
# print("|")
# 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("|")
return color_grid
......
......@@ -78,22 +78,23 @@ class ChessGame:
# initial setup of board
print("\n--- STARTING BOARD ---")
self.do_cv()
# use example fen if not starting from blank game
if self.test_img in example_fens:
self.board = chess.Board(example_fens[self.test_img])
print(self.board)
print(self.board.fen())
self.img_idx += 1
else:
self.do_cv()
# self.img_idx += 1
while(1): # game loop
# if self.board.turn == chess.WHITE:
self.player_turn()
if self.check_game_over():
break
# if self.board.turn == chess.BLACK:
self.ai_turn()
if self.check_game_over():
......@@ -176,7 +177,7 @@ class ChessGame:
return
# 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
if self.color_scheme == 'p/y':
......@@ -189,9 +190,10 @@ class ChessGame:
print("Bad FEN.")
return
# print the board to terminal
print(self.board)
print(self.board.fen())
# print the board to terminal:
if self.test_img not in example_fens or self.img_idx != 2:
print(self.board)
print(self.board.fen())
def check_game_over(self):
if self.board.is_checkmate():
......@@ -251,6 +253,9 @@ class ChessGame:
piece_type = chess.QUEEN
elif letter == "K":
piece_type = chess.KING
elif letter == "X":
temp_board.remove_piece_at(chess.square(j, 7 - i))
continue
piece_color = None
if color == color1:
......@@ -260,7 +265,9 @@ class ChessGame:
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))
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.)
# Would be more efficient to update the extra fen string parameters myself. Might want to change
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment