def get_squares(self, save=None, show=False):
"""
Extracts squares from the original images and applies the thresholding algorithm on the segment.
Args:
save (str): If specified, will save the extracted squares as JPG to the target directory with images named
according to the 0 based grid position. If the directory already exists, the results will be overwritten.
show (bool): If True, will display each square as they are extracted.
Returns:
np.array: Array of all `np.array` images for each square in the Sudoku grid.
"""
if save is not None:
recreate_grid_dir(save)
# Cut each rectangle from the image and save them to the indexed directory
squares = []
for i, rect in enumerate(self.grid_squares):
cut = cv.cut_from_rect(self.cropped, rect)
cut = grid_square_threshold(cut)
cut = cv2.resize(cut, (self.digit_size, self.digit_size))
squares.append(cut)
if show:
cv.show_image(cut)
if save is not None:
cv2.imwrite(os.path.join(save, '%s.jpg' % i), cut)
return np.array(squares)
def get_digits(self,
save=None,
show=False,
include_blanks=False,
raw=False):
"""
Saves only the extracted digits from a Sudoku grid to the digits subfolder.
Args:
save (str): If specified, will save the extracted squares as JPG to the target directory with images named
according to the 0 based grid position. If the directory already exists, the results will be overwritten.
show (bool): If True, will display each square as they are extracted.
include_blanks (bool): If True, will include squares where a digit could not be found as a completely black square.
raw (bool): If True, will extract the digits without any pre-processing.
Returns:
np.array: Array of all `np.array` images for each digit in the Sudoku grid.
"""
if save is not None:
recreate_grid_dir(save)
blank = cv.create_blank_image(
self.digit_size,
self.digit_size,
include_gray_channel=self.include_gray_channel)
digits = []
if self.classification_mode == 'digit-basic':
self.cropped = pre_process_image(self.cropped, skip_dilate=True)
h, w = self.cropped.shape[:2]
self.cropped.reshape(w, h, 1)
for i, rect in enumerate(self.grid_squares):
if 'raw' in self.classification_mode or raw:
digit = extract_cell_raw(
self.cropped,
rect,
self.digit_size,
include_gray_channel=self.include_gray_channel)
else:
digit = extract_digit(
self.cropped,
rect,
self.digit_size,
self.classification_mode,
include_gray_channel=self.include_gray_channel)
if digit is not None:
digits.append(digit)
if save is not None:
cv2.imwrite(os.path.join(save, '%s.jpg' % str(i)), digit)
if show:
cv.show_image(digit)
elif include_blanks:
digits.append(blank)
return np.array(digits)
def classify_digits(name, src):
print('%s Classification' % name)
# Some housekeeping
mkdir(os.path.join(CLASSIFIED_DIR, name))
for i in range(10):
digit_dir = os.path.join(CLASSIFIED_DIR, name, str(i))
mkdir(digit_dir)
# Sort files by their number otherwise we'll run into problems when classifying the digits
files = [f for f in os.listdir(src) if f.split('.')[1] == 'jpg']
files = sorted(files, key=lambda x: int(x.split('.')[0]))
for i, f in enumerate(files):
print('Classifying %s...' % i)
original = [v.replace('.', '0') for k, v in read_original_board(i, src).items()]
grid = Sudoku(os.path.join(src, f), include_gray_channel=True, skip_recog=True)
# Ignore completely blank images, not required in the training set
digits_idx = [(j, digit) for j, digit in enumerate(grid.digits) if not np.array_equal(digit, blank)]
# Modify the image so we can augment the training set for better neural networks
# Structured to use `imgaug.Sequential` objects to describe transformations
if imgaug_seq is not None:
digits = [x[1] for x in digits_idx]
digits = imgaug_seq.augment_images(digits)
digits_idx = [(digits_idx[k][0], digit) for k, digit in enumerate(digits)]
# Modify the image using a custom function outside of the imgaug library
if aug_fn is not None:
digits_idx = [(d[0], aug_fn(d[1], *aug_args)) for d in digits_idx]
for j, digit in digits_idx:
if show:
cv.show_image(digit)
if not dry: # If not a dry run, write the image to the relevant directory
# Only keep a small percentage of blank cells when augmenting to avoid huge skewing of dataset.
# Also should allow us to build larger, meaningful sets and train the models faster.
if original[j] == '0' and (imgaug_seq is not None or aug_fn is not None):
if np.random.randint(50) != 0:
continue
cv2.imwrite(os.path.join(CLASSIFIED_DIR, name, original[j], '%s.jpg' % get_next_file(name, original[j])),
digit)
def show_digits(self, save=None, show=True, colour=255, window_name=None):
"""Shows a preview of what the board looks like once digits have been extracted."""
rows = []
with_border = [
cv2.copyMakeBorder(img.copy(), 1, 1, 1, 1, cv2.BORDER_CONSTANT,
None, colour) for img in self.digits
]
for i in range(9):
row = np.concatenate(with_border[i * 9:((i + 1) * 9)], axis=1)
rows.append(row)
out = np.concatenate(rows)
if show:
cv.show_image(out, name=window_name)
if save is not None:
cv2.imwrite(save, out)
return out
def draw_numbers(self,
numbers,
colour=(50, 50, 255),
thickness=3,
crop=False,
show_known=False,
show=True,
save=None):
"""
Draws numbers onto the cropped or original Sudoku image and shows it on the screen.
Args:
numbers (iterable): Array of size 81 with the numbers to display on the grid.
colour (tuple): BGR (blue, green, red) values between 0 and 255.
thickness (int): Thickness of the font in pixels.
crop (bool): If True, will display the numbers on the cropped and warped image instead of the original.
show_known (bool): If True, will display the predictions for the given numbers on the board instead of the
empty cells.
show (bool): If True, will show the image on the screen.
save (str): If specified, will save the image to that location.
Returns:
np.array: Image with the missing numbers drawn on.
"""
if self.cropped_color is None:
self.cropped_color, rect, matrix = cv.crop_and_warp(self.original,
self.crop_rect,
square=True)
img = self.cropped_color.copy()
scale = int((self.grid_squares[0][1][0] - self.grid_squares[0][0][0]) *
0.075) # Dynamic scale for font
for i, square in enumerate(self.grid_squares):
condition = self.board_int[
i] == 0 # Don't draw numbers given on the board
if show_known: # Unless we want to see them
condition = not condition
if condition:
fh, fw = cv2.getTextSize(
str(numbers[i]), cv2.FONT_HERSHEY_PLAIN, scale,
thickness)[0] # Get font height and width
h_pad = int((square[1][0] - square[0][0] - fw) /
2) # Padding to centre the number
v_pad = int((square[1][1] - square[0][1] - fw) / 2)
h_pad -= self.border # No border on the original, so this compensates
v_pad += self.border
img = cv2.putText(
img,
str(numbers[i]),
(int(square[0][0]) + h_pad, int(square[1][1]) - v_pad),
cv2.FONT_HERSHEY_PLAIN,
fontScale=scale,
color=colour,
thickness=thickness)
# Display the cropped image and return
if crop:
if show:
cv.show_image(img)
if save is not None:
cv2.imwrite(save, img)
return img
# Perform an inverse of the crop and warp transformation to put the image back onto the original
height, width = self.original.shape[:2]
img = cv2.warpPerspective(img,
self.crop_matrix, (width, height),
flags=cv2.WARP_INVERSE_MAP,
dst=self.original,
borderMode=cv2.BORDER_TRANSPARENT)
if show:
cv.show_image(img)
if save is not None:
cv2.imwrite(save, img)
return img
def show_cropped(self):
cv.show_image(self.cropped)
def show_original(self):
cv.show_image(self.original)