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)
|