advancedskrald/runner.py

import cv2
import numpy as np
import glob
import os
from sklearn import cluster
from sklearn import metrics
from sklearn import svm
from sklearn.externals import joblib
from sklearn import neural_network
import heapq
from datetime import datetime
import utils



pieces = ["rook", "knight"]
colors = ['black', 'white']






def selective_search(image, use_fast=False, use_slow=False, image_name = None):
    # speed-up using multithreads
    cv2.setUseOptimized(True)
    cv2.setNumThreads(4)

    im = image

    img_out = im.copy()

    # create Selective Search Segmentation Object using default parameters
    ss = cv2.ximgproc.segmentation.createSelectiveSearchSegmentation()

    # set input image on which we will run segmentation
    ss.setBaseImage(im)

    ss.switchToSingleStrategy()

    if (use_fast):
        ss.switchToSelectiveSearchFast()

    elif (use_slow):
        ss.switchToSelectiveSearchQuality()

    # run selective search segmentation on input image
    rects = ss.process()

    # number of region proposals to show
    numShowRects = 150

    best_proposals = []

    while True:
        # create a copy of original image

        # itereate over all the region proposals
        for i, rect in enumerate(rects):
            imOut = im.copy()

            # draw rectangle for region proposal till numShowRects
            if (i < numShowRects):
                x, y, w, h = rect

                top_left = (x,y)
                bottom_left = (x, y+h)
                top_right = (x+w, y)
                bottom_right = (x+w, y+h)

                rect_width = bottom_right[0] - bottom_left[0]
                rect_height = bottom_right[1] - top_right[1]

                size = rect_width * rect_height

                best_proposals.append((rect, size))


            else:
                break

        height, width, channels = im.shape
        center_x = width // 2
        center_y = (height // 2)+5
        dists = []

        for i in heapq.nlargest(10, best_proposals, key=lambda x: x[1]):
            width, height, channels = im.shape

            x, y, w, h = i[0]

            if i[1] < (width*height)*0.90 and i[1] > (width*height)*0.25:


                top_left = (x,y)
                bottom_left = (x, y+h)
                top_right = (x+w, y)
                bottom_right = (x+w, y+h)

                box_center_x = (top_left[0]+bottom_left[0]+top_right[0]+bottom_right[0]) // 4
                box_center_y = (top_left[1]+bottom_left[1]+top_right[1]+bottom_right[1]) // 4

                dist = (center_x - box_center_x) ** 2 + (center_y - box_center_y) ** 2
                dists.append([i, dist])


        imCop = imOut.copy()


        print(image_name)
        best = heapq.nsmallest(1, dists, key=lambda x: x[1])
        x, y, w, h = best[0][0][0]
        cv2.rectangle(imCop, (x, y), (x + w, y + h), (0, 255, 0), 4, cv2.LINE_AA)

        top_left = (x, y)
        bottom_right = (x + w, y + h)


        cropped = img_out[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
        return cropped

        # show output
        cv2.imshow("Output", imOut)

        # record key press
        k = cv2.waitKey(0) & 0xFF

        # m is pressed
        if k == 109:
            # increase total number of rectangles to show by increment
            numShowRects += increment
        # l is pressed
        elif k == 108 and numShowRects > increment:
            # decrease total number of rectangles to show by increment
            numShowRects -= increment
        # q is pressed
        elif k == 113:
            break
    # close image show window
    cv2.destroyAllWindows()


def generate_centers(number_of_clusters, sift : cv2.xfeatures2d_SIFT):
    features = None
    for piece in pieces:
        for color in colors:
            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)

                if features is None:
                    features = np.array(desc)
                else:
                    print(f"{piece}, {color}, {filename}")
                    features = np.vstack((features, desc))

    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:
        histogram[0,i] = histogram[0,i] + 1
    histogram = histogram / np.sum(histogram)

    return histogram


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

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

    for piece in pieces:
        for color in colors:
            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(spec_piece, color):
    X = None
    Y = None

    for piece in pieces:
        piece_class = int(spec_piece == piece)
        for filename in glob.glob(os.path.join("training_data", piece, f"{color}_square", "*.npy")):
            data = np.load(filename)
            if X is None:
                X = np.array(data)
                Y = np.array([piece_class])
            else:
                X = np.vstack((X, data))
                Y = np.vstack((Y, [piece_class]))
    return X, Y


def train_empty_or_piece_var():
    pieces = ['empty', 'knight', 'rook']
    for color in colors:

        X = None
        Y = None

        total_weight = 0
        for piece in pieces:
            total_weight += len(glob.glob(os.path.join("training_images", f"{piece}", f"{color}_square", "*.png")))

        current_weight = len(glob.glob(os.path.join("training_images", 'empty', f"{color}_square", "*.png")))

        for piece in pieces:
            piece_class = int('empty' == piece)
            for filename in glob.glob(os.path.join("training_images", piece, f"{color}_square", "*.png")):
                img = cv2.imread(filename)

                magnitude_of_var = np.linalg.norm(cv2.meanStdDev(img)[1])

                if X is None:
                    X = np.array(magnitude_of_var)
                    Y = np.array([piece_class])
                else:
                    X = np.vstack((X, magnitude_of_var))
                    Y = np.vstack((Y, [piece_class]))

        classifier = svm.SVC(class_weight={0: current_weight, 1: total_weight - current_weight}, probability=True)
        classifier.fit(X, Y)
        joblib.dump(classifier, f"classifiers/classifier_empty_var/{color}.pkl")




def train_pieces_svm():
    for piece in pieces:
        for color in colors:
            # 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 compute_features(training_image):
    sift = cv2.xfeatures2d.SIFT_create()
    gray_training_image = cv2.cvtColor(training_image, cv2.COLOR_BGR2GRAY)
    kp = sift.detect(gray_training_image)
    kp, desc = sift.compute(gray_training_image, kp)

    cv2.drawKeypoints(training_image, kp, training_image)

    return training_image


def warp_board(camera_image, debug_image=None):
    #cv2.imwrite('camera_image.png', camera_image)

    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 i in range(len(matches))]

    good_matches = []

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

    draw_params = dict(matchColor=(0,255,0),
                       singlePointColor=(255,0,0),
                       matchesMask=matchesMask,
                       flags=0)

    img3 = cv2.drawMatchesKnn(camera_image,
                              camera_image_keypoints,
                              baseline,
                              baseline_keypoints,
                              matches,
                              None,
                              **draw_params)
    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

    # print(len(points2))

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

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

    # cv2.imwrite('homo_pls_fuck.jpg', im1Reg)
    return im1Reg


def get_square(warped_board, file, rank):
    files = "ABCDEFGH"
    file = files.index(file)
    rank = 8 - rank
    width, _, _ = warped_board.shape  # board is square anyway

    side = int(width * 0.04)
    size = width - 2 * side
    square_size = size // 8

    padding = 0

    x1 = side + (square_size * file)
    x2 = x1 + square_size
    y1 = max(0, side + (square_size * rank) - padding)
    y2 = min(width, y1 + square_size + padding)

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


def get_squares(warped_board):
    result = {}
    for file in "ABCDEFGH":
        for rank in range(1, 9):
            square =  get_square(warped_board, file, rank)
            result[f"{file}{rank}"] = square
            # cv2.imwrite(f"warped_square_{file}{rank}.png", square)
    return result

def load_data_nn(spec_piece):
    X = None
    Y = None
    for piece in pieces:
        piece_class = int(spec_piece == piece)
        for filename in glob.glob(os.path.join("training_images", piece, "*", "*.png")):
            image = cv2.imread(filename)
            image = cv2.resize(image, (64, 128))
            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:
                    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:
                    X = np.vstack((X, data))
                    Y = np.vstack((Y, [piece_class]))
    return (X, Y)

def train_nn():
    for piece in pieces:
        X, Y = load_data_nn(piece)
        classifier = neural_network.MLPClassifier(hidden_layer_sizes=64)
        classifier.fit(X, Y)
        joblib.dump(classifier, "classifiers/neural_net_" + piece + ".pkl")


def letter_to_int(letter):
    alphabet = list('ABCDEFGH')
    return alphabet.index(letter) + 1

def compute_color(file, rank):
    if ((letter_to_int(file)+rank) % 2):
        return 'white'
    else:
        return 'black'


def save_empty_fields(warped, skip_rank=None):
    alpha = "ABCDEFGH"
    ranks = [1, 2, 3, 4, 5, 6, 7, 8]

    if skip_rank is not None:
        ranks.remove(skip_rank)

    for file in alpha:
        for rank in ranks:
            square = get_square(warped, file, rank)
            color = compute_color(file, rank)

            utils.imwrite(f"training_images/empty/{color}_square/training_{file}{rank}_{datetime.utcnow().timestamp()}.png", square)