advancedskrald/main.py

from functools import lru_cache

import cv2
import runner
from sklearn.externals import joblib
import numpy as np
import operator
import glob
import os
import heapq
import math


pieces = ['rook', 'knight']
#pieces = ['rook', 'knight']
#piece_to_symbol = {'rook': 1, 'knight': 2, 'empty': 0}
piece_to_symbol = {'rook': 1, 'knight': 2}
colors = ['black', 'white']



def classify(image, sift : cv2.xfeatures2d_SIFT, file, rank, empty_bias=False):
    centers = np.load("training_data/centers.npy")
    probs = {'rook': {'black': 0, 'white': 0}, 'knight': {'black': 0, 'white': 0}, 'empty': {'black': 0, 'white': 0}}
    #probs = {'rook': 0, 'knight': 0, 'empty': 0}
    for piece in pieces:
        for color in colors:
            #color = runner.compute_color(file, rank)
            classifier = joblib.load(f"classifiers/classifier_{piece}/{color}.pkl")
            features = runner.generate_bag_of_words(image, centers, sift)
            prob = classifier.predict_proba(features)



            probs[piece][color] = prob[0, 1]


    if empty_bias:
        probs['empty'] *= 1.2

    return probs


def pred_test(file, rank, mystery_image=None, empty_bias=False):
    sift = cv2.xfeatures2d.SIFT_create()
    if mystery_image is None:
        mystery_image = cv2.imread("training_images/rook/white/rook_training_D4_2.png")
    probs = classify(mystery_image, sift, file, rank, empty_bias=empty_bias)
    return probs



def pre_process_and_train():
    runner.do_pre_processing()
    runner.train_pieces_svm()


def build_board_from_dict(board_dict : dict):
    sift = cv2.xfeatures2d.SIFT_create()
    board = [[0]*8 for _ in range(8)]
    counter = 0
    for idx, value in enumerate(board_dict.values()):
        probs = classify(value, sift)
        likely_piece = max(probs.items(), key=operator.itemgetter(1))[0]
        symbol = piece_to_symbol[likely_piece]

        column = idx // 8
        row = (idx % 7)


        board[row][column] = symbol

        print(probs)

        if likely_piece != 'empty':
            counter += 1

    print(counter)
    print(64/(counter-1))
    return board




def detect_using_nn(spec_image):
    probs = {'rook': 0, 'knight': 0}
    for piece in pieces:
        piece_class = piece_to_symbol[piece]
        win_size = (64, 64)
        classifier = joblib.load("classifiers/neural_net_" + piece + ".pkl")

        spec_image = cv2.resize(spec_image, (64, 128))
        features = np.reshape(spec_image, (1, np.product(spec_image.shape)))
        prob = classifier.predict_proba(features)
        print(piece)
        print(prob[0,1])




def test_entire_board():
    board = cv2.imread("homo_pls_fuck.jpg")
    warped = runner.warp_board(board)

    board_dict = runner.get_squares(warped)
    board = build_board_from_dict(board_dict)
    print(board)


def lel_test():
    # img = cv2.imread('training_images/rook/white/rook_training_D4_2.png')

    counter = 0

    for filename in glob.glob(os.path.join("training_images", "empty", "*", "*.png")):

        img = cv2.imread(filename)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        ret = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 3, 3)

        # binarize the image
        #ret, bw = cv2.threshold(gray, 100, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

        # find connected components
        connectivity = 4
        nb_components, output, stats, centroids = cv2.connectedComponentsWithStats(ret, connectivity, cv2.CV_32S)
        sizes = stats[1:, -1]
        nb_components = nb_components - 1
        min_size = 250  # threshhold value for objects in scene
        img2 = np.zeros((img.shape), np.uint8)


        for i in range(0, nb_components + 1):
            # use if sizes[i] >= min_size: to identify your objects
            color = np.random.randint(255, size=3)
            # draw the bounding rectangele around each object
            cv2.rectangle(img2, (stats[i][0], stats[i][1]), (stats[i][0] + stats[i][2], stats[i][1] + stats[i][3]),
                          (0, 255, 0), 2)
            img2[output == i + 1] = color

        #print(nb_components+1)

        if nb_components+1 >= 4:
            counter += 1
            print(filename)
            cv2.imshow("lel", img2)
            cv2.waitKey(0)


    print(counter)


def selective_search(image, use_fast=False, use_slow=False):
    # speed-up using multithreads
    cv2.setUseOptimized(True)
    cv2.setNumThreads(4)


    if type(image) == str:
        # read image
        im = cv2.imread(image)
    else:
        im = image
    # resize image
    #newHeight = 200
    #newWidth = int(im.shape[1] * 150 / im.shape[0])
    #im = cv2.resize(im, (newWidth, newHeight))

    #im = cv2.imread(image)


    #lel, im = cv2.threshold(im, 128, 255, cv2.THRESH_BINARY)



    # create Selective Search Segmentation Object using default parameters
    ss = cv2.ximgproc.segmentation.createSelectiveSearchSegmentation()

    # set input image on which we will run segmentation
    ss.setBaseImage(im)

    # Switch to fast but low recall Selective Search method

    ss.switchToSingleStrategy()

    if (use_fast):
        ss.switchToSelectiveSearchFast()

    # Switch to high recall but slow Selective Search method
    elif (use_slow):
        ss.switchToSelectiveSearchQuality()

    # run selective search segmentation on input image
    rects = ss.process()
    #print('Total Number of Region Proposals: {}'.format(len(rects)))

    # number of region proposals to show
    numShowRects = 150
    # increment to increase/decrease total number
    # of reason proposals to be shown
    increment = 1

    best_proposals = []

    while True:
        # create a copy of original image

        # itereate over all the region proposals
        for i, rect in enumerate(rects):
            imOut = im.copy()

            # draw rectangle for region proposal till numShowRects
            if (i < numShowRects):
                x, y, w, h = rect
                # cv2.rectangle(imOut, (x, y), (x + w, y + h), (0, 255, 0), 1, cv2.LINE_AA)

                # size = (max(w, x) - min(w, x)) * ((max(h, y) - min(h, y)))
                top_left = (x,y)
                bottom_left = (x, y+h)
                top_right = (x+w, y)
                bottom_right = (x+w, y+h)

                rect_width = bottom_right[0] - bottom_left[0]
                rect_height = bottom_right[1] - top_right[1]


                size = rect_width * rect_height
                #print(f"({x}, {y}), ({w}, {h})\n Of size: { size }")

                #cv2.rectangle(imOut, (x, y), (x + w, y + h), (0, 255, 0), 1, cv2.LINE_AA)
                #cv2.imshow("lel", imOut)
                #cv2.waitKey(0)

                best_proposals.append((rect, size))

                #if size > biggest_size:
                #    biggest_rect = (x, y, w, h)
                #    biggest_size = size
                #    print(f"New biggest: \n({x}, {y}), ({w}, {h})\nOf size: {biggest_size}")

            else:
                break

        height, width, channels = im.shape
        center_x = width // 2
        center_y = (height // 2)+5
        dists = []
        #print(f"Amount of best proposals:\n{len(best_proposals)}")

        #print(f"lel: {len(heapq.nlargest(10, best_proposals, key=lambda x: x[1]))}")
        for i in heapq.nlargest(10, best_proposals, key=lambda x: x[1]):
            width, height, channels = im.shape
            #print(width * height)
            #print(i[1])
            x, y, w, h = i[0]



            if i[1] <= (width*height)*0.8 and i[1] > (width*height)*0.25:


                imCop = imOut.copy()

                #cv2.rectangle(imCop, (x, y), (x + w, y + h), (0, 255, 0), 2, cv2.LINE_AA)
                #cv2.imshow("lel", imCop)
                #cv2.waitKey(0)


                #cv2.rectangle(imCop, (x, y), (x + w, y + h), (0, 255, 0), 4, cv2.LINE_AA)

                top_left = (x,y)
                bottom_left = (x, y+h)
                top_right = (x+w, y)
                bottom_right = (x+w, y+h)

                box_center_x = (top_left[0]+bottom_left[0]+top_right[0]+bottom_right[0]) // 4
                box_center_y = (top_left[1]+bottom_left[1]+top_right[1]+bottom_right[1]) // 4

                #print(f"{box_center_x}, {box_center_y}, {center_x}, {center_y}")

                dist = (center_x - box_center_x) ** 2 + (center_y - box_center_y) ** 2
                print(dist)
                dists.append([i, dist])


                cv2.drawMarker(imCop, position=(x+w, h+y), color=(255, 0, 0), thickness=3)
                cv2.drawMarker(imCop, position=(x+w, y), color=(255, 0, 0), thickness=3)
                cv2.drawMarker(imCop, position=(x, y), color=(255, 0, 0), thickness=3)
                cv2.drawMarker(imCop, position=(x, y+h), color=(255, 0, 0), thickness=3)

                cv2.drawMarker(imCop, position=(box_center_x, box_center_y), color=(0, 255, 0), thickness=3)

                cv2.drawMarker(imCop, position=(center_x, center_y), color=(0, 0, 255), thickness=3)


                #cv2.imshow("lel", imCop)
                #cv2.waitKey(0)

        #print("-------"*5)

        for pls in dists:
            imCop = imOut.copy()
            x, y, w, h = pls[0][0]
            #print(x,y,w,h)
            #print(pls[1])

            top_left = (x, y)
            bottom_left = (x, y + h)
            top_right = (x + w, y)
            bottom_right = (x + w, y + h)

            cv2.drawMarker(imCop, position=(x + w, h + y), color=(255, 0, 0), thickness=3)
            cv2.drawMarker(imCop, position=(x + w, y), color=(255, 0, 0), thickness=3)
            cv2.drawMarker(imCop, position=(x, y), color=(255, 0, 0), thickness=3)
            cv2.drawMarker(imCop, position=(x, y + h), color=(255, 0, 0), thickness=3)

            box_center_x = (top_left[0] + bottom_left[0] + top_right[0] + bottom_right[0]) // 4
            box_center_y = (top_left[1] + bottom_left[1] + top_right[1] + bottom_right[1]) // 4

            cv2.drawMarker(imCop, position=(box_center_x, box_center_y), color=(0, 255, 0), thickness=3)

            cv2.drawMarker(imCop, position=(center_y, center_x), color=(0, 0, 255), thickness=3)


            cv2.rectangle(imCop, (x, y), (x + w, y + h), (0, 255, 0), 2, cv2.LINE_AA)
            #cv2.imshow("lel", imCop)
            #cv2.waitKey(0)

        imCop = imOut.copy()




        best = heapq.nsmallest(1, dists, key=lambda x: x[1])

        if (len(best) == 0):
            return ((0, 0), (0, height), (width, 0), (width, height))

        x, y, w, h = best[0][0][0]
        cv2.rectangle(imCop, (x, y), (x + w, y + h), (0, 255, 0), 4, cv2.LINE_AA)

        top_left = (x, y)
        bottom_left = (x, y + h)
        top_right = (x + w, y)
        bottom_right = (x + w, y + h)

        cv2.drawMarker(imCop, position=(x + w, h + y), color=(255, 0, 0), thickness=3)
        cv2.drawMarker(imCop, position=(x + w, y), color=(255, 0, 0), thickness=3)
        cv2.drawMarker(imCop, position=(x, y), color=(255, 0, 0), thickness=3)
        cv2.drawMarker(imCop, position=(x, y + h), color=(255, 0, 0), thickness=3)

        box_center_x = (top_left[0] + bottom_left[0] + top_right[0] + bottom_right[0]) // 4
        box_center_y = (top_left[1] + bottom_left[1] + top_right[1] + bottom_right[1]) // 4

        cv2.drawMarker(imCop, position=(box_center_x, box_center_y), color=(0, 255, 0), thickness=3)

        cv2.drawMarker(imCop, position=(center_x, center_y), color=(0, 0, 255), thickness=3)


        #cv2.imshow("lel", imCop)
        #cv2.waitKey(0)
        return (top_left, bottom_left, top_right, bottom_right)

        # show output
        cv2.imshow("Output", imOut)

        # record key press
        k = cv2.waitKey(0) & 0xFF

        # m is pressed
        if k == 109:
            # increase total number of rectangles to show by increment
            numShowRects += increment
        # l is pressed
        elif k == 108 and numShowRects > increment:
            # decrease total number of rectangles to show by increment
            numShowRects -= increment
        # q is pressed
        elif k == 113:
            break
    # close image show window
    cv2.destroyAllWindows()

def predict(square, file, rank):
    color = runner.compute_color(file, rank)
    empty_var_classifier = load_classifier(f"classifiers/classifier_empty_var/{color}.pkl")

    magnitude_of_var = np.linalg.norm(cv2.meanStdDev(square)[1])

    prob = empty_var_classifier.predict_proba(np.array(magnitude_of_var).reshape(-1, 1))
    print(prob[0, 1])
    if (prob[0, 1]) > 0.5:
        return 'empty'

    return None



@lru_cache()
def load_classifier(filename):
    return joblib.load(filename)




if __name__ == '__main__':

    board = cv2.imread("whole_boards/board_102_1554110461.608167_.png")
    warped = runner.warp_board(board)

    files = "ABCDEFGH"
    ranks = [1,2,3,4,5,6,7,8]

    counter = 0

    for file in files:
        for rank in ranks:
            square = runner.get_square(warped, file, rank)
            if predict(square, file, rank) == 'empty':
                counter += 1

    print(counter)
    exit()


    square = runner.get_square(warped, "D", 2)

    gray_square = cv2.cvtColor(square, cv2.COLOR_BGR2GRAY)
    print(cv2.meanStdDev(gray_square)[1])
    print(cv2.meanStdDev(square)[1])
    cv2.imshow("square", square)
    cv2.waitKey(0)




    print(pred_test("C", 2, square))

    sift: cv2.xfeatures2d_SIFT = cv2.xfeatures2d.SIFT_create()
    gray = cv2.cvtColor(square, cv2.COLOR_BGR2GRAY)

    kp, desc = sift.detectAndCompute(gray, None)

    cv2.drawKeypoints(square, kp, square)

    cv2.imshow("kp", square)
    cv2.waitKey(0)

    exit()

    board = cv2.imread("whole_boards/board_202_1554154094.001122_.png")

    runner.fetch_empty_fields(board)
    exit()


    warped = runner.warp_board(board)

    counter = 0

    #square = runner.get_square(warped, "A", 3)

    #top_left, bottom_left, top_right, bottom_right = selective_search(square, use_fast=True)
    #cropped = square[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]



    for file in files:
        for rank in ranks:

            square = runner.get_square(warped, file, rank)


            top_left, bottom_left, top_right, bottom_right = selective_search(square, use_fast=True)
            cropped = square[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]

            rect_width = bottom_right[0] - bottom_left[0]
            rect_height = bottom_right[1] - top_right[1]

            size = rect_width * rect_height

            square_height, square_width, channels = square.shape

            empty_bias = (size == square_height*square_width)
            if size == square_height*square_width:
                print(f"{file}{rank} is likely empty")

            res = pred_test(file, rank, mystery_image=square, empty_bias=empty_bias)
            print(res)
            if (max(res.items(), key=operator.itemgetter(1))[0] == 'empty'):
                counter += 1



    print(f"Amount of empty fields: {counter}")
    #print("Non-cropped:\t",pred_test(square))
    #print("Cropped:\t",pred_test(cropped))

    #cv2.imshow("square", square)
    #cv2.waitKey(0)


    #runner.do_pre_processing()
    #runner.train()


    #img = "warped_square_B5.png"
    #detect_using_nn(img)

    #selective_search("training_images/empty/colorless/warped_square_A6.png", use_fast=True)
    #selective_search("warped_square_B5.png", use_fast=True)
    img = "training_images/rook/white/rook_training_D4_7.png"
    #img = "training_images/rook/white_square/rook_training_E4_10.png"
    #img = "training_images/knight/white_square/training_D5_134.png"

    #top_left, bottom_left, top_right, bottom_right = selective_search(img, use_fast=True)
    #cropped = cv2.imread(img)[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]

    #cv2.imshow("output", cropped)
    #print(pred_test(cropped))
    #cv2.waitKey(0)



    #lel_test()


    # test_entire_board()

    #board = [[0, 0, 1, 2, 0, 0, 0, 2], [0, 1, 2, 2, 1, 0, 0, 1], [0, 0, 0, 0, 1, 0, 2, 0], [0, 2, 2, 1, 1, 2, 2, 0], [0, 1, 0, 0, 1, 2, 0, 0], [0, 0, 0, 0, 0, 2, 2, 0], [0, 0, 0, 2, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0]]
    #for i in board:
    #    print(i)

    #warped = cv2.imread("homo_pls_fuck.jpg")
    #square = runner.get_square(warped, "D", 4)
    #print(pred_test(square))
    #cv2.imshow("lel", square)
    #cv2.waitKey(0)