10 |
9 |
10 |
17 | 
18 |
19 | Bot playing against lichess COMPUTER opponent 20 |
21 | 22 | 23 | 🔴 DISCLAIMER: THIS IS PURELY FOR EDUCATIONAL PURPOSES. DO NOT USE IT AGAINST ACTUAL HUMAN PLAYERS IN ONLINE CHESS WEBSITES (lichess / chess.com / ...). I DONT TAKE ANY RESPONSIBILITY FOR ANY MISUSE. 24 | 25 | 26 | 27 | -------------------------------------------------------------------------------- /image_helper.py: -------------------------------------------------------------------------------- 1 | import PIL 2 | import numpy as np 3 | 4 | 5 | # Load image from file, convert to grayscale float32 numpy array 6 | def load_image_grayscale(img_file): 7 | img = PIL.Image.open(img_file) 8 | 9 | # Convert to grayscale and return 10 | return img.convert("L") 11 | 12 | 13 | def grayscale_resized_image(img): 14 | # Resize if image larger than 2k pixels on a side 15 | if img.size[0] > 2000 or img.size[1] > 2000: 16 | # print(f"Image too big ({img.size[0]} x {img.size[1]})") 17 | new_size = 800.0 18 | if img.size[0] > img.size[1]: 19 | # resize by width to new limit 20 | ratio = new_size / img.size[0] 21 | else: 22 | # resize by height 23 | ratio = new_size / img.size[1] 24 | # print("Reducing by factor of %.2g" % (1. / ratio)) 25 | nx, ny = int(img.size[0] * ratio), int(img.size[1] * ratio) 26 | 27 | img = img.resize((nx, ny), PIL.Image.ADAPTIVE) 28 | # print(f"New size: ({img.size[0]}px x {img.size[1]}px)") 29 | 30 | # Convert to grayscale and array 31 | return np.asarray(img.convert('L'), dtype=np.float32) 32 | 33 | 34 | def grayscale_image(img): 35 | # Convert to grayscale and array 36 | return np.asarray(img.convert('L'), dtype=np.float32) -------------------------------------------------------------------------------- /predictor.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import PIL 3 | from tensorflow.keras.models import load_model 4 | from image_helper import grayscale_resized_image 5 | from board_detector import detect_chessboard, get_chess_tiles 6 | 7 | 8 | # load pre-trained model 9 | model = load_model('models/softmax_v1') 10 | model.summary() 11 | 12 | label = ' KQRBNPkqrbnp' 13 | 14 | 15 | # reads tiles object and returns FEN string 16 | def predict(tiles): 17 | fen = '' 18 | c = 0 19 | print(tiles.shape) 20 | for i in range(8): 21 | if i != 0 and i != 7: 22 | fen += '/' 23 | for j in range(8): 24 | img_arr = tiles[:, :, i * 8 + j] 25 | img = PIL.Image.fromarray(img_arr).resize([32, 32], PIL.Image.ADAPTIVE) 26 | img_arr = grayscale_resized_image(img)[:, :] / 255.0 27 | img_arr = img_arr.reshape(1, 32, 32) 28 | cls = np.argmax(model.predict(img_arr), axis=-1)[0] 29 | if cls == 0: 30 | c += 1 31 | else: 32 | if c != 0: 33 | fen += str(c) 34 | c = 0 35 | fen += label[cls] 36 | if c != 0: 37 | fen += str(c) 38 | c = 0 39 | return fen 40 | 41 | 42 | def test_predictor(): 43 | image_from_file = PIL.Image.open("example_screens/lichess_screen.png") 44 | img = grayscale_resized_image(image_from_file) 45 | is_match, lines_x, lines_y = detect_chessboard(img) 46 | if is_match: 47 | tiles = get_chess_tiles(img, lines_x, lines_y) 48 | print("found Chessboard!!") 49 | print(predict(tiles)) 50 | else: 51 | print("Chessboard not detected on screen!") 52 | 53 | 54 | if __name__ == "__main__": 55 | test_predictor() 56 | -------------------------------------------------------------------------------- /web2png.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from PyQt5.QtWidgets import QApplication 3 | from PyQt5.QtCore import Qt, QUrl, QTimer 4 | from PyQt5.QtWebEngineWidgets import QWebEngineView, QWebEngineSettings 5 | 6 | 7 | class QWebScreenshot(QWebEngineView): 8 | 9 | def capture(self, url, output_file): 10 | self.output_file = output_file 11 | self.load(QUrl(url)) 12 | self.loadFinished.connect(self.on_loaded) 13 | # Create hidden view without scrollbars 14 | self.setAttribute(Qt.WA_DontShowOnScreen) 15 | self.page().settings().setAttribute(QWebEngineSettings.ShowScrollBars, False) 16 | self.show() 17 | 18 | def on_loaded(self): 19 | size = self.page().contentsSize().toSize() 20 | self.resize(size) 21 | # Wait for resize 22 | QTimer.singleShot(1000, self.take_screenshot) 23 | 24 | def take_screenshot(self): 25 | self.grab().save(self.output_file, b'PNG') 26 | self.app.exit(0) 27 | 28 | 29 | app = QApplication(sys.argv) 30 | 31 | 32 | def take_chess_screenshot(fen_string=None, output_filename=None): 33 | # Take uncropped screenshot of lichess board of FEN string and save to file 34 | url_template = "http://en.lichess.org/editor/%s" 35 | s = QWebScreenshot() 36 | s.app = app 37 | s.capture(url_template % fen_string, output_filename) 38 | return app.exec_() 39 | 40 | 41 | def take_screenshot(url=None, output_filename=None): 42 | s = QWebScreenshot() 43 | s.app = app 44 | s.capture(url, output_filename) 45 | return app.exec_() 46 | 47 | 48 | # take_chess_screenshot(fen_string='rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR',output_filename='webpage.png') 49 | # take_chess_screenshot(fen_string='rnbqkbnr/pppppppp/7Q/8/4P3/8/PPPP1PPP/RNBQKBNR',output_filename='webpage2.png') 50 | -------------------------------------------------------------------------------- /tileset_generator.py: -------------------------------------------------------------------------------- 1 | import os 2 | import glob 3 | import numpy as np 4 | import PIL 5 | from image_helper import load_image_grayscale 6 | from board_detector import detect_chessboard, get_chess_tiles 7 | 8 | 9 | def save_tiles(tiles, img_save_dir, img_file): 10 | letters = 'ABCDEFGH' 11 | if not os.path.exists(img_save_dir): 12 | os.makedirs(img_save_dir) 13 | 14 | for i in range(64): 15 | sqr_filename = f"{img_save_dir}/{img_file}_{letters[i % 8]}{int(i / 8) + 1}.png" 16 | 17 | # Make resized 32x32 image from matrix and save 18 | if tiles.shape != (32, 32, 64): 19 | PIL.Image.fromarray(tiles[:, :, i]) \ 20 | .resize([32, 32], PIL.Image.ADAPTIVE) \ 21 | .save(sqr_filename) 22 | else: 23 | # Possibly saving floats 0-1 needs to change fromarray settings 24 | PIL.Image.fromarray((tiles[:, :, i] * 255).astype(np.uint8)) \ 25 | .save(sqr_filename) 26 | 27 | 28 | def generate_tileset(input_chessboard_folder, output_tile_folder): 29 | # Create output folder as needed 30 | if not os.path.exists(output_tile_folder): 31 | os.mkdir(output_tile_folder) 32 | 33 | # Get all image files of type png/jpg/gif 34 | img_files = set(glob.glob("%s/*.png" % input_chessboard_folder))\ 35 | .union(set(glob.glob("%s/*.jpg" % input_chessboard_folder)))\ 36 | .union(set(glob.glob("%s/*.gif" % input_chessboard_folder))) 37 | 38 | num_success = 0 39 | num_failed = 0 40 | num_skipped = 0 41 | 42 | for i, img_path in enumerate(img_files): 43 | print("------") 44 | print(f"{i+1}/{len(img_files)} : {img_path}") 45 | # Strip to just filename 46 | img_file = img_path[len(input_chessboard_folder)+1:-4] 47 | 48 | # Create output save directory or skip this image if it exists 49 | img_save_dir = f"{output_tile_folder}/tiles_{img_file}" 50 | 51 | if os.path.exists(img_save_dir): 52 | print("\tSkipping existing") 53 | num_skipped += 1 54 | continue 55 | 56 | # Load image 57 | print(f"Loading {img_path}...") 58 | img_arr = np.array(load_image_grayscale(img_path), dtype=np.float32) 59 | 60 | # Get tiles 61 | print(f"Generating tiles for {img_file}...") 62 | is_match, lines_x, lines_y = detect_chessboard(img_arr) 63 | 64 | if is_match: 65 | tiles = get_chess_tiles(img_arr, lines_x, lines_y) 66 | print(f"Saving tiles {img_file}...") 67 | save_tiles(tiles, img_save_dir, img_file) 68 | num_success += 1 69 | else: 70 | print("No Match, skipping") 71 | num_failed += 1 72 | 73 | print(f"\t{num_success}/{len(img_files) - num_skipped} generated, {num_failed} failures, {num_skipped} skipped.") -------------------------------------------------------------------------------- /screen_helper.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from PyQt5 import QtCore 3 | from PyQt5 import QtGui 4 | from PyQt5.QtCore import QRect 5 | from PyQt5.QtGui import QPen, QRegion 6 | from PyQt5.QtWidgets import QApplication, QDesktopWidget, QPushButton 7 | from PyQt5.QtWidgets import QWidget 8 | 9 | 10 | class ScreenHighlighter(QWidget): 11 | setVisibleSignal = QtCore.pyqtSignal(bool) 12 | resizeSignal = QtCore.pyqtSignal(int, int) 13 | moveSignal = QtCore.pyqtSignal(int, int) 14 | status_text = "" 15 | color = QtCore.Qt.red 16 | 17 | x0 = 0 18 | y0 = 0 19 | lines_x = [] 20 | lines_y = [] 21 | 22 | def safe_setVisible(self, v): 23 | self.setVisibleSignal.emit(v) 24 | 25 | def safe_resize(self, w, h): 26 | self.resizeSignal.emit(w, h) 27 | 28 | def safe_move(self, x, y): 29 | self.moveSignal.emit(x, y) 30 | 31 | def set_status(self, str): 32 | self.status_text = str 33 | 34 | def set_color(self, color): 35 | self.color = color 36 | 37 | def set_lines(self, x0, y0, lines_x, lines_y): 38 | self.x0 = x0 39 | self.y0 = y0 40 | self.lines_x = lines_x 41 | self.lines_y = lines_y 42 | 43 | def __init__(self, parent=None): 44 | super(ScreenHighlighter, self) \ 45 | .__init__(parent, QtCore.Qt.FramelessWindowHint 46 | | QtCore.Qt.WindowSystemMenuHint 47 | | QtCore.Qt.WindowStaysOnTopHint 48 | | QtCore.Qt.WindowTransparentForInput) 49 | 50 | self.setAttribute(QtCore.Qt.WA_NoSystemBackground) 51 | self.setAttribute(QtCore.Qt.WA_TranslucentBackground) 52 | self.setAttribute(QtCore.Qt.WA_TransparentForMouseEvents) 53 | self.resize(500, 500) 54 | 55 | self.resizeSignal.connect(self.resize) 56 | self.moveSignal.connect(self.move) 57 | self.setVisibleSignal.connect(self.setVisible) 58 | 59 | sizeObject = QDesktopWidget().screenGeometry(-1) 60 | print("Screen size : " + str(sizeObject.width()) + "x" + str(sizeObject.height())) 61 | 62 | self.setWindowTitle(self.tr("ChessBot v0.1")) 63 | 64 | def paintEvent(self, event): 65 | painter = QtGui.QPainter(self) 66 | painter.fillRect(QRect(0, 0, self.geometry().width(), 20),QtCore.Qt.black) 67 | painter.setRenderHint(QtGui.QPainter.Antialiasing) 68 | pen = QPen() 69 | pen.setColor(self.color) 70 | pen.setWidth(3) 71 | painter.setPen(pen) 72 | painter.setBrush(QtCore.Qt.NoBrush) 73 | painter.drawRect(QRect(0, 0, self.geometry().width(), self.geometry().height())) 74 | 75 | # if len(self.lines_x) > 0: 76 | # for i in range(len(self.lines_x)): 77 | # painter.drawLine(self.lines_x[i] - self.x0, 20, self.lines_x[i] - self.x0, self.geometry().height()) 78 | # painter.drawLine(0, self.lines_y[i] - self.y0 + 20, self.geometry().width(), self.lines_y[i] - self.y0 + 20) 79 | 80 | if self.status_text != "": 81 | painter.drawText(5, 15, self.status_text) 82 | 83 | def sizeHint(self): 84 | return QtCore.QSize(200, 200) 85 | 86 | def resizeEvent(self, event): 87 | pass 88 | # r = 3 89 | # w = self.frameGeometry().width() 90 | # h = self.frameGeometry().height() 91 | # reg = QRegion(self.frameGeometry()) \ 92 | # .subtracted( 93 | # QRegion(r, r, w - 2 * r, h - 2 * r, QtGui.QRegion.Rectangle)) 94 | # self.setMask(reg) 95 | 96 | 97 | if __name__ == '__main__': 98 | app = QApplication(sys.argv) 99 | sh = ScreenHighlighter() 100 | sh.show() 101 | sys.exit(app.exec_()) 102 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | import sys 3 | import time 4 | import PIL 5 | import numpy as np 6 | import pyautogui 7 | import pyscreenshot as pss 8 | import matplotlib.pyplot as plt 9 | import tensorflow as tf 10 | import chess 11 | import chess.engine 12 | 13 | from PyQt5 import QtCore 14 | from PyQt5.QtWidgets import QApplication, QWidget, QPushButton, QTextEdit 15 | from PyQt5.QtGui import QIcon 16 | from PyQt5.QtCore import QSize 17 | 18 | from board_detector import detect_chessboard, get_chess_tiles 19 | from image_helper import grayscale_image 20 | from screen_helper import ScreenHighlighter 21 | 22 | app = QApplication(sys.argv) 23 | screen_h = ScreenHighlighter() 24 | screen_h.show() 25 | screen_h.safe_resize(300, 20) 26 | screen_h.safe_move(0, 0) 27 | screen_h.safe_setVisible(False) 28 | 29 | rect_x0 = None 30 | rect_y0 = None 31 | rect_w = None 32 | rect_h = None 33 | 34 | best_move = None 35 | 36 | ch_lines_x = None 37 | ch_lines_y = None 38 | 39 | te_log = None 40 | 41 | engine = None 42 | model = tf.keras.models.load_model('models/softmax_v1') 43 | 44 | 45 | def window(): 46 | global te_log, engine 47 | widget = QWidget() 48 | widget.setFixedSize(310, 200) 49 | 50 | b_locate = QPushButton('', widget) 51 | b_locate.move(0, 0) 52 | b_locate.clicked.connect(b_locate_clicked) 53 | b_locate.setIcon(QIcon('ui/board.png')) 54 | b_locate.setIconSize(QSize(20, 20)) 55 | 56 | b_test1 = QPushButton('', widget) 57 | b_test1.move(60, 0) 58 | b_test1.clicked.connect(test1) 59 | b_test1.setIcon(QIcon('ui/board-full.png')) 60 | b_test1.setIconSize(QSize(20, 20)) 61 | 62 | b_play = QPushButton('', widget) 63 | b_play.move(120,0) 64 | b_play.clicked.connect(b_auto_play_clicked) 65 | b_play.setIcon(QIcon('ui/play.png')) 66 | b_play.setIconSize(QSize(20, 20)) 67 | 68 | b_exit = QPushButton('', widget) 69 | b_exit.move(180, 0) 70 | b_exit.clicked.connect(b_exit_clicked) 71 | b_exit.setIcon(QIcon('ui/exit.png')) 72 | b_exit.setIconSize(QSize(20, 20)) 73 | 74 | te_log = QTextEdit('ChessBot v1.0 by @arminkz', widget) 75 | te_log.move(5, 38) 76 | te_log.setReadOnly(True) 77 | te_log.setFixedSize(300, 150) 78 | 79 | widget.setGeometry(50, 50, 320, 200) 80 | widget.setWindowTitle("Chessbot") 81 | widget.show() 82 | 83 | logln("") 84 | logln("[Engine] Starting Stockfish engine ...") 85 | engine = chess.engine.SimpleEngine.popen_uci("/usr/local/bin/stockfish") 86 | sys.exit(app.exec_()) 87 | 88 | 89 | def b_auto_play_clicked(): 90 | print("TODO") 91 | 92 | 93 | def b_locate_clicked(): 94 | global rect_x0, rect_y0, rect_w, rect_h 95 | global ch_lines_x, ch_lines_y 96 | logln("[Vision] Detecting chessboard on screen ...") 97 | img = pss.grab() 98 | img = img.resize([int(0.5 * s) for s in img.size]) 99 | img = grayscale_image(img) 100 | 101 | is_match, lines_x, lines_y = detect_chessboard(img) 102 | if is_match: 103 | print(lines_x) 104 | print(lines_y) 105 | 106 | ch_lines_x = lines_x 107 | ch_lines_y = lines_y 108 | 109 | # Highlight Detected Area 110 | screen_h.safe_setVisible(True) 111 | 112 | stepx = np.int32(np.round(np.mean(np.diff(lines_x)))) 113 | stepy = np.int32(np.round(np.mean(np.diff(lines_y)))) 114 | 115 | x0 = lines_x[0] - stepx 116 | y0 = lines_y[0] - stepy 117 | w = lines_x[-1] - x0 + stepx 118 | h = lines_y[-1] - y0 + stepy 119 | 120 | rect_x0 = x0 121 | rect_y0 = y0 122 | rect_w = w 123 | rect_h = h 124 | 125 | screen_h.set_lines(x0, y0, lines_x, lines_y) 126 | screen_h.safe_move(x0 - 1, y0 - 21) 127 | screen_h.safe_resize(w + 2, h + 22) 128 | 129 | tiles = get_chess_tiles(img, lines_x, lines_y) 130 | logln(f"[Vision] Chessboard has been detected!") 131 | screen_h.set_status(f"Chessboard detected") 132 | screen_h.set_color(QtCore.Qt.green) 133 | screen_h.safe_setVisible(True) 134 | screen_h.update() 135 | else: 136 | logln(f"[Vision] Error detecting chessboard!") 137 | # screen_h.set_color(QtCore.Qt.red) 138 | # screen_h.safe_move(0, 20) 139 | # screen_h.safe_resize(300, 20) 140 | screen_h.safe_setVisible(False) 141 | # screen_h.update() 142 | 143 | 144 | def b_exit_clicked(): 145 | app.exit(1) 146 | 147 | 148 | def logln(str): 149 | global te_log 150 | te_log.setText(te_log.toPlainText() + "\n" + str) 151 | te_log.verticalScrollBar().setValue(te_log.verticalScrollBar().maximum()) 152 | 153 | 154 | def log(str): 155 | global te_log 156 | te_log.setText(te_log.toPlainText() + str) 157 | te_log.verticalScrollBar().setValue(te_log.verticalScrollBar().maximum()) 158 | 159 | 160 | def display_array(a, rng=[0, 1]): 161 | a = (a - rng[0]) / float(rng[1] - rng[0]) * 255 162 | a = np.uint8(np.clip(a, 0, 255)) 163 | plt.imshow(PIL.Image.fromarray(a)) 164 | plt.show() 165 | 166 | 167 | def test1(): 168 | global best_move 169 | # print('model loaded') 170 | # print(model.summary()) 171 | logln("[Piece Detection] Predicting using neural net ...") 172 | img = pss.grab() 173 | img = img.resize([int(0.5 * s) for s in img.size]) 174 | img = grayscale_image(img) 175 | tiles = get_chess_tiles(img, ch_lines_x, ch_lines_y) 176 | # print(tiles.shape) 177 | norm_tiles = np.zeros((64, 32, 32)) 178 | for i in range(64): 179 | resized_image = PIL.Image.fromarray(tiles[:, :, i]).resize([32, 32]) 180 | norm_tiles[i, :, :] = np.array(resized_image) / 255.0 181 | # print(norm_tiles.shape) 182 | preds = model.predict(norm_tiles) 183 | labels = np.argmax(preds, axis=1) 184 | 185 | logln("[Piece Detection] FEN: ") 186 | fen_text = label2FEN(labels) 187 | log(fen_text) 188 | 189 | # feed FEN to stockfish 190 | board = chess.Board(f"{fen_text} w") 191 | result = engine.play(board, chess.engine.Limit(time=0.1)) 192 | logln(f"[Engine] Best Move: {result.move}") 193 | best_move = result.move 194 | 195 | move = str(best_move) 196 | start_coord = move[0:2] 197 | end_coord = move[2:4] 198 | sx = ord(start_coord[0]) - ord('a') 199 | sy = 8 - int(start_coord[1]) 200 | ex = ord(end_coord[0]) - ord('a') 201 | ey = 8 - int(end_coord[1]) 202 | cell = rect_w / 8 203 | tsx = rect_x0 + sx*cell + (cell/2) 204 | tsy = rect_y0 + sy*cell + (cell/2) 205 | tex = rect_x0 + ex*cell + (cell/2) 206 | tey = rect_y0 + ey*cell + (cell/2) 207 | pyautogui.click(tsx, tsy, clicks=2, interval=0.2) 208 | time.sleep(0.3) 209 | pyautogui.dragTo(tex, tey, button='left', duration=0.3) 210 | time.sleep(1.5) 211 | test1() 212 | 213 | 214 | 215 | def label2FEN(labels): 216 | fen = "" 217 | spc = 0 218 | lastsp = False 219 | for i in range(8): 220 | for j in range(8): 221 | lbl = ' KQRBNPkqrbnp'[labels[(7 - i) * 8 + j]] 222 | if lbl == ' ': 223 | lastsp = True 224 | spc += 1 225 | else: 226 | if lastsp: 227 | fen += str(spc) 228 | spc = 0 229 | lastsp = False 230 | fen += lbl 231 | if lastsp: 232 | fen += str(spc) 233 | spc = 0 234 | lastsp = False 235 | if i != 7: fen += '/' 236 | return fen 237 | 238 | 239 | if __name__ == '__main__': 240 | window() 241 | -------------------------------------------------------------------------------- /board_detector.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import PIL.Image 3 | import tensorflow as tf 4 | import scipy.signal 5 | from image_helper import grayscale_resized_image 6 | 7 | 8 | # convert kernel matrix to tensor-compatible filter 9 | def make_tf_kernel(k): 10 | k = np.asarray(k) 11 | # reshape it to tensorflow 4-D filter 12 | k = k.reshape(list(k.shape) + [1, 1]) 13 | return tf.constant(k, dtype=tf.float32) 14 | 15 | 16 | # Simple 2D convolution 17 | def simple_conv2d(x, k): 18 | x = tf.expand_dims(tf.expand_dims(x, 0), -1) 19 | y = tf.nn.depthwise_conv2d(x, k, [1, 1, 1, 1], padding='SAME') 20 | return y[0, :, :, 0] 21 | 22 | 23 | def gradient_x(x): 24 | k = make_tf_kernel([[-1, 0, 1], 25 | [-1, 0, 1], 26 | [-1, 0, 1]]) 27 | return simple_conv2d(x, k) 28 | 29 | 30 | def gradient_y(x): 31 | k = make_tf_kernel([[-1, -1, -1], 32 | [0, 0, 0], 33 | [1, 1, 1]]) 34 | return simple_conv2d(x, k) 35 | 36 | 37 | # checks whether there exists 7 lines of consistent increasing order in set of lines 38 | def check_match(lineset): 39 | linediff = np.diff(lineset) 40 | x = 0 41 | cnt = 0 42 | for line in linediff: 43 | # Within 5 px of the other (allowing for minor image errors) 44 | if np.abs(line - x) < 5: 45 | cnt += 1 46 | else: 47 | cnt = 0 48 | x = line 49 | return cnt == 5 50 | 51 | 52 | # prunes a set of lines to 7 in consistent increasing order (chessboard) 53 | def prune_lines(lineset): 54 | linediff = np.diff(lineset) 55 | x = 0 56 | cnt = 0 57 | start_pos = 0 58 | for i, line in enumerate(linediff): 59 | # Within 5 px of the other (allowing for minor image errors) 60 | if np.abs(line - x) < 5: 61 | cnt += 1 62 | if cnt == 5: 63 | end_pos = i + 2 64 | return lineset[start_pos:end_pos] 65 | else: 66 | cnt = 0 67 | x = line 68 | start_pos = i 69 | return lineset 70 | 71 | 72 | # return skeletonized 1d array (thin to single value, favor to the right) 73 | def skeletonize_1d(arr): 74 | _arr = arr.copy() # create a copy of array to modify without destroying original 75 | # Go forwards 76 | for i in range(_arr.size - 1): 77 | # Will right-shift if they are the same 78 | if arr[i] <= _arr[i + 1]: 79 | _arr[i] = 0 80 | 81 | # Go reverse 82 | for i in np.arange(_arr.size - 1, 0, -1): 83 | if _arr[i - 1] > _arr[i]: 84 | _arr[i] = 0 85 | return _arr 86 | 87 | 88 | # returns pixel indices for the 7 internal chess lines in x and y axes 89 | def get_chess_lines(hdx, hdy, hdx_thresh, hdy_thresh): 90 | # Blur 91 | gausswin = scipy.signal.gaussian(21, 4) 92 | gausswin /= np.sum(gausswin) 93 | 94 | # Blur where there is a strong horizontal or vertical line (binarize) 95 | blur_x = np.convolve(hdx > hdx_thresh, gausswin, mode='same') 96 | blur_y = np.convolve(hdy > hdy_thresh, gausswin, mode='same') 97 | 98 | skel_x = skeletonize_1d(blur_x) 99 | skel_y = skeletonize_1d(blur_y) 100 | 101 | # Find points on skeletonized arrays (where returns 1-length tuple) 102 | lines_x = np.where(skel_x)[0] # vertical lines 103 | lines_y = np.where(skel_y)[0] # horizontal lines 104 | 105 | # Prune inconsistent lines 106 | lines_x = prune_lines(lines_x) 107 | lines_y = prune_lines(lines_y) 108 | 109 | is_match = len(lines_x) == 7 and len(lines_y) == 7 and check_match(lines_x) and check_match(lines_y) 110 | 111 | return lines_x, lines_y, is_match 112 | 113 | 114 | # Gets a numpy grayscale image and returns lines_x and lines_y 115 | def detect_chessboard(img): 116 | grey = img 117 | 118 | dX = gradient_x(grey) 119 | dY = gradient_y(grey) 120 | 121 | dX_pos = tf.clip_by_value(dX, 0., 255., name="dx_positive") 122 | dX_neg = tf.clip_by_value(dX, -255., 0., name="dx_negative") 123 | dY_pos = tf.clip_by_value(dY, 0., 255., name="dy_positive") 124 | dY_neg = tf.clip_by_value(dY, -255., 0., name="dy_negative") 125 | 126 | dX_hough = tf.reduce_sum(dX_pos, 0) * tf.reduce_sum(-dX_neg, 0) / (grey.shape[0] * grey.shape[0]) 127 | dY_hough = tf.reduce_sum(dY_pos, 1) * tf.reduce_sum(-dY_neg, 1) / (grey.shape[1] * grey.shape[1]) 128 | 129 | # Arbitrarily choose half of max value as threshold, since they're such strong responses 130 | dX_hough_thresh = tf.reduce_max(dX_hough) * 0.5 131 | dY_hough_thresh = tf.reduce_max(dY_hough) * 0.5 132 | 133 | lines_x, lines_y, is_match = get_chess_lines(tf.keras.backend.flatten(dX_hough), 134 | tf.keras.backend.flatten(dY_hough), 135 | dX_hough_thresh * .9, 136 | dY_hough_thresh * .9) 137 | 138 | if is_match: 139 | print("Chessboard found") 140 | else: 141 | print("Couldn't find Chessboard") 142 | 143 | return is_match, lines_x, lines_y 144 | 145 | 146 | # Split up input grayscale array into 64 tiles stacked in a 3D matrix using the chess linesets 147 | def get_chess_tiles(img, lines_x, lines_y): 148 | # Find average square size, round to a whole pixel for determining edge pieces sizes 149 | stepx = np.int32(np.round(np.mean(np.diff(lines_x)))) 150 | stepy = np.int32(np.round(np.mean(np.diff(lines_y)))) 151 | 152 | # Pad edges as needed to fill out chessboard (for images that are partially over-cropped) 153 | # print stepx, stepy 154 | # print "x",lines_x[0] - stepx, "->", lines_x[-1] + stepx, a.shape[1] 155 | # print "y", lines_y[0] - stepy, "->", lines_y[-1] + stepy, a.shape[0] 156 | padr_x = 0 157 | padl_x = 0 158 | padr_y = 0 159 | padl_y = 0 160 | 161 | if lines_x[0] - stepx < 0: 162 | padl_x = np.abs(lines_x[0] - stepx) 163 | if lines_x[-1] + stepx > img.shape[1] - 1: 164 | padr_x = np.abs(lines_x[-1] + stepx - img.shape[1]) 165 | if lines_y[0] - stepy < 0: 166 | padl_y = np.abs(lines_y[0] - stepy) 167 | if lines_y[-1] + stepx > img.shape[0] - 1: 168 | padr_y = np.abs(lines_y[-1] + stepy - img.shape[0]) 169 | 170 | # New padded array 171 | # print "Padded image to", ((padl_y,padr_y),(padl_x,padr_x)) 172 | a2 = np.pad(img, ((padl_y, padr_y), (padl_x, padr_x)), mode='edge') 173 | 174 | setsx = np.hstack([lines_x[0] - stepx, lines_x, lines_x[-1] + stepx]) + padl_x 175 | setsy = np.hstack([lines_y[0] - stepy, lines_y, lines_y[-1] + stepy]) + padl_y 176 | 177 | a2 = a2[setsy[0]:setsy[-1], setsx[0]:setsx[-1]] 178 | setsx -= setsx[0] 179 | setsy -= setsy[0] 180 | # display_array(a2, rng=[0,255]) 181 | # print "X:",setsx 182 | # print "Y:",setsy 183 | 184 | # Matrix to hold images of individual squares (in grayscale) 185 | # print "Square size: [%g, %g]" % (stepy, stepx) 186 | squares = np.zeros([np.round(stepy), np.round(stepx), 64], dtype=np.uint8) 187 | 188 | # For each row 189 | for i in range(0, 8): 190 | # For each column 191 | for j in range(0, 8): 192 | # Vertical lines 193 | x1 = setsx[i] 194 | x2 = setsx[i + 1] 195 | padr_x = 0 196 | padl_x = 0 197 | padr_y = 0 198 | padl_y = 0 199 | 200 | if (x2 - x1) > stepx: 201 | if i == 7: 202 | x1 = x2 - stepx 203 | else: 204 | x2 = x1 + stepx 205 | elif (x2 - x1) < stepx: 206 | if i == 7: 207 | # right side, pad right 208 | padr_x = stepx - (x2 - x1) 209 | else: 210 | # left side, pad left 211 | padl_x = stepx - (x2 - x1) 212 | # Horizontal lines 213 | y1 = setsy[j] 214 | y2 = setsy[j + 1] 215 | 216 | if (y2 - y1) > stepy: 217 | if j == 7: 218 | y1 = y2 - stepy 219 | else: 220 | y2 = y1 + stepy 221 | elif (y2 - y1) < stepy: 222 | if j == 7: 223 | # right side, pad right 224 | padr_y = stepy - (y2 - y1) 225 | else: 226 | # left side, pad left 227 | padl_y = stepy - (y2 - y1) 228 | # slicing a, rows sliced with horizontal lines, cols by vertical lines so reversed 229 | # Also, change order so its A1,B1...H8 for a white-aligned board 230 | # Apply padding as defined previously to fit minor pixel offsets 231 | squares[:, :, (7 - j) * 8 + i] = np.pad(a2[y1:y2, x1:x2], ((padl_y, padr_y), (padl_x, padr_x)), mode='edge') 232 | return squares 233 | --------------------------------------------------------------------------------