advancedskrald/main.py

import sys
import warnings
from datetime import datetime
from typing import List, Tuple

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

        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 predict_empty(square, position) != PIECE.EMPTY:
            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)