Skip to content
GitLab
Explore
Sign in
Primary navigation
Search or go to…
Project
A
AI Chess Robot
Manage
Activity
Members
Labels
Plan
Issues
Issue boards
Milestones
Wiki
Code
Merge requests
Repository
Branches
Commits
Tags
Repository graph
Compare revisions
Snippets
Build
Pipelines
Jobs
Pipeline schedules
Artifacts
Deploy
Releases
Package Registry
Model registry
Operate
Environments
Terraform modules
Monitor
Incidents
Analyze
Value stream analytics
Contributor analytics
CI/CD analytics
Repository analytics
Model experiments
Help
Help
Support
GitLab documentation
Compare GitLab plans
Community forum
Contribute to GitLab
Provide feedback
Keyboard shortcuts
?
Snippets
Groups
Projects
Show more breadcrumbs
jhur22
AI Chess Robot
Commits
73599af6
Commit
73599af6
authored
10 months ago
by
zalonzo2
Browse files
Options
Downloads
Plain Diff
adding planning
parents
a6559e7e
dce59845
No related branches found
Branches containing commit
No related tags found
No related merge requests found
Changes
2
Hide whitespace changes
Inline
Side-by-side
Showing
2 changed files
board_detector.py
+94
-43
94 additions, 43 deletions
board_detector.py
game.py
+43
-15
43 additions, 15 deletions
game.py
with
137 additions
and
58 deletions
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
=
7
0
diff_mask
[
diff_mask
>
diff_thresh
]
=
255
diff_mask
[
diff_mask
<=
diff_thresh
]
=
0
#
diff_thresh =
6
0
#
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
...
...
This diff is collapsed.
Click to expand it.
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
This diff is collapsed.
Click to expand it.
Preview
0%
Loading
Try again
or
attach a new file
.
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Save comment
Cancel
Please
register
or
sign in
to comment