font-identifier / arrange_train_test_images.py
Gabor Cselle
Added more fonts, bold and italics
db35ca2
raw
history blame
1.38 kB
# moves the font images into train and test folders
# TODO(gabor): maybe we should copy these instead, so we don't have to regenerate the images every times?
import os
import shutil
import random
from consts import TRAIN_TEST_IMAGES_DIR, GEN_IMAGES_DIR
train_dir = os.path.join(TRAIN_TEST_IMAGES_DIR, 'train')
test_dir = os.path.join(TRAIN_TEST_IMAGES_DIR, 'test')
# create directories if they don't exist
os.makedirs(train_dir, exist_ok=True)
os.makedirs(test_dir, exist_ok=True)
# make a list of all the font names
fonts = [f.split('_')[0] for f in os.listdir(GEN_IMAGES_DIR) if f.endswith('.png')]
fonts = list(set(fonts)) # getting unique font names
for font in fonts:
font_train_dir = os.path.join(train_dir, font)
font_test_dir = os.path.join(test_dir, font)
os.makedirs(font_train_dir, exist_ok=True)
os.makedirs(font_test_dir, exist_ok=True)
font_files = [f for f in os.listdir(GEN_IMAGES_DIR) if f.startswith(font + "_")]
random.shuffle(font_files)
train_files = font_files[:int(0.8 * len(font_files))]
test_files = font_files[int(0.8 * len(font_files)):]
# Move training files
for train_file in train_files:
shutil.move(os.path.join(GEN_IMAGES_DIR, train_file), font_train_dir)
# Move test files
for test_file in test_files:
shutil.move(os.path.join(GEN_IMAGES_DIR, test_file), font_test_dir)