SudokuSolver / realtime_solver.py
LTPhat's picture
code
1f1fc6b
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
from PIL import Image
from helper_number_page import valid_board
classifier = torch.load('digit_classifier.h5',map_location ='cpu')
classifier.eval()
frameWidth = 960
frameHeight = 720
cap = cv2.VideoCapture(0)
frame_rate = 60
# width is id number 3, height is id 4
cap.set(3, frameWidth)
cap.set(4, frameHeight)
# change brightness
cap.set(10, 150)
prev = 0
while cap.isOpened():
time_elapsed = time.time() - prev
success, img = cap.read()
if time_elapsed > 1. / frame_rate:
prev = time.time()
final_img = img.copy()
to_process_img = img.copy()
#Processing
thresholded_img = preprocess(to_process_img) # Gray-scale img
corners_img, corners_list, org_img = find_contours(thresholded_img, to_process_img)
if corners_list:
# 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
grid_mask = create_grid_mask(horizontal, vertical)
# Resize will get better result ??
grid_mask = cv2.resize(grid_mask,(600,600), cv2.INTER_AREA)
# Extract number
number_img = cv2.bitwise_and(cv2.resize(warped_processed, (600,600), cv2.INTER_AREA), grid_mask)
# 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(classifier, norm_resized, org_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)
final_img = unwarp_image(warp_with_nums, corners_img, corners_list, time.time() - start_time)
cv2.imshow("Result", final_img)
if valid_board(solved_board):
cv2.waitKey(1000)
else:
cv2.imshow("Result", final_img)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cv2.destroyAllWindows()
cap.release()