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
ec7eb073
Commit
ec7eb073
authored
10 months ago
by
zalonzo2
Browse files
Options
Downloads
Patches
Plain Diff
Commented easyocr section of board_detector.py
parent
c9ff8141
No related branches found
Branches containing commit
No related tags found
No related merge requests found
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
board_detector.py
+51
-129
51 additions, 129 deletions
board_detector.py
with
51 additions
and
129 deletions
board_detector.py
+
51
−
129
View file @
ec7eb073
...
...
@@ -299,112 +299,96 @@ def find_pieces(warped_img, sorted_warped_points):
for
i
in
range
(
0
,
8
):
color_grid
.
append
([])
for
j
in
range
(
0
,
8
):
# establish corners of current square
tl
=
sorted_warped_points
[
i
][
j
]
tr
=
sorted_warped_points
[
i
][
j
+
1
]
bl
=
sorted_warped_points
[
i
+
1
][
j
]
br
=
sorted_warped_points
[
i
+
1
][
j
+
1
]
# # create a polygon mask for current grid square
# # (this is because square is not always perfectly square so I can't just loop pixels from tl to tr then bl to br)
# # (this might not be worth the extra computation required to shave off a few pixels from being considered)
# height, width, _ = warped_img.shape
# rect_mask = np.zeros((height, width), dtype=np.uint8)
# poly = np.array([[tl, tr, br, bl]], dtype=np.int32)
# cv2.fillPoly(rect_mask, poly, 255)
# num_pixels = 1
# hue = 0
# # loop through a perfect square that is slightly bigger than actual "square." obtain average hue of color in grid square
# 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 rect_mask[y,x] == 255 and hsv_after[y, x, 0] != 0:
# num_pixels += 1
# hue += hsv_after[y, x, 0]
# avg_hue = hue / num_pixels
# create mask of the square
height
,
width
,
_
=
warped_img
.
shape
rect_mask
=
np
.
zeros
((
height
,
width
),
dtype
=
np
.
uint8
)
poly
=
np
.
array
([[
tl
,
tr
,
br
,
bl
]],
dtype
=
np
.
int32
)
cv2
.
fillPoly
(
rect_mask
,
poly
,
255
)
# get contours inside square mask
masked_hsv_after
=
cv2
.
bitwise_and
(
gray_after
,
gray_after
,
mask
=
rect_mask
)
# display_img([masked_hsv_after])
contours
,
hierarchy
=
cv2
.
findContours
(
masked_hsv_after
,
cv2
.
RETR_TREE
,
cv2
.
CHAIN_APPROX_NONE
)
# print(contours)
num_pixels
=
1
hue
=
0
avg_hue
=
0
cur_bounding_box
=
None
if
contours
is
not
None
:
if
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
)
# print(cv2.contourArea(largest_contour))
# if cv2.contourArea(largest_contour) < 50:
# largest_contour = None
except
ValueError
:
# print("No contours")
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
)
# cv2.drawContours(warped_img, largest_contour, -1, (255, 255, 0), 2)
# display_img([filled_contour_mask])
num_pixels
=
1
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
# for pixel in largest_contour:
# y,x = pixel[0]
# # display_img([hsv_after])
# if hsv_after[y, x, 0] != 0:
# num_pixels += 1
# hue += hsv_after[y, x, 0]
# avg_hue = hue / num_pixels
# if there is a color in square, then label based on custom hue thresholds and add 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
if
largest_contour
is
not
None
:
if
avg_hue
!=
0
:
for
color
,
(
lower
,
upper
)
in
hue_thresh_dict
.
items
()
:
if
lower
<=
avg_hue
<
upp
er
:
x
,
y
,
w
,
h
=
cur_bounding_box
border
=
10
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
result
=
reader
.
readtext
(
image
=
img_to_read
,
allowlist
=
"
PKQRBN
"
,
rotation_info
=
[
180
]
,
text_threshold
=
0.1
,
low_text
=
0.1
,
min_size
=
5
)
if
len
(
result
)
!=
0
:
bound_box
,
letter
,
confidence
=
result
[
0
]
if
show_cv
:
print
(
letter
,
confidence
)
else
:
letter
=
None
if
largest_contour
is
not
None
:
# loop through all large contours (which should be spots with letters)
for
color
,
(
lower
,
upper
)
in
hue_thresh_dict
.
items
():
# for each color threshold
if
lower
<=
avg_hue
<
upper
:
# create img of lett
er
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
:
display_img
([
img_to_read
])
color_grid
[
i
].
append
([
color
,
avg_hue
,
num_pixels
,
letter
])
piece_found
=
True
if
show_cv
and
num_pixels
>
pixel_thresh
:
y
,
x
=
tl
[
0
]
+
5
,
tl
[
1
]
+
5
warped_img_draw
.
text
((
y
,
x
),
color
,
fill
=
(
255
,
0
,
0
))
# draw color found onto image
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
# draw color found onto cv image
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
:
color_grid
[
i
].
append
([
None
,
avg_hue
,
num_pixels
,
None
])
...
...
@@ -418,63 +402,6 @@ 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
])
# reader = easyocr.Reader(['en'], gpu=False)
# results = reader.readtext(
# image=warped_img,
# allowlist="PKQRBNpkqrbn",
# # rotation_info=[180],
# # batch_size=2,
# # text_threshold=0.01,
# # slope_ths=0.0
# )
# print(len(results))
# for text, _, _ in results:
# print(text)
# gray_masked = cv2.bitwise_and(gray_after, gray_after, mask=filled_contour_mask[:,:,0])
# gray_masked[gray_masked != 0] = 255
# scale = 1
# size = img_size * scale
# img_to_read = cv2.resize(gray_masked, (int(size), int(size)))
# reader = easyocr.Reader(['en'])
# results = reader.readtext(
# image=img_to_read,
# allowlist="PKQRBN",
# rotation_info=[180],
# batch_size=500,
# text_threshold=0.3,
# link_threshold = 100000000000,
# slope_ths=0.0,
# ycenter_ths=0.01,
# height_ths=0.01,
# width_ths = 0.01,
# min_size=int(size/150)
# )
# img_to_read = cv2.cvtColor(img_to_read, cv2.COLOR_GRAY2BGR)
# img_to_read = cv2.resize(img_to_read, (img_size, img_size))
# # if show_cv:
# img_to_read_pil = cv2_to_pil(img_to_read)
# img_to_draw = ImageDraw.Draw(img_to_read_pil)
# for result in results:
# bound_box, letter = result[0:2]
# y_min, x_min = [int(min(val)/scale) for val in zip(*bound_box)]
# y_max, x_max = [int(max(val)/scale) for val in zip(*bound_box)]
# # cv2.circle(img_to_read, (int(x_min + (x_max - x_min)/2), int(y_min + (y_max - y_min)/2)), 1, (0, 255, 0), 2)
# # cv2.circle(img_to_read, (x_min, y_min), 1, (0, 255, 0), 2)
# # cv2.putText(img_to_read, letter, (x_min, y_min), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
# y_center, x_center = (int(y_min + (y_Fmax - y_min)/2), int(x_min + (x_max - x_min)/2))
# x_grid = int(x_center / img_size * 8)
# y_grid = int(y_center / img_size * 8)
# # print(y_grid, x_grid, letter)
# color_grid[x_grid][y_grid][3] = letter
# # if show_cv:
# img_to_draw.text((y_min - 10, x_min), letter, fill=(255, 0, 0))
# 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|")
...
...
@@ -488,11 +415,6 @@ def find_pieces(warped_img, sorted_warped_points):
# print("\t\t|", end="")
# print("|")
# img_to_draw._image.save('game_images/ocr_results.jpg')
# if show_cv:
# img_to_draw = pil_to_cv2(img_to_draw._image)
# display_img([img_to_draw])
return
color_grid
def
sort_square_grid_coords
(
coordinates
,
unpacked
):
...
...
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