main.py - randomcrap - random crap programs of varying quality
(HTM) git clone git://git.codemadness.org/randomcrap
(DIR) Log
(DIR) Files
(DIR) Refs
(DIR) README
(DIR) LICENSE
---
main.py (15971B)
---
1 from skimage import data, img_as_float
2
3 from operator import itemgetter
4
5 import cv2
6
7 import glob
8 import math
9 import numpy as np
10 import os
11 import sys
12 import time
13
14 config = {
15 "debug": True,
16 "debugimage": True,
17 "checkmoveindicator": True,
18 "boardstyle": "", # "5334_polgar", "lichess", "woodpecker"
19 "guessstyle": True,
20 "flip": False, # scanned board is flipped?
21 "flipdisplay": False # output should be flipped?
22 }
23
24 CHESS_BOARD_STYLE_DIR = glob.glob('board_styles/*.png')
25
26 boardstyles = {}
27
28 # Initiate SIFT detector
29 sift = cv2.SIFT_create()
30
31 # [ x, y, square color, piece ],
32 # table to read each piece on a different squared color.
33 # and also an empty square for each color.
34 readtab = [
35 # black
36 [ 0, 0, "w", "r" ],
37 [ 1, 0, "b", "n" ],
38 [ 2, 0, "w", "b" ],
39 [ 3, 0, "b", "q" ],
40 [ 4, 0, "w", "k" ],
41 [ 5, 0, "b", "b" ],
42 [ 6, 0, "w", "n" ],
43 [ 7, 0, "b", "r" ],
44 [ 0, 1, "b", "p" ],
45 [ 1, 1, "w", "p" ],
46 [ 0, 2, "w", "" ], # empty
47 [ 1, 2, "b", "" ],
48 [ 3, 3, "w", "q" ],
49 [ 4, 3, "b", "k" ],
50 # white
51 [ 0, 7, "b", "R" ],
52 [ 1, 7, "w", "N" ],
53 [ 2, 7, "b", "B" ],
54 [ 3, 7, "w", "Q" ],
55 [ 4, 7, "b", "K" ],
56 [ 5, 7, "w", "B" ],
57 [ 6, 7, "b", "N" ],
58 [ 7, 7, "w", "R" ],
59 [ 0, 6, "w", "P" ],
60 [ 1, 6, "b", "P" ],
61 [ 0, 5, "b", "" ], # empty
62 [ 1, 5, "w", "" ],
63 [ 3, 4, "b", "Q" ],
64 [ 4, 4, "w", "K" ],
65 ]
66
67 # read input board setup data to compare to.
68 for path in CHESS_BOARD_STYLE_DIR:
69 basename = os.path.basename(path)
70 name = basename[:-4] # remove ".png"
71
72 # do not load board styles that are not used anyway.
73 if "guessstyle" in config and config["guessstyle"] == False and \
74 "boardstyle" in config and config["boardstyle"] != "" and name != config["boardstyle"]:
75 continue
76
77 # read board style image
78 img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
79
80 # split fields as images in board image
81 height = len(img)
82 width = len(img[0])
83 # square width and height, fixed-size, 8x8 squares.
84 sw = width / 8
85 sh = height / 8
86
87 style = {
88 "w": {},
89 "b": {}
90 }
91
92 # get each field and expected configuration, store them as template:
93 # pieces (per square), empty squares.
94 for t in readtab:
95 # x, y position is floored/truncated
96 x1 = int(t[0] * sw)
97 y1 = int(t[1] * sh)
98 x2 = int(x1 + sw)
99 y2 = int(y1 + sh)
100
101 region = img[y1:y2, x1:x2]
102
103 # style[squarecolor][piece] = image data
104 style[t[2]][t[3]] = region
105
106 boardstyles[name] = style
107
108 def debug(s):
109 if not config["debug"]:
110 return
111 print("DEBUG: " + s, file=sys.stderr, flush=True)
112
113 def iswhitesquare(x, y):
114 if (y & 1) == 0:
115 return ((x + 1) & 1)
116 else:
117 return (x & 1)
118
119 def img_resize(img, sw, sh, dw, dh):
120 # resize piece relative to board size (approximately).
121 #interp = cv2.INTER_CUBIC # enlarge: use cv2.INTER_LINEAR (fast) or cv2.INTER_CUBIC (slow)
122
123 # NOTE: not scaled for quality, because we use MSE. Sampling would change pixel value and
124 # reduce matching results.
125 interp = cv2.INTER_AREA
126 #interp = cv2.INTER_LINEAR
127
128 return cv2.resize(img, (dw, dh), interpolation=interp)
129
130 # https://docs.opencv.org/3.4/d1/de0/tutorial_py_feature_homography.html
131 def find_pieces_flann(desimg, despiece):
132 FLANN_INDEX_KDTREE = 1
133 index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 5)
134 search_params = dict(checks = 50)
135 flann = cv2.FlannBasedMatcher(index_params, search_params)
136
137 matches = flann.knnMatch(desimg, despiece, k=2)
138 # store all the good matches as per Lowe's ratio test.
139 # https://docs.opencv.org/3.4/d5/d6f/tutorial_feature_flann_matcher.html
140 good = []
141 for m, n in matches:
142 # NOTE: was: 0.7 (reference value from above webpage)
143 if m.distance < 0.3 * n.distance:
144 good.append(m)
145
146 return good
147
148 def board2txt(boardstate, flip):
149 output = ""
150 empty = 0
151 matches = boardstate["matches"]
152 nrows = len(matches)
153 for iy in range(nrows):
154 if flip:
155 y = 7 - iy
156 else:
157 y = iy
158 if iy > 0:
159 output = output + ("\n---+---+---+---+---+---+----+---\n")
160 row = matches[y]
161 nitems = len(row)
162 for ix in range(nitems):
163 if flip:
164 x = 7 - ix
165 else:
166 x = ix
167 match = matches[y][x]
168 c = match["piece"]
169 if c == "":
170 c = " "
171
172 output = output + (" %s |" % (c))
173
174 output = output + ("\n---+---+---+---+---+---+----+---\n")
175
176 return output
177
178 def board2fen(boardstate, flip):
179 output = ""
180 empty = 0
181 matches = boardstate["matches"]
182 nrows = len(matches)
183 for iy in range(nrows):
184 if flip:
185 y = 7 - iy
186 else:
187 y = iy
188 if empty > 0:
189 output = output + ("%d" % empty)
190 empty = 0
191
192 if iy > 0:
193 output = output + ("/")
194
195 row = matches[y]
196 nitems = len(row)
197 for ix in range(nitems):
198 if flip:
199 x = 7 - ix
200 else:
201 x = ix
202 match = matches[y][x]
203 if not "piece" in match or match["piece"] == "":
204 empty = empty + 1
205 continue
206
207 if empty > 0:
208 output = output + ("%d" % empty)
209 empty = 0
210 output = output + match["piece"]
211
212 if empty > 0:
213 output = output + ("%d" % empty)
214
215 castle = "KQkq"
216 if "castle" in boardstate:
217 castle = boardstate["castle"]
218 if castle == "":
219 castle = "-"
220
221 # workaround: we do not know the turn or move number (for now).
222 output = output + " " + boardstate["side"] + " " + castle + " - 0 1"
223
224 return output
225
226 def board2lichess(boardstate, flip):
227 output = board2fen(boardstate, flip)
228 output = output.replace(" ", "_")
229 output = "https://lichess.org/editor/" + output + "?color=white"
230 # output = "https://lichess.org/analysis/" + output + "?color=white"
231
232 return output
233
234 # sift cache for pieces, if dimensions are the same.
235 siftcache = {}
236 def detectpieces(boardimage, boardstyle, onlykings, flip):
237 status = True
238
239 # pieces from template
240 chesspieceimages = boardstyles[boardstyle]
241 boardheight = len(boardimage)
242 boardwidth = len(boardimage[0])
243
244 # pieces (scaled) to board size.
245 piecerelative = {
246 "w": {},
247 "b": {}
248 }
249
250 pieceimage = chesspieceimages["b"]["K"]
251 piecewidth = len(pieceimage[0])
252 usepiecewidth = int(boardwidth / 8)
253
254 # resize all template piece images.
255 # scale pieces relative to the board
256 for color in chesspieceimages:
257 for piece in chesspieceimages[color]:
258 if onlykings and piece != "K" and piece != "k":
259 continue
260 pieceimage = chesspieceimages[color][piece]
261 pieceheight = len(pieceimage)
262 piecewidth = len(pieceimage[0])
263
264 ratio = float(piecewidth) / pieceheight
265 dw = usepiecewidth
266 dh = int(dw * ratio)
267 img = img_resize(pieceimage, piecewidth, pieceheight, dw, dh)
268
269 piecerelative[color][piece] = img # use new size
270
271 sw = int(boardwidth / 8)
272 sh = int(boardheight / 8)
273
274 piecessift = {}
275 for color in ["b", "w"]:
276 piecessift[color] = {}
277 for piece in piecerelative[color]:
278 if onlykings and piece != "K" and piece != "k":
279 continue
280
281 key = ("%s|%s|%s|%s|%s" % (boardstyle, sw, sh, color, piece))
282 if key in siftcache:
283 s = siftcache[key]
284 piecessift[color][piece] = s
285 continue
286
287 h = piecerelative[color][piece]
288 pimg = cv2.cvtColor(h, cv2.COLOR_BGR2GRAY)
289 kp, des = sift.detectAndCompute(pimg, None)
290 if des is None or len(des) < 2:
291 continue
292 s = {"kp": kp, "des": des}
293 piecessift[color][piece] = s
294
295 # store in global cache.
296 siftcache[key] = s
297
298
299 boardstate = {"matches":[], "side": "w"}
300 matches = []
301 for y in range(8):
302 row = []
303 for x in range(8):
304 # default: no match
305 match = { "color": "", "piece": "", "score": 0.0 }
306 row.append(match)
307 matches.append(row)
308
309 for iy in range(8):
310 if flip:
311 y = 7 - iy
312 else:
313 y = iy
314 for ix in range(8):
315 if flip:
316 x = 7 - ix
317 else:
318 x = ix
319 # x, y position is floored/truncated
320 x1 = int(x * sw)
321 y1 = int(y * sh)
322 x2 = int(x1 + sw)
323 y2 = int(y1 + sh)
324
325 region = boardimage[y1:y2, x1:x2]
326
327 regiongray = cv2.cvtColor(region, cv2.COLOR_BGR2GRAY)
328 kpregion, desregion = sift.detectAndCompute(regiongray, None)
329 if desregion is None:
330 continue
331
332 bestmatch = -1
333 match = {}
334 color = "w"
335 if not iswhitesquare(x, y):
336 color = "b"
337
338 for piece in piecerelative[color]:
339 if onlykings and piece != "K" and piece != "k":
340 continue
341
342 h = piecerelative[color][piece]
343 if not piece in piecessift[color]:
344 continue
345 s = piecessift[color][piece]
346 if not "des" in s or s["des"] is None or len(s["des"]) < 2:
347 continue
348 despiece = s["des"]
349
350 # FLANN
351 fmatches = find_pieces_flann(desregion, despiece)
352 score = len(fmatches)
353 if score > 0: # threshold
354 # higher matches is better
355 if bestmatch == -1 or score > bestmatch:
356 bestmatch = score
357 match = { "color": color, "piece": piece, "score": score }
358
359 if "score" in match:
360 matches[y][x] = match
361
362
363 boardstate["matches"] = matches
364 boardstate["status"] = status
365 return boardstate
366
367 def boarddetect(img):
368 areas = []
369 gimg = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
370
371 # dilate/erode: improves detection with broken lines.
372 kernel = np.ones((3,3), np.uint8)
373 d_im = cv2.dilate(gimg, kernel, iterations=1)
374 e_im = cv2.erode(d_im, kernel, iterations=1)
375
376 # canny edge (without gaussian blur)
377 cimg = cv2.Canny(e_im, 100, 200)
378
379 contours, hierarchy = cv2.findContours(cimg,
380 cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) # simple only stores points
381
382 for c in contours:
383 x, y, w, h = cv2.boundingRect(c)
384 # skip too small areas (for example 8 * piece size)
385 if w < 200 or h < 200:
386 continue
387
388 # quick check: must be somewhat square at least
389 ratio = w / h
390 if ratio < 0.95 or ratio > 1.05:
391 continue
392
393 areas.append({"x": x, "y": y, "w": w, "h": h })
394
395 return areas
396
397 def isvalidboard(boardstate):
398 wk = False
399 bk = False
400 matches = boardstate["matches"]
401 # check if a white and black king is found.
402 for y in range(len(matches)):
403 row = matches[y]
404 for x in range(len(row)):
405 c = row[x]["piece"]
406 if c == "K":
407 wk = True
408 if c == "k":
409 bk = True
410
411 return wk and bk
412
413 # check castle rights based on the image, but it does not check if in check
414 # or if castling is legal.
415 def fixcastlerights(boardstate):
416 m = boardstate["matches"] # 8x8 matches
417
418 wk = "K"
419 wq = "Q"
420 bk = "k"
421 bq = "q"
422
423 # check if king moved.
424 if m[7][4]["piece"] != "K": # white king
425 wk = ""
426 wq = ""
427 else:
428 # check if rooks moved.
429 if m[7][0]["piece"] != "R":
430 wq = ""
431 if m[7][7]["piece"] != "R":
432 wk = ""
433
434 if m[0][4]["piece"] != "k": # black king
435 bk = ""
436 bq = ""
437 else:
438 # check if rooks moved.
439 if m[0][0]["piece"] != "r":
440 bq = ""
441 if m[0][7]["piece"] != "r":
442 bk = ""
443
444 boardstate["castle"] = wk + wq + bk + bq
445
446 # main
447 filename = "stdin"
448 title = "stdin"
449 fullimg = cv2.imread("/dev/stdin")
450 fullimgh, fullimgw, _ = fullimg.shape
451
452 areas = boarddetect(fullimg)
453
454 # try to somewhat sort by visible (columns: left bottom, to right down).
455 areas = sorted(areas, key=itemgetter("x", "y"))
456 debug("areas: %d" % len(areas))
457
458 # filter away the most obvious duplicate areas.
459 uareas = [] # unique areas
460 maxdiff = 10
461 arealen = len(areas)
462 for i in range(arealen):
463 dup = False
464 for j in range(i, arealen):
465 if i == j:
466 continue # self
467 if areas[j]["x"] >= (areas[i]["x"] - maxdiff) and areas[j]["x"] <= (areas[i]["x"] + maxdiff) and \
468 areas[j]["y"] >= (areas[i]["y"] - maxdiff) and areas[j]["y"] <= (areas[i]["y"] + maxdiff):
469 dup = True
470 if not dup:
471 uareas.append(areas[i])
472 areas = uareas
473 # if no areas found, use whole image.
474 if len(areas) == 0:
475 areas.append({"x": 0, "y": 0, "w": fullimgw, "h": fullimgh })
476
477 debug("unique areas: %d" % len(areas))
478
479 if config["debug"] or config["debugimage"]:
480 debugimg = fullimg
481
482 flip = "flip" in config and config["flip"]
483 dareas = []
484 boardstates = []
485 for area in areas:
486 x1 = area["x"] + 1
487 y1 = area["y"] + 1
488 x2 = x1 + area["w"] - 1
489 y2 = y1 + area["h"] - 1
490 region = fullimg[y1:y2, x1:x2]
491
492 boardstyle = config["boardstyle"]
493 if config["guessstyle"]:
494 # detect board style, use the boardstyle with the highest score.
495 # only scans for king pieces.
496 best = 0
497 beststyle = ""
498 for boardstyle in boardstyles:
499 boardstate = detectpieces(region, boardstyle, True, flip)
500 if boardstate["status"] == False:
501 continue
502 l = len(boardstate["matches"])
503 wk = 0.0
504 bk = 0.0
505 for y in range(l):
506 row = boardstate["matches"][y]
507 for x in range(len(row)):
508 match = row[x]
509 c = match["piece"]
510 score = match["score"]
511 if c == "K":
512 wk = score
513 if c == "k":
514 bk = score
515 if wk + bk > best:
516 best = wk + bk
517 beststyle = boardstyle
518
519 debug("best guess for board style: \"%s\"" % (beststyle))
520 boardstyle = beststyle
521
522 if boardstyle != "":
523 starttime = time.perf_counter_ns()
524 boardstate = detectpieces(region, boardstyle, False, flip)
525 endtime = time.perf_counter_ns()
526 ms = (endtime - starttime) / 1000000
527 debug("detectpieces() took: %d ms" % ms)
528
529 boardstate["area"] = area
530 if boardstate["status"] == False:
531 boardstates.append(boardstate)
532 continue
533
534 # check move indicator: check for any figure next to the region (about one square).
535 if boardstyle in ["5334_polgar", "woodpecker"] and config["checkmoveindicator"]:
536 # top right: black to move
537 mx1 = x2 + 4
538 mx2 = int(x2 + (area["w"] / 8))
539 my1 = y1
540 my2 = int(y1 + (area["h"] / 8))
541 if mx2 < fullimgw and my2 < fullimgh:
542 region = fullimg[my1:my2, mx1:mx2]
543 gimg = cv2.cvtColor(region, cv2.COLOR_BGR2GRAY)
544 canny = cv2.Canny(gimg, 100, 100) # canny edge (without gaussian blur)
545 contours, hierarchy = cv2.findContours(canny,
546 cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) # simple only stores points
547 blen = len(contours)
548
549 # bottom right: white to move
550 mx1 = x2 + 4
551 mx2 = int(mx1 + (area["w"] / 8))
552 my1 = int(y2 - (area["h"] / 8))
553 my2 = y2
554
555 if mx2 < fullimgw and my2 < fullimgh:
556 region = fullimg[my1:my2, mx1:mx2]
557 gimg = cv2.cvtColor(region, cv2.COLOR_BGR2GRAY)
558 canny = cv2.Canny(gimg, 100, 100) # canny edge (without gaussian blur)
559 contours, hierarchy = cv2.findContours(canny,
560 cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) # simple only stores points
561 wlen = len(contours)
562
563 if blen > 0 and blen > wlen:
564 boardstate["side"] = "b" # black to move
565 elif wlen > 0 and wlen > blen:
566 boardstate["side"] = "w" # white to move"
567
568 if boardstyle == "":
569 continue
570
571 if not isvalidboard(boardstate):
572 continue
573
574 fixcastlerights(boardstate)
575 boardstates.append(boardstate)
576
577 flipdisplay = "flipdisplay" in config and config["flipdisplay"]
578 # print(board2txt(boardstate, flipdisplay), flush=True)
579 print(board2fen(boardstate, flipdisplay), flush=True)
580
581 # draw a debug image with detected board regions and pieces.
582 if config["debug"] or config["debugimage"]:
583 for b in boardstates:
584 if "area" in b:
585 area = b["area"]
586 sw = area["w"] / 8 # square width
587 sh = area["h"] / 8 # square height
588
589 # board area
590 x1 = area["x"] + 1
591 y1 = area["y"] + 1
592 x2 = x1 + area["w"] - 1
593 y2 = y1 + area["h"] - 1
594 color = (0, 255, 255)
595 if b["status"] == False:
596 color = (0, 0, 255)
597 cv2.rectangle(debugimg, (x1, y1), (x2, y2), color, 2)
598
599 # draw pieces as text characters.
600 if "matches" in b:
601 matches = b["matches"]
602 l = len(matches)
603 for y in range(len(matches)):
604 row = matches[y]
605 for x in range(len(row)):
606 m = row[x]
607 c = m["piece"]
608 if c == "":
609 continue
610
611 px = int(x1 + (x * sw))
612 py = int(y1 + (y * sh))
613
614 # text color for white or black.
615 if c.isupper():
616 color = (255, 0, 0) # BGR: blue
617 else:
618 color = (0, 0, 255) # BGR: red
619
620 # fontscale = 1.0 / min(area["w"], area["h"]) * (sh * 6)
621 fontscale = 1.0
622 cv2.putText(debugimg, c, (px, py + int(sh)), cv2.FONT_HERSHEY_SIMPLEX, fontscale, color, 2, cv2.LINE_AA)
623
624 if boardstyle in ["5334_polgar", "woodpecker"] and config["checkmoveindicator"]:
625 # side to move indicator.
626 if "side" in b and b["side"] == "w":
627 mx1 = x2 + 4
628 mx2 = int(mx1 + sw)
629 my1 = int(y2 - sh)
630 my2 = y2
631 cv2.rectangle(debugimg, (mx1, my1), (mx2, my2), (255, 0, 0), 2)
632 elif "side" in b and b["side"] == "b":
633 mx1 = x2 + 4
634 mx2 = int(x2 + sw)
635 my1 = y1
636 my2 = int(y1 + sh)
637 cv2.rectangle(debugimg, (mx1, my1), (mx2, my2), (0, 0, 255), 2)
638
639 filename = "/tmp/debug.png"
640 cv2.imwrite(filename, debugimg)
641 debug("debug image written: %s" % (filename))