YuzuMarker.FontDetection / font_ds_stat.py
gyrojeff's picture
feat: add exclusion rules
a9b34cd
import sys
import traceback
import pickle
import os
import concurrent.futures
from tqdm import tqdm
from font_dataset.font import load_fonts
from font_dataset.layout import generate_font_image
from font_dataset.text import CorpusGeneratorManager
from font_dataset.background import background_image_generator
cjk_ratio = 3
train_cnt = 100
val_cnt = 5
test_cnt = 30
train_cnt_cjk = int(train_cnt * cjk_ratio)
val_cnt_cjk = int(val_cnt * cjk_ratio)
test_cnt_cjk = int(test_cnt * cjk_ratio)
dataset_path = "./dataset/font_img"
os.makedirs(dataset_path, exist_ok=True)
fonts, exclusion_rule = load_fonts()
cnt = 0
for font in fonts:
if exclusion_rule(font):
print(f"Excluded font: {font.path}")
continue
if font.language == "CJK":
cnt += cjk_ratio
else:
cnt += 1
print("Total training images:", train_cnt * cnt)
print("Total validation images:", val_cnt * cnt)
print("Total testing images:", test_cnt * cnt)
if os.path.exists(os.path.join(dataset_path, "train")):
num_file_train = len(os.listdir(os.path.join(dataset_path, "train")))
else:
num_file_train = 0
if os.path.exists(os.path.join(dataset_path, "val")):
num_file_val = len(os.listdir(os.path.join(dataset_path, "val")))
else:
num_file_val = 0
if os.path.exists(os.path.join(dataset_path, "test")):
num_file_test = len(os.listdir(os.path.join(dataset_path, "test")))
else:
num_file_test = 0
print("Total files generated:", num_file_train + num_file_val + num_file_test)
print("Total files target:", (train_cnt + val_cnt + test_cnt) * cnt * 2)
print(
f"{(num_file_train + num_file_val + num_file_test) / ((train_cnt + val_cnt + test_cnt) * cnt * 2) * 100:.2f}% completed"
)