advancedskrald/main.py

import cv2
import sys
from collections import defaultdict
from datetime import datetime

import matplotlib.pyplot as plt
import numpy as np

import runner
from util import load_classifier, PIECE, COLOR, POSITION, Board, Squares, PieceAndColor
from sklearn.exceptions import DataConversionWarning
import warnings

warnings.filterwarnings(action='ignore', category=DataConversionWarning)
np.set_printoptions(threshold=sys.maxsize)


def identify_piece(image: np.ndarray, sift : cv2.xfeatures2d_SIFT, empty_bias=False) -> PieceAndColor:
    centers = np.load("training_data/centers.npy")
    probs = defaultdict(lambda: defaultdict(float))
    best = 0
    best_piece = best_color = None
    for piece in PIECE:
        for color in COLOR:
            #color = runner.compute_color(file, rank)
            classifier = load_classifier(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 prob[0, 1] > best:
                best_piece, best_color = piece, color
    print(probs)

    if empty_bias:
        probs[PIECE.EMPTY] *= 1.2

    return best_piece, best_color


def pred_test(position: POSITION, 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 = identify_piece(mystery_image, sift, empty_bias=empty_bias)
    return probs


def pre_process_and_train() -> None:
    runner.do_pre_processing()
    runner.train_pieces_svm()


def build_board_from_squares(squares: Squares) -> Board:
    sift = cv2.xfeatures2d.SIFT_create()
    board = Board()
    counter = 0
    for position, square in squares.values():
        likely_piece = identify_piece(square, sift)
        board[position] = likely_piece
        if likely_piece != PIECE.EMPTY:
            counter += 1

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


def test_entire_board() -> None:
    board_img = cv2.imread("homo_pls_fuck.jpg")
    warped = runner.warp_board(board_img)

    squares = runner.get_squares(warped)
    board = build_board_from_squares(squares)
    print(board)


def predict(square: np.ndarray, position: POSITION) -> PIECE:
    y, x = np.histogram(square.ravel(), bins=32, range=[0, 256])

    left, right = x[:-1], x[1:]
    X = np.array([left, right]).T.flatten()
    Y = np.array([y, y]).T.flatten()
    area = sum(np.diff(x) * y)
    plt.plot(X, Y)
    plt.xlabel(f"{position}")
    #plt.show()


    #for color in COLOR:
    empty_classifier = load_classifier(f"classifiers/classifier_empty/white_piece_on_{position.color}_square.pkl")
    prob = empty_classifier.predict_proba(np.array(y).reshape(1, -1))
    print(f"{position}, {position.color}: {prob[0, 1]}")
    if prob[0, 1] > 0.95:
        print(f"{position} is empty")
        return PIECE.EMPTY

    return None

def remove_most_empties(warped):
    empty = 0
    non_empties = []
    for position in POSITION:
        counter = 0

        src = runner.get_square(warped, position)
        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)
            pls.append(len(y))
            top, bottom, left, right = min(y), max(y), min(x), max(x)

            dst = masked.filled()[top: bottom + 1, left: right + 1]
            cv2.imwrite(f"segment_test/segment_{datetime.utcnow().timestamp()}_{position}.png", 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"{position} is nonempty")
            non_empties.append([position, src])
            empty += 1

    print(64 - empty)

    return non_empties

if __name__ == '__main__':


    #board = cv2.imread("whole_boards/boards_for_empty/board_1554286488.605142_rank_3.png")
    board = cv2.imread("whole_boards/boards_for_empty/board_1554288606.075646_rank_1.png")

    warped = runner.warp_board(board)

    non_empties = remove_most_empties(warped)


    #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()


    counter = 0
    completely_non_empties = []
    for position, square in non_empties:
        #predict(square, position)

        #y, x = np.histogram(square.ravel(), bins=32, range=[0, 256])
        #left, right = x[:-1], x[1:]
        #X = np.array([left, right]).T.flatten()
        #Y = np.array([y, y]).T.flatten()
        #plt.plot(X, Y)
        #plt.xlabel(f"{position}")
        #plt.show()



        if predict(square,position) == PIECE.EMPTY:
                counter += 1
        else:
            completely_non_empties.append([position, square])




    print(counter)
    for position, square in completely_non_empties:
        cv2.imshow(f"{position}", square)
    cv2.waitKey(0)
    exit()


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

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




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

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

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

    cv2.drawKeypoints(square_img, kp, square_img)

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