Python-opencv+KNN求解數獨

最近一直在玩數獨，突發奇想實作影像識別求解數獨，輸入到輸出平均需要0.5s，

整體思路大概就是識別出圖中數字生成list，然后求解，

輸入輸出demo

數獨采用的是微軟自帶的Microsoft sudoku軟體隨便截取的影像，如下圖所示：
經程序式求解后，得到的結果如下圖所示：

程式具體流程

程式整體流程如下圖所示：

讀入影像后，根據求解輪廓資訊找到數字所在位置，以及不包含數字的空白位置，提取數字資訊通過KNN識別，識別出數字；無數字資訊的在list中置0；生成未求解數獨list，之后求解數獨，將資訊在原圖中顯示出來，

# -*-coding:utf-8-*-
import os
import cv2 as cv
import numpy as np
import time

####################################################
#尋找數字生成list
def find_dig_(img, train_set):
    if img is None:
        print("無效的圖片！")
        os._exit(0)
        return
    _, thre = cv.threshold(img, 230, 250, cv.THRESH_BINARY_INV)
    _, contours, hierarchy = cv.findContours(thre, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
    sudoku_list = []
    boxes = []
    for i in range(len(hierarchy[0])):
        if hierarchy[0][i][3] == 0:  # 表示父輪廓為 0
            boxes.append(hierarchy[0][i])
    # 提取數字
    nm = []
    for j in range(len(boxes)):    # 此處len(boxes)=81
        if boxes[j][2] != -1:
            x, y, w, h = cv.boundingRect(contours[boxes[j][2]])
            nm.append([x, y, w, h])
            # 在原圖中框選各個數字
            cropped = img[y:y + h, x:x + w]
            im = img_pre(cropped)			#預處理
            AF = incise(im)				#切割數字影像
            result = identification(train_set, AF, 7)		#knn識別
            sudoku_list.insert(0, int(result))				#生成list
        else:
            sudoku_list.insert(0, 0)
            
    if len(sudoku_list) == 81:
        sudoku_list= np.array(sudoku_list)
        sudoku_list= sudoku_list.reshape((9, 9))
        print("old_sudoku -> \n", sudoku_list)
        return sudoku_list, contours, hierarchy
    else:
        print("無效的圖片！")
        os._exit(0)

######################################################
#KNN演算法識別數字
def img_pre(cropped):
    # 預處理數字影像
    im = np.array(cropped)  # 轉化為二維陣列
    for i in range(im.shape[0]):  # 轉化為二值矩陣
        for j in range(im.shape[1]):
            # print(im[i, j])
            if im[i, j] != 255:
                im[i, j] = 1
            else:
                im[i, j] = 0
    return im


# 提取圖片特征
def feature(A):
    midx = int(A.shape[1] / 2) + 1
    midy = int(A.shape[0] / 2) + 1
    A1 = A[0:midy, 0:midx].mean()
    A2 = A[midy:A.shape[0], 0:midx].mean()
    A3 = A[0:midy, midx:A.shape[1]].mean()
    A4 = A[midy:A.shape[0], midx:A.shape[1]].mean()
    A5 = A.mean()
    AF = [A1, A2, A3, A4, A5]
    return AF


# 切割圖片并回傳每個子圖片特征
def incise(im):
    # 豎直切割并回傳切割的坐標
    a = [];
    b = []
    if any(im[:, 0] == 1):
        a.append(0)
    for i in range(im.shape[1] - 1):
        if all(im[:, i] == 0) and any(im[:, i + 1] == 1):
            a.append(i + 1)
        elif any(im[:, i] == 1) and all(im[:, i + 1] == 0):
            b.append(i + 1)
    if any(im[:, im.shape[1] - 1] == 1):
        b.append(im.shape[1])
    # 水平切割并回傳分割圖片特征
    names = locals();
    AF = []
    for i in range(len(a)):
        names['na%s' % i] = im[:, range(a[i], b[i])]
        if any(names['na%s' % i][0, :] == 1):
            c = 0
        else:
            for j in range(names['na%s' % i].shape[0]):
                if j < names['na%s' % i].shape[0] - 1:
                    if all(names['na%s' % i][j, :] == 0) and any(names['na%s' % i][j + 1, :] == 1):
                        c = j
                        break
                else:
                    c = j
        if any(names['na%s' % i][names['na%s' % i].shape[0] - 1, :] == 1):
            d = names['na%s' % i].shape[0] - 1
        else:
            for j in range(names['na%s' % i].shape[0]):
                if j < names['na%s' % i].shape[0] - 1:
                    if any(names['na%s' % i][j, :] == 1) and all(names['na%s' % i][j + 1, :] == 0):
                        d = j + 1
                        break
                else:
                    d = j
        names['na%s' % i] = names['na%s' % i][range(c, d), :]
        AF.append(feature(names['na%s' % i]))  # 提取特征
        for j in names['na%s' % i]:
            pass
    return AF


# 訓練已知圖片的特征
def training():
    train_set = {}
    for i in range(9):
        value = []
        for j in range(15):
            ima = cv.imread('E:/test_image/knn_test/{}/{}.png'.format(i + 1, j + 1), 0)
            im = img_pre(ima)
            AF = incise(im)
            value.append(AF[0])
        train_set[i + 1] = value

    return train_set


# 計算兩向量的距離
def distance(v1, v2):
    vector1 = np.array(v1)
    vector2 = np.array(v2)
    Vector = (vector1 - vector2) ** 2
    distance = Vector.sum() ** 0.5
    return distance


# 用最近鄰演算法識別單個數字
def knn(train_set, V, k):
    key_sort = [11] * k
    value_sort = [11] * k
    for key in range(1, 10):
        for value in train_set[key]:
            d = distance(V, value)
            for i in range(k):
                if d < value_sort[i]:
                    for j in range(k - 2, i - 1, -1):
                        key_sort[j + 1] = key_sort[j]
                        value_sort[j + 1] = value_sort[j]
                    key_sort[i] = key
                    value_sort[i] = d
                    break
    max_key_count = -1
    key_set = set(key_sort)
    for key in key_set:
        if max_key_count < key_sort.count(key):
            max_key_count = key_sort.count(key)
            max_key = key
    return max_key


# 生成數字
def identification(train_set, AF, k):
    result = ''
    for i in AF:
        key = knn(train_set, i, k)
        result = result + str(key)
    return result



######################################################
######################################################
#求解數獨
def get_next(m, x, y):
    # 獲得下一個空白格在數獨中的坐標，
    :param m 數獨矩陣
    :param x 空白格行數
    :param y 空白格列數
    """
    for next_y in range(y + 1, 9):  # 下一個空白格和當前格在一行的情況
        if m[x][next_y] == 0:
            return x, next_y
    for next_x in range(x + 1, 9):  # 下一個空白格和當前格不在一行的情況
        for next_y in range(0, 9):
            if m[next_x][next_y] == 0:
                return next_x, next_y
    return -1, -1  # 若不存在下一個空白格，則回傳 -1，-1


def value(m, x, y):
    # 回傳符合"每個橫排和豎排以及九宮格內無相同數字"這個條件的有效值，
  
    i, j = x // 3, y // 3
    grid = [m[i * 3 + r][j * 3 + c] for r in range(3) for c in range(3)]
    v = set([x for x in range(1, 10)]) - set(grid) - set(m[x]) - \
        set(list(zip(*m))[y])
    return list(v)


def start_pos(m):
    # 回傳第一個空白格的位置坐標
    for x in range(9):
        for y in range(9):
            if m[x][y] == 0:
                return x, y
    return False, False  # 若數獨已完成，則回傳 False, False


def try_sudoku(m, x, y):
    # 試著填寫數獨
    for v in value(m, x, y):
        m[x][y] = v
        next_x, next_y = get_next(m, x, y)
        if next_y == -1:  # 如果無下一個空白格
            return True
        else:
            end = try_sudoku(m, next_x, next_y)  # 遞回
            if end:
                return True
            m[x][y] = 0  # 在遞回的程序中，如果數獨沒有解開，
            # 則回溯到上一個空白格


def sudoku_so(m):
    x, y = start_pos(m)
    try_sudoku(m, x, y)
    print("new_sudoku -> \n", m)
    return m

###################################################
# 將結果繪制到原圖
def draw_answer(img, contours, hierarchy, new_sudoku_list ):
    new_sudoku_list = new_sudoku_list .flatten().tolist()
    for i in range(len(contours)):
        cnt = contours[i]
        if hierarchy[0, i, -1] == 0:
            num = new_soduku_list.pop(-1)
            if hierarchy[0, i, 2] == -1:
                x, y, w, h = cv.boundingRect(cnt)
                cv.putText(img, "%d" % num, (x + 19, y + 56), cv.FONT_HERSHEY_SIMPLEX, 1.8, (0, 0, 255), 2)  # 填寫數字
    cv.imwrite("E:/answer.png", img)


if __name__ == '__main__':
    t1 = time.time()
    train_set = training()
    img = cv.imread('E:/test_image/python_test_img/Sudoku.png')
    img_gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    sudoku_list, contours, hierarchy = find_dig_(img_gray, train_set)
    new_sudoku_list = sudoku_so(sudoku_list)
    draw_answer(img, contours, hierarchy, new_sudoku_list )
    print("time :",time.time()-t1)

PS：
使用KNN演算法需要創建訓練集，數獨中共涉及9個數字，“1,2,3,4,5,6,7,8,9”各15幅圖放入檔案夾中，如下圖所示，