advancedskrald/main.py

import sys
import warnings

import cv2
import matplotlib.pyplot as plt
import numpy as np
from sklearn.exceptions import DataConversionWarning

import runner
from util import load_classifier, PIECE, COLOR, POSITION, Board, Squares, PieceAndColor, OUR_PIECES

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


def identify_piece(image: np.ndarray, position: POSITION, sift: cv2.xfeatures2d_SIFT) -> PieceAndColor:
    centers = np.load("training_data/centers.npy")
    best = 0
    probs = {p.name: {} for p in OUR_PIECES}

    best_piece = best_color = None
    for piece in OUR_PIECES:
        for color in COLOR:
            #color = runner.compute_color(file, rank)
            #classifier = load_classifier(f"classifiers/neural_net_{piece}/{color}.pkl")
            classifier = load_classifier(f"classifiers/classifier_{piece}/{color}.pkl")
            features = runner.generate_bag_of_words(image, centers, sift)
            prob = classifier.predict_proba(features)

            #image = cv2.resize(image, (172, 172))
            #data = np.reshape(image, (1, np.product(image.shape)))
            #prob = classifier.predict_proba(data)

            probs[piece.name][color.name] = prob[0, 1]
            print(f"{piece}, {color}, {prob[0, 1]}")
            #if prob[0, 1] > best and color == position.color:  # can only be best if correct color. Iterating through both colors for debugging only
            if prob[0, 1] > best:
                best = prob[0, 1]
                best_piece, best_color = piece, color
    #print(probs)
    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, position, sift)
    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, position, 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_empty(square: np.ndarray, position: POSITION) -> bool:
    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()

    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]}")
    return prob[0, 1] > 0.95


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

        img_src = runner.get_square(warped, position)
        width, height, _ = img_src.shape
        src = img_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"tmp_seg/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.9469 for x in pls])):
            #print(f"{position} is nonempty")
            non_empties.append([position, img_src])
            empty += 1

    #print(64 - empty)

    return non_empties


def find_occupied_squares(warped: np.ndarray) -> Squares:
    non_empties = remove_most_empties(warped)

    completely_non_empties = {}
    for position, square in non_empties:
        if not predict_empty(square, position):
            completely_non_empties[position] = square

    return completely_non_empties


if __name__ == '__main__':

    #runner.train_pieces_svm()

    board = cv2.imread("whole_boards/boards_for_empty/board_1554288891.129901_rank_8.png")
    #board = cv2.imread("whole_boards/boards_for_empty/board_1554286515.323962_rank_3.png")
    warped = runner.warp_board(board)

    tmp = find_occupied_squares(warped)
    for pos, square in tmp:
        cv2.imshow(f"{pos}", square)
    cv2.waitKey(0)
    exit()

    """
    rook_square = runner.get_square(warped, POSITION.H3)
    knight_square = runner.get_square(warped, POSITION.D3)
    cv2.imshow("lel", rook_square)
    cv2.imshow("lil", knight_square)
    #rook_out = cv2.Canny(rook_square, 50, 55, L2gradient=True)
    knight_out = cv2.Canny(knight_square, 50, 55, L2gradient=True)
    knight_out_l = cv2.Canny(knight_square, 50, 55, L2gradient=False)
    cv2.imshow("lal", knight_out)
    cv2.imshow("lul", knight_out_l)
    cv2.waitKey(0)
    exit()
    """
    occupied = find_occupied_squares(warped)

    sift = cv2.xfeatures2d.SIFT_create()
    for position, square in occupied:
        print("---"*15)
        piece, color = identify_piece(square, position, sift)
        print(f"{piece} on {position}")
        text_color = 255 if color == COLOR.WHITE else 0
        cv2.putText(square, f"{position} {piece.name}", (0, 50), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(text_color,)*3, thickness=3)
        cv2.imshow(f"{position}", square)
    cv2.waitKey(0)