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

Fast on raspi (diff has been fixed mostly, just don't move chess board)

parent a5526e35
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
board_detector.py
+ 102
− 42
View file @ 59238386
......@@ -9,6 +9,7 @@ import time
import chess
import easyocr
import os
from copy import deepcopy
# 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
......@@ -268,29 +269,46 @@ def find_board(img):
return warped_img, sorted_warped_points
prev_warped_img = None
prev_hsv_img = None
prev_hsv_mask_sat = None
prev_hsv_mask_bright = None
prev_hsv_mask = None
prev_hsv_after = None
prev_bgr_after = None
prev_gray_after = None
prev_filled_contour_mask = None
prev_masked_hsv_after = None
# prev_contours = None
prev_largest_contour = None
# prev_bounding_box_array = None
def find_pieces(warped_img, sorted_warped_points, img_idx):
global prev_warped_img
global prev_hsv_img
global prev_hsv_mask_sat
global prev_hsv_mask_bright
global prev_hsv_mask
global prev_hsv_after
global prev_bgr_after
global prev_gray_after
global prev_filled_contour_mask
global prev_masked_hsv_after
# global prev_contours
global prev_largest_contour
# global prev_bounding_box_array
# save warped img
img_path = os.path.join('game_images', 'warped' + str(img_idx) + '.jpg')
cv2.imwrite(img_path, 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])
# prev_warped_img = None
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)
# 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
......@@ -299,19 +317,37 @@ def find_pieces(warped_img, sorted_warped_points, img_idx):
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)
# 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)
diff_mask_contours = np.zeros_like(gray_after)
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)
# 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)
# display_img([diff_mask])
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)
diff_contours_sorted = diff_contours_sorted[0:5]
cv2.drawContours(diff_mask_contours, diff_contours_sorted, -1, (255), thickness=cv2.FILLED)
# display_img([diff_mask, diff_mask_contours])
# 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),
......@@ -321,13 +357,16 @@ def find_pieces(warped_img, sorted_warped_points, img_idx):
warped_img_pil = cv2_to_pil(warped_img)
warped_img_draw = ImageDraw.Draw(warped_img_pil)
reader = easyocr.Reader(['en'], gpu=False, verbose=False)
if compare_prev_warped:
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)
# if compare_prev_warped:
# prev_filled_contour_mask = np.zeros_like(prev_hsv_after)
# cur_bounding_box_array = np.empty((8,8), dtype=object)
# cur_bounding_box_array.fill([])
pixel_thresh = 40
color_grid = []
# loop through each square of chess board
......@@ -350,9 +389,9 @@ def find_pieces(warped_img, sorted_warped_points, img_idx):
# 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])
# 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
......@@ -375,6 +414,14 @@ def find_pieces(warped_img, sorted_warped_points, img_idx):
prev_largest_contour = max(prev_contours, key=cv2.contourArea)
except ValueError:
prev_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 or prev_largest_contour is not None:
......@@ -385,37 +432,35 @@ def find_pieces(warped_img, sorted_warped_points, img_idx):
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)
# cur_bounding_box_array[j][i] = cur_bounding_box
hue = 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)
# 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_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
if num_pixels > pixel_thresh:
if num_diff_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:
avg_hue = 0
else:
# loop through all pixels in square and find average hue
num_pixels = 1
......@@ -430,8 +475,8 @@ def find_pieces(warped_img, sorted_warped_points, img_idx):
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
# 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
......@@ -478,6 +523,7 @@ def find_pieces(warped_img, sorted_warped_points, img_idx):
break
if not piece_found:
color_grid[i].append([None, avg_hue, num_pixels, letter])
# print("stopping ocr")
else:
color_grid[i].append([None, avg_hue, num_pixels, None])
......@@ -506,6 +552,20 @@ def find_pieces(warped_img, sorted_warped_points, img_idx):
# print("\t\t|", end="")
# print("|")
prev_warped_img = deepcopy(warped_img)
prev_hsv_img = deepcopy(hsv_img)
prev_hsv_mask_sat = deepcopy(hsv_mask_sat)
prev_hsv_mask_bright = deepcopy(hsv_mask_bright)
prev_hsv_mask = deepcopy(hsv_mask)
prev_hsv_after = deepcopy(hsv_after)
prev_bgr_after = deepcopy(bgr_after)
prev_gray_after = deepcopy(gray_after)
prev_filled_contour_mask = deepcopy(filled_contour_mask)
prev_masked_hsv_after = deepcopy(masked_hsv_after)
# prev_contours = deepcopy(contours)
prev_largest_contour = deepcopy(largest_contour)
# prev_bounding_box_array = cur_bounding_box_array.copy()
return color_grid
def sort_square_grid_coords(coordinates, unpacked):
......
game.py
+ 4
− 2
View file @ 59238386
......@@ -84,6 +84,8 @@ class ChessGame:
self.board = chess.Board(example_fens[self.test_img])
print(self.board)
print(self.board.fen())
elif self.img_idx == 2:
self.board = chess.Board()
# self.img_idx += 1
while(1): # game loop
......@@ -131,9 +133,9 @@ class ChessGame:
def ai_turn(self):
print("\n--- CHESS ROBOT'S TURN ---")
# print("Board before chess_AI:", self.board.fen())
ESP_input, best_move = chess_AI(self.board.fen(), self.prev_board.fen(), 0)
# ESP_input, best_move = chess_AI(self.board.fen(), self.prev_board.fen(), 0)
print("Best Move: ", best_move)
# print("Best Move: ", best_move)
# convert best move to coordinates to send
input("[SEND BEST MOVE TO ESP32 AND WAIT]")
# self.switch_turn()
......
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