advancedskrald/main.py

from functools import lru_cache

import cv2
import runner
from sklearn.externals import joblib
import numpy as np
import operator
from matplotlib import pyplot as plt
import glob
import os
import heapq
import math
from datetime import datetime
import sys

np.set_printoptions(threshold=sys.maxsize)

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):
    square_color = runner.compute_color(file, rank)

    y, x = np.histogram(square.ravel(), bins=256, range=[0, 256])  # TODO: Maybe img.ravel() ?

    for color in ['black', 'white']:
        empty_classifier = load_classifier(f"classifiers/classifier_empty/white_piece_on_{color}_square.pkl")
        prob = empty_classifier.predict_proba(np.array(y).reshape(1, -1))
        print(f"{file}{rank}, {color}: {prob[0, 1]}")
        if prob[0, 1] > 0.5:
            return 'empty'

    return None



@lru_cache()
def load_classifier(filename):
    print(f"loading {filename}")
    return joblib.load(filename)


if __name__ == '__main__':

    board = cv2.imread("whole_boards/boards_for_empty/board_1554286488.605142_.png")
    warped = runner.warp_board(board)

    empty = 0

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


    non_empties = []

    for file in files:
        for rank in ranks:
            counter = 0

            src = runner.get_square(warped, file, rank)
            width, height, _ = src.shape
            src = src[width//25:, height//25:]
            # src = src[:-width//200, :-height//200]
            segmentator = cv2.ximgproc.segmentation.createGraphSegmentation(sigma=0.8, k=150, min_size=700)
            segment = segmentator.processImage(src)

            mask = segment.reshape(list(segment.shape) + [1]).repeat(3, axis=2)
            masked = np.ma.masked_array(src, fill_value=0)
            pls = []
            for i in range(np.max(segment)):
                masked.mask = mask != i

                y, x = np.where(segment == i)

                top, bottom, left, right = min(y), max(y), min(x), max(x)

                dst = masked.filled()[top: bottom + 1, left: right + 1]
                lel = (bottom - top) * (right - left)
                #print(f"this is lel: {lel} ")
                #print(f"this is meh: {np.sum(mask[:,:,0])} ")
                if f"{file}{rank}" == "H7":
                    print("--"*20)
                    print("H7")
                    print(lel)
                    print(len(y))
                    print(np.max(segment))
#                    print(lel)
#                    print(np.sum(mask[:, :, 0]))
                    print("--"*20)


                pls.append(len(y))
                if len(y) < (164**2)*0.65:
                    counter += 1

                cv2.imwrite(f"segment_test/segment_{datetime.utcnow().timestamp()}_{file}{rank}.jpg", dst)





            if np.max(segment) > 0 and not np.all([x < (164**2)*0.2 for x in pls]) and (np.max(segment) >= 3 or np.all([x < (164**2)*0.942 for x in pls])):
                print(f"{file}{rank} is nonempty")
                non_empties.append([f"{file}{rank}", src])
                print(counter)
                print(np.max(segment))
                empty += 1
    print("++"*20)
    print(counter)
    print(64-empty)

    for non_empty in non_empties:
        cv2.imshow(non_empty[0], non_empty[1])
    cv2.waitKey(0)
    exit()

    #empty_classifier = load_classifier(f"classifiers/classifier_empty/white_piece_on_white_square.pkl")
    #print(empty_classifier.predict_proba(np.array([0]*16).reshape(1, -1))[0, 1])

    #exit()


    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)