segment finding is nice

main.py

@@ -5,11 +5,12 @@ import runner
 from sklearn.externals import joblib
 import numpy as np
 import operator
+from matplotlib import pyplot as plt
 import glob
 import os
 import heapq
 import math
-
+from datetime import datetime
 
 pieces = ['rook', 'knight']
 #pieces = ['rook', 'knight']
@@ -381,16 +382,18 @@ def selective_search(image, use_fast=False, use_slow=False):
     # close image show window
     cv2.destroyAllWindows()
 
+
 def predict(square, file, rank):
-    color = runner.compute_color(file, rank)
-    empty_var_classifier = load_classifier(f"classifiers/classifier_empty_var/{color}.pkl")
+    square_color = runner.compute_color(file, rank)
 
-    magnitude_of_var = np.linalg.norm(cv2.meanStdDev(square)[1])
+    y, x = np.histogram(square.ravel(), bins=256, range=[0, 256])  # TODO: Maybe img.ravel() ?
 
-    prob = empty_var_classifier.predict_proba(np.array(magnitude_of_var).reshape(-1, 1))
-    print(prob[0, 1])
-    if (prob[0, 1]) > 0.5:
-        return 'empty'
+    for color in ['black', 'white']:
+        empty_classifier = load_classifier(f"classifiers/classifier_empty/white_piece_on_{color}_square.pkl")
+        prob = empty_classifier.predict_proba(np.array(y).reshape(1, -1))
+        print(f"{file}{rank}, {color}: {prob[0, 1]}")
+        if prob[0, 1] > 0.5:
+            return 'empty'
 
     return None
 
@@ -398,6 +401,7 @@ def predict(square, file, rank):
 
 @lru_cache()
 def load_classifier(filename):
+    print(f"loading {filename}")
     return joblib.load(filename)
 
 
@@ -405,11 +409,56 @@ def load_classifier(filename):
 
 if __name__ == '__main__':
 
+    board = cv2.imread("whole_boards/boards_for_empty/board_1554288951.972197_.png")
+    warped = runner.warp_board(board)
+
+    empty = 0
+
+
+    files = "ABCDEFGH"
+    ranks = [1, 2, 3, 4, 5, 6, 7, 8]
+
+
+    non_empties = []
+
+    for file in files:
+        for rank in ranks:
+            src = runner.get_square(warped, file, rank)
+            #src = cv2.GaussianBlur(src, (5, 5), 0)
+            segmentator = cv2.ximgproc.segmentation.createGraphSegmentation(sigma=0.75, k=175, min_size=750)
+            segment = segmentator.processImage(src)
+            mask = segment.reshape(list(segment.shape) + [1]).repeat(3, axis=2)
+            masked = np.ma.masked_array(src, fill_value=0)
+            for i in range(np.max(segment)):
+                masked.mask = mask != i
+                y, x = np.where(segment == i)
+                top, bottom, left, right = min(y), max(y), min(x), max(x)
+                dst = masked.filled()[top: bottom + 1, left: right + 1]
+                cv2.imwrite(f"segment_test/segment_{datetime.utcnow().timestamp()}_{file}{rank}.jpg", dst)
+
+            if np.max(segment) >= 2:
+                print(f"{file}{rank} is nonempty")
+                non_empties.append([f"{file}{rank}", src])
+                print(np.max(segment))
+                empty += 1
+    print(64-empty)
+
+    for non_empty in non_empties:
+        cv2.imshow(non_empty[0], non_empty[1])
+    cv2.waitKey(0)
+    exit()
+
+    #empty_classifier = load_classifier(f"classifiers/classifier_empty/white_piece_on_white_square.pkl")
+    #print(empty_classifier.predict_proba(np.array([0]*16).reshape(1, -1))[0, 1])
+
+    #exit()
+
+
     board = cv2.imread("whole_boards/board_102_1554110461.608167_.png")
     warped = runner.warp_board(board)
 
     files = "ABCDEFGH"
-    ranks = [1,2,3,4,5,6,7,8]
+    ranks = [1, 2, 3, 4, 5, 6, 7, 8]
 
     counter = 0
 

runner.py

@@ -1,3 +1,5 @@
+from functools import lru_cache
+
 import cv2
 import numpy as np
 import glob
@@ -9,9 +11,13 @@ from sklearn.externals import joblib
 from sklearn import neural_network
 import heapq
 from datetime import datetime
+
+from sklearn.pipeline import make_pipeline
+from sklearn.preprocessing import StandardScaler
+
 import utils
 
-
+import random
 
 pieces = ["rook", "knight"]
 colors = ['black', 'white']
@@ -214,37 +220,28 @@ def load_training_data(spec_piece, color):
 
 
 def train_empty_or_piece_var():
-    pieces = ['empty', 'knight', 'rook']
-    for color in colors:
-
+    for square_color in ["black", "white"]:
         X = None
         Y = None
+        for piece in ['empty', 'rook', 'knight']:
 
-        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")):
+            for filename in glob.glob(os.path.join("training_images", f"{piece}", f"{square_color}_square", "*.png")):
+                piece_class = 'empty' == piece
                 img = cv2.imread(filename)
 
-                magnitude_of_var = np.linalg.norm(cv2.meanStdDev(img)[1])
+                y, x = np.histogram(img.ravel(), bins=256, range=[0, 256])  # TODO: Maybe img.ravel() ?
 
                 if X is None:
-                    X = np.array(magnitude_of_var)
+                    X = np.array(y)
                     Y = np.array([piece_class])
                 else:
-                    X = np.vstack((X, magnitude_of_var))
+                    X = np.vstack((X, y))
                     Y = np.vstack((Y, [piece_class]))
 
-        classifier = svm.SVC(class_weight={0: current_weight, 1: total_weight - current_weight}, probability=True)
+        classifier = make_pipeline(StandardScaler(),
+                                   svm.SVC(C=10.0, gamma=0.01, probability=True))
         classifier.fit(X, Y)
-        joblib.dump(classifier, f"classifiers/classifier_empty_var/{color}.pkl")
-
-
+        joblib.dump(classifier, f"classifiers/classifier_empty/white_piece_on_{square_color}_square.pkl")
 
 
 def train_pieces_svm():
@@ -349,11 +346,11 @@ def get_square(warped_board, file, rank):
     file = files.index(file)
     rank = 8 - rank
     width, _, _ = warped_board.shape  # board is square anyway
-
-    side = int(width * 0.04)
+    side = int(width * 0.045)
     size = width - 2 * side
     square_size = size // 8
 
+
     padding = 0
 
     x1 = side + (square_size * file)
@@ -413,6 +410,8 @@ def letter_to_int(letter):
     alphabet = list('ABCDEFGH')
     return alphabet.index(letter) + 1
 
+
+@lru_cache(maxsize=64)
 def compute_color(file, rank):
     if ((letter_to_int(file)+rank) % 2):
         return 'white'
@@ -435,4 +434,8 @@ def save_empty_fields(warped, skip_rank=None):
             utils.imwrite(f"training_images/empty/{color}_square/training_{file}{rank}_{datetime.utcnow().timestamp()}.png", square)
 
 
+if __name__ == '__main__':
+    train_empty_or_piece_var()
+
+