advancedskrald/runner.py

import glob
import os
import time
from datetime import datetime
from pathlib import Path
from typing import Tuple

import copyreg


import cv2
import numpy as np
from sklearn import cluster, metrics, svm, neural_network
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, OUR_PIECES

here: Path = Path(__file__).parent
BASELINE = cv2.imread(str(here.joinpath("new_baseline_board.png")))
BASELINE_GRAY = cv2.cvtColor(BASELINE, cv2.COLOR_BGR2GRAY)
SIFT = cv2.xfeatures2d.SIFT_create()
BASELINE_KEYPOINTS = SIFT.detect(BASELINE_GRAY)
BASELINE_KEYPOINTS, BASELINE_DES = SIFT.compute(BASELINE_GRAY, BASELINE_KEYPOINTS)


def generate_centers(number_of_clusters, sift: cv2.xfeatures2d_SIFT):
    features = None
    for piece in OUR_PIECES:
        for color in COLOR:
            for filename in glob.glob(f"training_images/{piece}/{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}")

                if features is None:
                    features = np.array(desc)
                else:
                    print(f"{piece}, {color}, {filename}")
                    features = np.vstack((features, 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 OUR_PIECES:
        for color in COLOR:
            for filename in glob.glob(f"training_images/{piece}/{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 OUR_PIECES:
        for filename in glob.glob(f"training_data/{piece}/{color}_square/*.npy"):
            data = np.load(filename)
            X.append(data[0])
            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 OUR_PIECES + (PIECE.EMPTY,):
            print(f"training training_images/{piece}/{square_color}_square/*.png")
            for filename in glob.glob(f"training_images/{piece}/{square_color}_square/*.png"):
                img = cv2.imread(filename)
                y, x = np.histogram(img.ravel(), bins=32, 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 OUR_PIECES:
        for color in COLOR:
            total_weights = len(glob.glob(f"training_images/{piece}/{color}_square/*.png"))

    for piece in OUR_PIECES:
        for color in COLOR:
            current_weight = len(glob.glob(f"training_images/{piece}/{color}_square/*.png"))
            print(f"Training for piece: {piece}")
            X, Y = load_training_data(piece, color)
            classifier = svm.SVC(gamma=0.01, class_weight={0: current_weight, 1: total_weights}, probability=True)
            classifier.fit(X, Y)
            joblib.dump(classifier, f"classifiers/classifier_{piece}/{color}.pkl")


def train_pieces_svm_canny() -> None:
    for square_color in COLOR:
        X = []
        Y = []
        for piece in OUR_PIECES + (PIECE.EMPTY,):
            for filename in glob.glob(f"training_images/{piece}/{square_color}_square/*.png"):
                img = cv2.imread(filename)
                y, x = np.histogram(img.ravel(), bins=32, 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 find_keypoints(camera_image: np.ndarray, baseline: np.ndarray, debug=False) -> Tuple[np.ndarray, np.ndarray]:
    """
    Find keypoints in raw camera image of board.

    :return: (src points, dest points)
    """
    cv2.imwrite("camera_image.png", camera_image)
    camera_image_gray = cv2.cvtColor(camera_image, cv2.COLOR_BGR2GRAY)

    #sift = cv2.xfeatures2d.SURF_create()

    kp_start = time.time()
    camera_image_keypoints = SIFT.detect(camera_image_gray, None)



    camera_image_keypoints, des = SIFT.compute(camera_image_gray, camera_image_keypoints)
    #print("kp:",time.time() - kp_start)

    def_flan = time.time()
    # 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_start = time.time()
    flann = cv2.FlannBasedMatcher(index_params, search_params)
    #print("end_def:", time.time() - def_flan)
    matches = flann.knnMatch(des, BASELINE_DES, k=2)
    #print("flann:",time.time() - flann_start)

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

    #good_matches = list(filter(lambda x: x[0].distance < 0.55 * x[1].distance, matches))

    if debug:
        # Save keypoints
        keypoints_image = camera_image.copy()
        cv2.drawKeypoints(camera_image, keypoints=camera_image_keypoints, outImage=keypoints_image)
        cv2.imwrite("keypoints.png", keypoints_image)
        # Save matches
        matches_image = 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.png", matches_image)

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


    for i, (m, n) in enumerate(good_matches):
        src_points[i, :] = camera_image_keypoints[m.queryIdx].pt
        dst_points[i, :] = BASELINE_KEYPOINTS[m.trainIdx].pt

    return src_points, dst_points


def find_homography(camera_image: np.ndarray,
                    baseline: np.ndarray = BASELINE,
                    debug=False) -> np.ndarray:
    src_points, dst_points = find_keypoints(camera_image, baseline, debug=debug)
    h, mask = cv2.findHomography(src_points, dst_points, cv2.RANSAC)

    return h


def warp_board(camera_image: np.ndarray, homography: np.ndarray = None, debug=False) -> np.ndarray:
    if homography is None:
        homography = find_homography(camera_image, BASELINE, debug=debug)

    height, width, channels = BASELINE.shape
    return cv2.warpPerspective(camera_image, homography, (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 = 2

    x1 = side + (square_size * (position.file - 1))
    x2 = x1 + square_size + padding
    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, fourk=False) -> None:
    for position in POSITION:
        if position.rank == skip_rank:
            continue
        square = get_square(warped_board, position)
        if fourk:
            imwrite(f"training_images/4k/empty/{position.color}_square/training_{position}_{datetime.utcnow().timestamp()}.png", square)
        else:
            imwrite(f"training_images/empty/{position.color}_square/training_{position}_{datetime.utcnow().timestamp()}.png", square)


def load_data_nn(spec_piece, color):
    X = None
    Y = None
    for piece in OUR_PIECES:
        piece_class = int(spec_piece == piece)

        for filename in glob.glob(f"training_images/{piece}/{color}_square/*.png"):
            image = cv2.imread(filename)
            data = np.reshape(image, (1, np.product(image.shape)))
            if X is None:
                if piece_class == 1:
                    for _ in range(10):
                        X = np.array(data)
                        Y = np.array([piece_class])
                else:
                    for _ in range(5):
                        X = np.array(data)
                        Y = np.array([piece_class])
            else:
                if piece_class == 1:
                    for _ in range(10):
                        X = np.vstack((X, data))
                        Y = np.vstack((Y, [piece_class]))
                else:
                    for _ in range(5):
                        X = np.vstack((X, data))
                        Y = np.vstack((Y, [piece_class]))
    return X, Y


def train_nn():
    for piece in OUR_PIECES:
        for color in COLOR:
            X, Y = load_data_nn(piece, color)

            classifier = neural_network.MLPClassifier(hidden_layer_sizes=256)

            classifier.fit(X, Y)

            joblib.dump(classifier, f"classifiers/neural_net_{piece}/{color}.pkl")


if __name__ == '__main__':
    #train_nn()
    do_pre_processing()
    train_pieces_svm()