processing.py

import cv2
import numpy as np
from threshold import preprocess
from utils import find_corners, draw_circle_at_corners, grid_line_helper, draw_line
from utils import clean_square_helper, classify_one_digit

#----------------Process pipe line------------------------------#

# 1) Threshold Adaptive to get gray-scale image to find contours
# 2) Find contours from original image
# 3) Image alignment (warp image) on original image
# 4) Get horizontal, vertical line and create grid mask
# 5) Extract numbers and split gray-scale image into 81 squares
# 6) Clean noise pixels of each square
# 7) Recognize digits
# 8) Solve sudoku
# 9) Draw solved board on warped image
# 10) Unwarped image --> Result


def find_contours(img, original):
    """
    contours: A tuple of all point creating contour lines, each contour is a np array of points (x,y).
    hierachy: [Next, Previous, First_Child, Parent]
    contour approximation: https://pyimagesearch.com/2021/10/06/opencv-contour-approximation/
    """

    # find contours on threshold image
    contours, hierachy =  cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    #sort the largest contour to find the puzzle
    contours = sorted(contours, key = cv2.contourArea, reverse = True)
    polygon = None
    # find the largest rectangle-shape contour to make sure this is the puzzle
    for con in contours:
        area = cv2.contourArea(con)
        perimeter = cv2.arcLength(con, closed = True)
        approx = cv2.approxPolyDP(con, epsilon=0.01 * perimeter, closed =  True)
        num_of_ptr = len(approx)
        if num_of_ptr == 4 and area > 1000:
            polygon = con   #finded puzzle
            break
    if polygon is not None:
        # find corner
        top_left = find_corners(polygon, limit_func= min, compare_func= np.add)
        top_right = find_corners(polygon, limit_func= max, compare_func= np.subtract)
        bot_left = find_corners(polygon,limit_func=min, compare_func= np.subtract)
        bot_right = find_corners(polygon,limit_func=max, compare_func=np.add)
        #Check polygon is square, if not return []
        #Set threshold rate for width and height to determine square bounding box
        if not (0.5 < ((top_right[0]-top_left[0]) / (bot_right[1]-top_right[1]))<1.5):
            print("Exception 1 : Get another image to get square-shape puzzle")
            return [],[],[]
        if bot_right[1] - top_right[1] == 0:
            print("Exception 2 : Get another image to get square-shape puzzle")
            return [],[],[]
        corner_list = [top_left, top_right, bot_right, bot_left]
        draw_original = original.copy()
        cv2.drawContours(draw_original, [polygon], 0, (0,255,0), 3)
        #draw circle at each corner point
        for x in corner_list:
            draw_circle_at_corners(draw_original, x)

        return draw_original, corner_list, original
        # draw_original: Img which drown contour and corner
        # corner_list: list of 4 corner points
        # original: Original imgs
    print("Can not detect puzzle")
    return [],[],[]



def warp_image(corner_list, original):
    """
    Input: 4 corner points and threshold grayscale image
    Output: Perspective transformation matrix and transformed image
    Perspective transformation: https://theailearner.com/tag/cv2-warpperspective/
    """
    try:
        corners = np.array(corner_list, dtype= "float32")
        top_left, top_right, bot_left, bot_right = corners[0], corners[1], corners[2], corners[3]
        #Get the largest side to be the side of squared transfromed puzzle
        side = int(max([
            np.linalg.norm(top_right - bot_right),
            np.linalg.norm(top_left - bot_left),
            np.linalg.norm(bot_right - bot_left),
            np.linalg.norm(top_left - top_right)
        ]))
        out_ptr = np.array([[0,0],[side-1,0],[side-1,side-1], [0,side-1]],dtype="float32")
        transfrom_matrix = cv2.getPerspectiveTransform(corners, out_ptr)
        transformed_image = cv2.warpPerspective(original, transfrom_matrix, (side, side))
        return transformed_image, transfrom_matrix
    except IndexError:
        print("Can not detect corners")
    except:
        print("Something went wrong. Try another image")




def get_grid_line(img, length = 10):
    """
    Get horizontal and vertical lines from warped image
    """

    horizontal = grid_line_helper(img, shape_location= 1)
    vertical = grid_line_helper(img, shape_location=0)
    return vertical, horizontal




def create_grid_mask(horizontal, vertical):
    """
    Completely detect all lines by using Hough Transformation
    Create grid mask to extract number by using bitwise_and with warped images
    """
    # combine two line to make a grid
    grid = cv2.add(horizontal, vertical)
    # Apply threshold to cover more area

    # grid = cv2.adaptiveThreshold(grid, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 235, 2)
    morpho_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3,3))
    grid = cv2.dilate(grid, morpho_kernel, iterations=2)
    # find the line by Houghline transfromation
    lines = cv2.HoughLines(grid, 0.3, np.pi/90, 200) 
    lines_img = draw_line(grid, lines)
    # Extract all the lines
    mask = cv2.bitwise_not(lines_img)
    return mask




def split_squares(number_img):
    """
    Split number img into 81 squares.
    """
    square_list = []
    side = number_img.shape[0] // 9

    #find each square and append to square_list
    for j in range(0,9):
        for i in range(0,9):
            top_left_square = (i * side, j * side)
            bot_right_square = ((i+1) * side, (j+1) * side)
            square_list.append(number_img[top_left_square[1]:bot_right_square[1], top_left_square[0]: bot_right_square[0]])

    return square_list




def clean_square(square_list):
    """
    Return cleaned-square list and number of digits available in the image
    Clean-square list has both 0 and images
    """

    cleaned_squares = []
    count = 0

    for sq in square_list:
        new_img, is_num = clean_square_helper(sq)
        if is_num:
            cleaned_squares.append(new_img)
            count += 1
        else:
            cleaned_squares.append(0)
    return cleaned_squares, count



def clean_square_all_images(square_list):
    """
    Return cleaned-square list 
    Clean-square list has all images(images with no number with be black image after cleaning)
    """

    square_cleaned_list = []
    for i in square_list:
        clean_square, _ = clean_square_helper(i)
        square_cleaned_list.append(clean_square)
    return square_cleaned_list

def recognize_digits(model, resized, org_img):
    res_str = ""
    for img in resized:
        digit = classify_one_digit(model, img, org_img)
        res_str += str(digit)
    return res_str



def draw_digits_on_warped(warped_img, solved_board, unsolved_board):
    """
    Function to draw digits from solved board to warped img
    """
    
    width = warped_img.shape[0] // 9

    img_w_text = np.zeros_like(warped_img)


    for j in range(9):
        for i in range(9):
            if unsolved_board[j][i] == 0: # Only draw new number to blank cell in warped image, avoid overlapping 

                p1 = (i * width, j * width)  # Top left corner of a bounding box
                p2 = ((i + 1) * width, (j + 1) * width)  # Bottom right corner of bounding box

                # Find the center of square to draw digit
                center = ((p1[0] + p2[0]) // 2, (p1[1] + p2[1]) // 2)
                text_size, _ = cv2.getTextSize(str(solved_board[j][i]), cv2.FONT_HERSHEY_SIMPLEX, 1, 6)
                text_origin = (center[0] - text_size[0] // 2, center[1] + text_size[1] // 2)

                cv2.putText(warped_img, str(solved_board[j][i]),
                            text_origin, cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 0, 255), 6)
            
    return img_w_text, warped_img



def unwarp_image(img_src, img_dest, pts, time):
    pts = np.array(pts)

    height, width = img_src.shape[0], img_src.shape[1]
    pts_source = np.array([[0, 0], [width - 1, 0], [width - 1, height - 1], [0, width - 1]],
                          dtype='float32')

    matrix, status = cv2.findHomography(pts_source, pts)
    # Covert to original view perspective
    warped = cv2.warpPerspective(img_src, matrix, (img_dest.shape[1], img_dest.shape[0]))
    # Draw a black rectangle in img_dest

    cv2.fillConvexPoly(img_dest, pts, 0, 16)
    dst_img = cv2.add(img_dest, warped)
    dst_img_height, dst_img_width = dst_img.shape[0], dst_img.shape[1]
    cv2.putText(dst_img, "Time solved: {} s".format(str(np.round(time,4))), (dst_img_width - 360, 60), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)

    return dst_img