advancedskrald/runner.py

import cv2
import glob
import os
from datetime import datetime
from typing import Tuple

import numpy as np
from sklearn import cluster, metrics, svm
from sklearn.externals import joblib
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

from util import RANK, POSITION, imwrite, PIECE, COLOR, Squares


def generate_centers(number_of_clusters, sift: cv2.xfeatures2d_SIFT):
    features = []
    for piece in PIECE:
        for color in COLOR:
            for filename in glob.glob(os.path.join("training_images", piece, f"{color}_square", "*.png")):
                image = cv2.imread(filename)
                #image = selective_search(image, use_fast=True)
                gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
                kp, desc = sift.detectAndCompute(gray, None)
                print(f"{piece}, {color}, {filename}")
                features.append(desc)
    features = np.array(features)
    k_means = cluster.KMeans(number_of_clusters)
    k_means.fit(features)
    return k_means.cluster_centers_


def generate_bag_of_words(image, centers, sift: cv2.xfeatures2d_SIFT):
    num_centers = centers.shape[0]
    histogram = np.zeros((1, num_centers))

    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    kp, desc = sift.detectAndCompute(gray_image, None)

    if not kp:
        return histogram

    distances = metrics.pairwise.pairwise_distances(desc, centers)
    best_centers = np.argmin(distances, axis=1)

    for i in best_centers:  # TODO: Could do this way faster in one line with numpy somehow
        histogram[0, i] += + 1

    return histogram / np.sum(histogram)


def do_pre_processing() -> None:
    sift = cv2.xfeatures2d.SIFT_create()
    centers = generate_centers(8, sift)

    np.save("training_data/centers", centers)

    for piece in PIECE:
        for color in COLOR:
            for filename in glob.glob(os.path.join("training_images", piece, f"{color}_square", "*.png")):
                image = cv2.imread(filename)
                #image = selective_search(image, image_name=filename, use_fast=True)
                bow_features = generate_bag_of_words(image, centers, sift)
                np.save(f"training_data/{piece}/{color}_square/{os.path.basename(filename)}", bow_features)


def load_training_data(piece: PIECE, color: COLOR) -> Tuple[np.array, np.array]:
    X = []
    Y = []
    for p in PIECE:
        for filename in glob.glob(os.path.join("training_data", piece, f"{color}_square", "*.npy")):
            data = np.load(filename)
            X.append(data)
            Y.append(p == piece)
    return np.array(X), np.array(Y)


def train_empty_or_piece_hist() -> None:
    for square_color in COLOR:
        X = []
        Y = []
        for piece in (PIECE.EMPTY, PIECE.ROOK, PIECE.KNIGHT):
            for filename in glob.glob(os.path.join("training_images", f"{piece}", f"{square_color}_square", "*.png")):
                img = cv2.imread(filename)
                y, x = np.histogram(img.ravel(), bins=256, range=[0, 256])
                X.append(y)
                Y.append(piece == PIECE.EMPTY)

        classifier = make_pipeline(StandardScaler(),
                                   svm.SVC(C=10.0, gamma=0.01, probability=True))
        classifier.fit(np.array(X), np.array(Y))
        joblib.dump(classifier, f"classifiers/classifier_empty/white_piece_on_{square_color}_square.pkl")


def train_pieces_svm() -> None:
    for piece in PIECE:
        for color in COLOR:
            # TODO: Consider removing empty from total_weights, so all classifiers do not consider empty pieces
            total_weights = len(glob.glob(os.path.join("training_images", "*", f"{color}_square", "*.png")))
            current_weight = len(glob.glob(os.path.join("training_images", piece, f"{color}_square", "*.png")))

            print(f"Trainig for piece: {piece}")
            X, Y = load_training_data(piece, color)
            classifier = svm.SVC(class_weight={0: current_weight, 1: total_weights - current_weight}, probability=True)
            classifier.fit(X, Y)
            joblib.dump(classifier, f"classifiers/classifier_{piece}/{color}.pkl")


def warp_board(camera_image, debug_image=None) -> np.ndarray:
    baseline = cv2.imread("new_baseline_board.png")

    camera_image_gray = cv2.cvtColor(camera_image, cv2.COLOR_BGR2GRAY)
    baseline_gray = cv2.cvtColor(baseline, cv2.COLOR_BGR2GRAY)

    sift = cv2.xfeatures2d.SIFT_create()
    camera_image_keypoints = sift.detect(camera_image_gray, None)
    baseline_keypoints = sift.detect(baseline_gray, None)

    camera_image_keypoints, des = sift.compute(camera_image_gray, camera_image_keypoints)
    baseline_keypoints, des2 = sift.compute(baseline_gray, baseline_keypoints)

    if debug_image is not None:
        cv2.drawKeypoints(camera_image, keypoints=camera_image_keypoints, outImage=debug_image)
        cv2.imwrite("keypoints_img.jpg", camera_image)

    # FLANN parameters
    FLANN_INDEX_KDTREE = 0
    index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=8)
    search_params = dict(checks=100)   # or pass empty dictionary

    flann = cv2.FlannBasedMatcher(index_params,search_params)
    matches = flann.knnMatch(des, des2, k=2)

    # Need to draw only good matches, so create a mask
    matchesMask = [[0, 0] for _ in range(len(matches))]

    # Ratio test as per Lowe's paper
    good_matches = []
    for i, (m, n) in enumerate(matches):
        if m.distance < 0.55*n.distance:
            matchesMask[i] = [1, 0]
            good_matches.append([m, n])

    img3 = cv2.drawMatchesKnn(
        camera_image,
        camera_image_keypoints,
        baseline,
        baseline_keypoints,
        matches,
        None,
        matchColor=(0, 255, 0),
        singlePointColor=(255, 0, 0),
        matchesMask=matchesMask,
        flags=0
    )
    cv2.imwrite("matches.jpg", img3)

    # Extract location of good matches
    points1 = np.zeros((len(good_matches), 2), dtype=np.float32)
    points2 = np.zeros((len(good_matches), 2), dtype=np.float32)

    for i, (m, n) in enumerate(good_matches):
        points1[i, :] = camera_image_keypoints[m.queryIdx].pt
        points2[i, :] = baseline_keypoints[m.trainIdx].pt

    h, mask = cv2.findHomography(points1, points2, cv2.RANSAC)

    height, width, channels = baseline.shape
    return cv2.warpPerspective(camera_image, h, (width, height))


def get_square(warped_board: np.ndarray, position: POSITION) -> np.ndarray:
    width, _, _ = warped_board.shape  # board is square anyway
    side = int(width * 0.045)
    size = width - 2 * side
    square_size = size // 8
    padding = 0

    x1 = side + (square_size * position.file)
    x2 = x1 + square_size
    y1 = max(0, side + (square_size * (8 - position.rank)) - padding)  # 8 - rank because chessboard is from 8 to 1
    y2 = min(width, y1 + square_size + padding)

    square = warped_board[y1:y2, x1:x2]
    return square


def get_squares(warped_board: np.ndarray) -> Squares:
    # cv2.imwrite(f"warped_square_{square}.png", square)
    return {position: get_square(warped_board, position)
            for position in POSITION}


def save_empty_fields(warped_board: np.ndarray, skip_rank: RANK = None) -> None:
    for position in POSITION:
        if position.rank == skip_rank:
            continue
        square = get_square(warped_board, position)
        imwrite(f"training_images/empty/{position.color}_square/training_{position}_{datetime.utcnow().timestamp()}.png", square)


if __name__ == '__main__':
    train_empty_or_piece_hist()