40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
import captcha_settings
|
||
import os
|
||
import random
|
||
|
||
from captcha.image import ImageCaptcha
|
||
from PIL import Image
|
||
|
||
from tqdm import trange
|
||
|
||
def random_captcha_text(char_set=captcha_settings.NUMBER + captcha_settings.ALPHABET, captcha_size=4):
|
||
captcha_text = []
|
||
for i in range(captcha_size):
|
||
c = random.choice(char_set)
|
||
captcha_text.append(c)
|
||
return "".join(captcha_text)
|
||
|
||
def gen_captcha_text_and_image():
|
||
image = ImageCaptcha()
|
||
captcha_text = random_captcha_text()
|
||
captcha_image = Image.open(image.generate(captcha_text))
|
||
return captcha_text, captcha_image
|
||
|
||
if __name__ == "__main__":
|
||
dataset_type = input("请输入数据集类型（1 - train / 2 - test / 3 - predict）：")
|
||
dataset_len = input("请输入数据集长度：")
|
||
|
||
paths = [captcha_settings.TRAIN_DATASET_PATH, captcha_settings.TEST_DATASET_PATH, captcha_settings.PREDICT_DATASET_PATH]
|
||
|
||
dataset_type = int(dataset_type)
|
||
count = int(dataset_len)
|
||
path = (
|
||
paths[dataset_type - 1]
|
||
)
|
||
if not os.path.exists(path):
|
||
os.makedirs(path)
|
||
for i in trange(count):
|
||
text, image = gen_captcha_text_and_image()
|
||
filename = f"{str(i).zfill(5)}_{text}.png"
|
||
image.save(path + os.path.sep + filename)
|