-
Notifications
You must be signed in to change notification settings - Fork 1
/
generate_tiles.py
executable file
·95 lines (83 loc) · 3.32 KB
/
generate_tiles.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
#!/usr/bin/env python3
# Generate tile images for all chessboard images in input folder
# Used for building training datasets
import os
from glob import glob
import argparse
import math
import numpy as np
import PIL.Image
from constants import CHESSBOARDS_DIR, TILES_DIR, USE_GRAYSCALE
from chessboard_image import get_chessboard_tiles
OVERWRITE = False
def _img_filename_prefix(chessboard_img_path):
""" part of the image filename that shows which piece is on which square:
RRqpBnNr-QKPkrQPK-PpbQnNB1-nRRBpNpk-Nqprrpqp-kKKbNBPP-kQnrpkrn-BKRqbbBp
"""
return chessboard_img_path.split("/")[4][:-4]
def _img_sub_dir(chessboard_img_path):
""" The sub-directory where the chessboard tile images will be saved
"""
sub_dir = chessboard_img_path.split("/")[3]
return os.path.join(TILES_DIR, sub_dir)
def _img_save_dir(chessboard_img_path):
""" The directory within the sub-directory that will contain the
tile images
"""
return os.path.join(
_img_sub_dir(chessboard_img_path),
_img_filename_prefix(chessboard_img_path),
)
def save_tiles(tiles, chessboard_img_path):
""" Saves all 64 tiles as 32x32 PNG files with this naming convention:
a1_R.png (white rook on a1)
d8_q.png (black queen on d8)
c4_1.png (nothing on c4)
"""
sub_dir = _img_sub_dir(chessboard_img_path)
if not os.path.exists(sub_dir):
os.makedirs(sub_dir)
img_save_dir = _img_save_dir(chessboard_img_path)
print("\tSaving tiles to {}\n".format(img_save_dir))
if not os.path.exists(img_save_dir):
os.makedirs(img_save_dir)
piece_positions = _img_filename_prefix(chessboard_img_path).split('-')
files = 'abcdefgh'
for i in range(64):
piece = piece_positions[math.floor(i / 8)][i % 8]
sqr_id = '{}{}'.format(files[i % 8], 8 - math.floor(i / 8))
tile_img_filename = '{}/{}_{}.png'.format(img_save_dir, sqr_id, piece)
tiles[i].save(tile_img_filename, format='PNG')
def generate_tiles_from_all_chessboards():
""" Generates 32x32 PNGs for each square of all chessboards
in CHESSBOARDS_DIR
"""
if not os.path.exists(TILES_DIR):
os.makedirs(TILES_DIR)
chessboard_img_filenames = glob("{}/*/*.png".format(CHESSBOARDS_DIR))
num_chessboards = len(chessboard_img_filenames)
num_success = 0
num_skipped = 0
num_failed = 0
for i, chessboard_img_path in enumerate(chessboard_img_filenames):
print("%3d/%d %s" % (i + 1, num_chessboards, chessboard_img_path))
img_save_dir = _img_save_dir(chessboard_img_path)
if os.path.exists(img_save_dir) and not OVERWRITE:
print("\tIgnoring existing {}\n".format(img_save_dir))
num_skipped += 1
continue
tiles = get_chessboard_tiles(chessboard_img_path, use_grayscale=USE_GRAYSCALE)
if len(tiles) != 64:
print("\t!! Expected 64 tiles. Got {}\n".format(len(tiles)))
num_failed += 1
continue
save_tiles(tiles, chessboard_img_path)
num_success += 1
print(
'Processed {} chessboard images ({} generated, {} skipped, {} failed)'.format(
num_chessboards, num_success, num_skipped, num_failed
)
)
if __name__ == '__main__':
np.set_printoptions(suppress=True, precision=2)
generate_tiles_from_all_chessboards()