1 files changed, 130 insertions, 0 deletions
diff --git a/captcha.py b/captcha.py
new file mode 100644
index 0000000..4b41015
--- /dev/null
+++ b/captcha.py
@@ -0,0 +1,130 @@
+import os
+import random
+from wand.image import Image
+from wand.drawing import Drawing
+from wand.color import Color
+
+def load_dataset(captcha_dir: str) -> list:
+    dataset = []
+    for entry in os.scandir(captcha_dir):
+        if not entry.is_dir():
+            continue
+        count = sum(
+            1 for f in os.scandir(entry.path)
+            if f.is_file() and f.name.endswith(".png")
+        )
+        if count > 0:
+            dataset.append([entry.name, count])
+
+    if len(dataset) < 2:
+        raise ValueError(
+            f"captcha_dir '{captcha_dir}' must contain at least 2 category "
+            f"sub-folders with .png images; found {len(dataset)}."
+        )
+    return dataset
+
+
+def generate_captcha_image(
+    captcha_dir: str,
+    font_path: str = "./data/fonts/captcha.ttf",
+    minimum_ans: int = 3,
+    maximum_ans: int = 5
+) -> tuple[list[int], bytes]:
+    captcha_dataset = load_dataset(captcha_dir)
+    pool = list(range(1, 16)) #[1, 2, ..., 15]
+    picks = random.randint(minimum_ans, maximum_ans)
+    answer_pos = []
+
+    for i in range(picks): # pick random answer slots first
+        idx = random.randint(0, len(pool)-1)
+        answer_pos.append(pool.pop(idx))
+
+    c = len(captcha_dataset)
+    chosen = captcha_dataset[random.randint(0, c-1)] # choose a dataset
+
+    choices = [ds for ds in captcha_dataset if ds[0] != chosen[0]] # wrong answers
+    grid = []
+    for i in range(16):
+        if i in answer_pos:
+            grid.append(chosen)
+        else:
+            grid.append(choices[random.randint(0, len(choices) - 1)])
+
+    # top banner bg
+    if random.randint(0, 1) == 0:
+        theme = {"bg": "#ebdbb2", "fg": "#1d2021"}
+    else:
+        theme = {"bg": "#1d2021", "fg": "#ebdbb2"}
+
+    noise_types = ["gaussian", "laplacian"]
+    distort_types = ["affine", "shepards"]
+
+    with Image(width=400, height=427, background=Color(theme["bg"])) as im:
+        im.format = "jpeg"
+        im.background_color = Color(theme["bg"])
+
+        for idx in range(16):
+            row = idx // 4
+            col = idx % 4
+
+            img_path = os.path.join(
+                captcha_dir,
+                grid[idx][0],
+                f"{random.randint(1, grid[idx][1])}.png",
+            )
+
+            with Image(filename=img_path) as tile:
+                tile.background_color = Color("black")
+                tile.alpha_channel = "remove"
+                if random.randint(0, 1) == 1:
+                    tile.flop()
+
+                tile.distort(
+                    distort_types[random.randint(0, 1)],
+                    [
+                        0,   0,   random.randint(-15, 15), random.randint(-15, 15),
+                        100, 0,   random.randint(80, 120), random.randint(-15, 15),
+                        100, 100, random.randint(80, 120), random.randint(80, 120),
+                        0,   100, random.randint(-15, 15), random.randint(80, 120),
+                    ],
+                    best_fit=False,
+                )
+                tile.noise(noise_types[random.randint(0, 1)])
+                im.composite(tile, left=col * 100, top=(row * 100) + 27)
+
+        label = f"Pick {picks} images of {chosen[0].replace('_', ' ')}"
+
+        with Drawing() as draw:
+            draw.font = font_path
+            draw.font_size = 24
+            draw.fill_color = Color(theme["fg"])
+            full_metrics = draw.get_font_metrics(im, label)
+            pos_x = 200 - (full_metrics.text_width / 2)
+
+            for char in label:
+                char_metrics = draw.get_font_metrics(im, char)
+                draw.push()
+                draw.translate(pos_x, 22)
+                draw.rotate(random.randint(-15, 15))
+                draw.text(0, 0, char)
+                draw.pop()
+
+                pos_x += char_metrics.text_width
+            draw(im)
+        im.compression_quality = 90
+        return answer_pos, im.make_blob("jpeg")
+
+
+
+
+if __name__ == "__main__":
+    import sys
+    positions, blob = generate_captcha_image(
+        captcha_dir="../captcha",
+        font_path="../data/fonts/captcha.ttf",
+    )
+    out = "captcha_out.jpg"
+    with open(out, "wb") as fh:
+        fh.write(blob)
+    print(f"Answer positions : {positions}", file=sys.stderr)
+    print(f"Image saved to   : {out}", file=sys.stderr)