JacobLinCool/captcha-recognizer

1import os2from PIL import Image, ImageFont, ImageDraw3from random import choice, random4from string import ascii_lowercase5from src.shared import genereated_dir6 7WIDTH, HEIGHT = 108, 308 9count = 10010 11font = ImageFont.truetype(12    os.path.join(os.path.dirname(__file__), "..", "fonts", "NotoSans-Regular.ttf"), 1813)14 15 16def gen():17    image = Image.new("RGB", (WIDTH, HEIGHT))18 19    if random() > 0.5:20        # alphabetical21        a = choice(ascii_lowercase)22        b = choice(ascii_lowercase)23        c = choice(ascii_lowercase)24        d = choice(ascii_lowercase)25        ans = a + b + c + d26 27        image = draw(image, [a, b, c, d])28 29    else:30        # arithmetic31        a = choice(range(10))32        b = choice(range(10))33        op = choice(["+", "-", "x"])34 35        if op == "+":36            ans = f"{a}+{b}="37        elif op == "-":38            ans = f"{a}-{b}="39        else:40            ans = f"{a}x{b}="41 42        if op == "x" and random() > 0.5:43            op = "X"44 45        image = draw(image, [str(a), op, str(b), "="])46 47    return image, str(ans)48 49 50def draw(image: Image, text: list[str]) -> Image:51    draw = ImageDraw.Draw(image)52    draw.rectangle((0, 0, WIDTH, HEIGHT), fill=(255, 255, 255))53 54    for i, t in enumerate(text):55        txt = Image.new("RGBA", (30, 30))56        d = ImageDraw.Draw(txt)57        d.text(58            (choice(range(0, 15)), -5 + choice(range(0, 15))),59            t,60            font=font,61            fill=(255, 0, 0),62        )63        image.paste(txt, (14 + (i * 20), 0), txt)64 65    # draw noise lines66    for i in range(30):67        fill = choice([120, 200])68        x = random() * WIDTH69        y = random() * HEIGHT70 71        draw.line(72            (73                x,74                y,75                x + 15 * (random() - 1),76                y + 15 * (random() - 1),77            ),78            fill=(fill, fill, fill),79            width=1,80        )81 82    return image83 84 85if __name__ == "__main__":86    for i in range(count):87        image, ans = gen()88        image.save(os.path.join(genereated_dir, f"{i}.png"))89        with open(os.path.join(genereated_dir, f"{i}.txt"), "w") as f:90            f.write(ans)91            print(f"Generated {i}.png and {i}.txt")92 93    print("Done")94