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

Fixed diff issue by only doing hough lines once and saving that (assuming...

Fixed diff issue by only doing hough lines once and saving that (assuming board will not be moved). Also for loops are much more efficient now
parent a515e16c
No related branches found
No related tags found
No related merge requests found
......@@ -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]])
......@@ -330,12 +335,19 @@ def find_pieces(warped_img, sorted_warped_points, reader, img_idx):
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 +355,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 +366,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 +380,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 +394,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 +453,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 +461,33 @@ 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
# 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 +496,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 +537,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 +571,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)
......
......@@ -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
......@@ -49,6 +50,10 @@ class ChessGame:
self.img_size = 512
# store warping info
self.corners_sorted = None
self.intersection = None
if not skip_camera:
self.picam2 = Picamera2()
else:
......@@ -175,11 +180,16 @@ 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
......@@ -246,7 +256,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 +270,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
......@@ -298,6 +308,10 @@ class ChessGame:
print("Found move!")
else:
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)
......
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