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

Commented easyocr section of board_detector.py

parent c9ff8141
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
Showing
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): ...@@ -299,112 +299,96 @@ def find_pieces(warped_img, sorted_warped_points):
for i in range(0,8): for i in range(0,8):
color_grid.append([]) color_grid.append([])
for j in range(0,8): for j in range(0,8):
# establish corners of current square # establish corners of current square
tl = sorted_warped_points[i][j] tl = sorted_warped_points[i][j]
tr = sorted_warped_points[i][j+1] tr = sorted_warped_points[i][j+1]
bl = sorted_warped_points[i+1][j] bl = sorted_warped_points[i+1][j]
br = sorted_warped_points[i+1][j+1] br = sorted_warped_points[i+1][j+1]
# # create a polygon mask for current grid square # create mask of the 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
height, width, _ = warped_img.shape height, width, _ = warped_img.shape
rect_mask = np.zeros((height, width), dtype=np.uint8) rect_mask = np.zeros((height, width), dtype=np.uint8)
poly = np.array([[tl, tr, br, bl]], dtype=np.int32) poly = np.array([[tl, tr, br, bl]], dtype=np.int32)
cv2.fillPoly(rect_mask, poly, 255) cv2.fillPoly(rect_mask, poly, 255)
# get contours inside square mask
masked_hsv_after = cv2.bitwise_and(gray_after, gray_after, mask=rect_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) contours, hierarchy = cv2.findContours(masked_hsv_after, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
# print(contours)
num_pixels = 1 num_pixels = 1
hue = 0 hue = 0
avg_hue = 0 avg_hue = 0
cur_bounding_box = None 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: try:
largest_contour = max(contours, key=cv2.contourArea) largest_contour = max(contours, key=cv2.contourArea)
# print(cv2.contourArea(largest_contour))
# if cv2.contourArea(largest_contour) < 50:
# largest_contour = None
except ValueError: except ValueError:
# print("No contours")
largest_contour = None largest_contour = None
# add contour to contour mask
if largest_contour is not None: if largest_contour is not None:
cv2.drawContours(filled_contour_mask, [largest_contour], -1, (255, 0, 0), thickness=cv2.FILLED) cv2.drawContours(filled_contour_mask, [largest_contour], -1, (255, 0, 0), thickness=cv2.FILLED)
cur_bounding_box = cv2.boundingRect(largest_contour) 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 num_pixels = 1
hue = 0 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 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])): 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: if filled_contour_mask[y, x, 0] > 0 and hsv_after[y, x, 0] != 0:
num_pixels += 1 num_pixels += 1
hue += hsv_after[y, x, 0] hue += hsv_after[y, x, 0]
avg_hue = hue / num_pixels avg_hue = hue / num_pixels
# only save the contour if it has enough pixels, otherwise erase it
if num_pixels < pixel_thresh: if num_pixels < pixel_thresh:
cv2.drawContours(filled_contour_mask, [largest_contour], -1, (0, 0, 0), thickness=cv2.FILLED) cv2.drawContours(filled_contour_mask, [largest_contour], -1, (0, 0, 0), thickness=cv2.FILLED)
largest_contour = None largest_contour = None
# for pixel in largest_contour: # if there is a color in square, then label based on custom hue thresholds, run ocr, and add info to color_grid
# 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
piece_found = False piece_found = False
if largest_contour is not None: if largest_contour is not None: # loop through all large contours (which should be spots with letters)
if avg_hue != 0: for color, (lower, upper) in hue_thresh_dict.items(): # for each color threshold
for color, (lower, upper) in hue_thresh_dict.items(): if lower <= avg_hue < upper:
if lower <= avg_hue < upper: # create img of letter
x, y, w, h = cur_bounding_box x, y, w, h = cur_bounding_box
border = 10 border = 10 # widen bounding box by _ pixels for each letter
bl = (max(y-border,0),max(x-border,0)) bl = (max(y-border,0),max(x-border,0))
tr = (min(y+h+border,img_size),min(x+w+border,img_size)) 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 = masked_hsv_after[bl[0]:tr[0], bl[1]:tr[1]]
img_to_read[img_to_read != 0] = 255 img_to_read[img_to_read != 0] = 255 # bitmap of img of current letter
result = reader.readtext(
image=img_to_read, # detect letter in current square
allowlist="PKQRBN", result = reader.readtext(
rotation_info=[180], image=img_to_read,
text_threshold=0.1, allowlist="PKQRBN", # only want these letters
low_text = 0.1, rotation_info=[180], # either rightside up or upside down
min_size = 5 text_threshold=0.1,
) low_text = 0.1,
if len(result) != 0: min_size = 5
bound_box, letter, confidence = result[0] )
if show_cv:
print(letter, confidence) # get letter if found
else: if len(result) != 0:
letter = None bound_box, letter, confidence = result[0]
if show_cv: if show_cv:
display_img([img_to_read]) print(letter, confidence)
color_grid[i].append([color, avg_hue, num_pixels, letter]) else:
piece_found = True letter = None
if show_cv and num_pixels > pixel_thresh: if show_cv:
y,x = tl[0] + 5, tl[1] + 5 display_img([img_to_read])
warped_img_draw.text((y,x), color, fill=(255, 0, 0)) # draw color found onto image
# 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: if piece_found == False:
color_grid[i].append([None, avg_hue, num_pixels, None]) color_grid[i].append([None, avg_hue, num_pixels, None])
...@@ -418,63 +402,6 @@ def find_pieces(warped_img, sorted_warped_points): ...@@ -418,63 +402,6 @@ def find_pieces(warped_img, sorted_warped_points):
cv2.circle(bgr_after_intersections, point, 1, (255, 255, 255), -1) 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])
# 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) # 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|") # # print("|avg_hue, num_pixels, letter|")
...@@ -488,11 +415,6 @@ def find_pieces(warped_img, sorted_warped_points): ...@@ -488,11 +415,6 @@ def find_pieces(warped_img, sorted_warped_points):
# print("\t\t|", end="") # print("\t\t|", end="")
# print("|") # 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 return color_grid
def sort_square_grid_coords(coordinates, unpacked): def sort_square_grid_coords(coordinates, unpacked):
......
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