Spaces:
Runtime error
Runtime error
File size: 2,630 Bytes
1f1fc6b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | import numpy as np
import torch
from utils import *
from processing import *
from threshold import preprocess
import time
import cv2
from sudoku_solve import Sudoku_solver
import matplotlib.pyplot as plt
# This module performs sudoku solver which input is a image file.
classifier = torch.load('digit_model.h5',map_location ='cpu')
classifier.eval()
def image_solver(img, model):
original_img = img.copy()
threshold = preprocess(img)
corners_img, corners_list, org_img = find_contours(threshold, img)
try:
# Warped original img
warped, matrix = warp_image(corners_list, corners_img)
# Threshold warped img
warped_processed = preprocess(warped) # warped_processed is gray-scaled img
#Get lines
horizontal = grid_line_helper(warped_processed, shape_location=0)
vertical = grid_line_helper(warped_processed, shape_location=1)
# Create mask
if img.shape[0] > 600 or img.shape[1] > 600:
# Resize will get better result ??
grid_mask = create_grid_mask(horizontal, vertical)
grid_mask = cv2.resize(grid_mask,(600,600), cv2.INTER_AREA)
number_img = cv2.bitwise_and(cv2.resize(warped_processed, (600,600), cv2.INTER_AREA), grid_mask)
else:
grid_mask = create_grid_mask(horizontal, vertical)
# Extract number
number_img = cv2.bitwise_and(warped_processed, grid_mask)
# Split into squares
squares = split_squares(number_img)
cleaned_squares = clean_square_all_images(squares)
# Resize and scale pixel
resized_list = resize_square(cleaned_squares)
norm_resized = normalize(resized_list)
# # Recognize digits
rec_str = recognize_digits(model, norm_resized, original_img)
board = convert_str_to_board(rec_str)
# Solve
unsolved_board = board.copy()
sudoku = Sudoku_solver(board, 9)
start_time = time.time()
sudoku.solve()
solved_board = sudoku.board
# Unwarp
_, warp_with_nums = draw_digits_on_warped(warped, solved_board, unsolved_board)
dst_img = unwarp_image(warp_with_nums, corners_img, corners_list, time.time() - start_time)
return dst_img, solved_board
except TypeError:
print("Can not warp image. Please try another image")
if __name__ == "__main__":
url = "streamlit_app\image_from_user\Test40.jpg" # Url for test image
res, solved_board = image_solver(url, classifier)
cv2.imshow("Result", cv2.resize(res, (700,700), cv2.INTER_AREA))
cv2.waitKey(0)
|