Generate a captcha image

captchaimage.py

from __future__ import annotations

import platform
from collections.abc import Mapping, Sequence
from os import PathLike, environ, path
from pathlib import Path
from random import choice, randint
from typing import TYPE_CHECKING, Any, BinaryIO

import numpy as np
from PIL import Image, ImageDraw, ImageFilter, ImageFont

if TYPE_CHECKING:
    from _typeshed import StrOrBytesPath

__all__ = ["captcha_image"]


def random_color(s=1, e=255):
    """Random color default color range [1255]"""
    return randint(s, e), randint(s, e), randint(s, e)


def captcha_image(
    s: str,
    *,
    blur_radius: float = 1.0,
    font: StrOrBytesPath | BinaryIO | Sequence[StrOrBytesPath | BinaryIO] | None = None,
    font_size: int = 64,
    padding: int = 4,
    rotate: float = 45.0,
    rotate_expand: bool = False,
    save: str | Path | BinaryIO | None = None,
    save_format: str | None = None,
    save_params: Mapping[str, Any] | None = None,
) -> Image.Image:
    """Generate a captcha image

    Parameters
    ----------
    s (str): String of the captcha
        blur_radius: Standard deviation of the Gaussian kernel
    font: A singe or list of filename or file-like object containing a TrueType font.
        If the file is not found in this filename,
        the loader may also search in other directories,
        such as the :file:`fonts/` directory on Windows,
        or :file:`/Library/Fonts/`, :file:`/System/Library/Fonts/` and :file:`~/Library/Fonts/` on macOS.
        :see: Pillow's :func:`ImageFont.truetype`
    font_size: The requested font size, in pixels.
        :see: Pillow's :func:`ImageFont.truetype`
    padding: The padding of each character, in pixels.
    rotate: The random degree range to rotate the character. ``0`` means no rotate.
    rotate_expand: If true, expands the output image to make it large enough to hold the entire rotated image.
           If false or omitted, make the output image the same size as the input image.
           Note that the expand flag assumes rotation around the center and no translation.
    save: A filename (string), pathlib.Path object or file object.
        :see: Pillow's :meth:`Image.Image.save`
    save_format: Optional format override.
        If omitted, the format to use is determined from the filename extension.
        If a file object was used instead of a filename, this
        parameter should always be used.
        :see: Pillow's :meth:`Image.Image.save`
    save_params: Extra parameters to the image writer.
        :see: Pillow's :meth:`Image.Image.save`

    Returns
    -------
    Generated Pillow image object
    """
    # Create a font object
    if font is None:
        system = platform.system()
        if system == "Windows":
            font_dir = path.join(environ["WINDIR"], "Fonts")
            default_font_path = path.join(font_dir, "arial.ttf")
        elif system == "Linux":
            font_dir = "/usr/share/fonts/truetype"
            default_font_path = path.join(font_dir, "dejavu", "DejaVuSans.ttf")
        elif system == "Darwin":  # macOS
            font_dir = "/Library/Fonts"
            default_font_path = path.join(font_dir, "Arial.ttf")
        else:
            raise NotImplementedError(f"Unsupported OS: {system}")
        if path.exists(default_font_path):
            font = default_font_path
        else:
            raise FileNotFoundError(
                f"Default font not found at expected path: {default_font_path}"
            )
    if isinstance(font, Sequence) and not isinstance(
        font, (str, bytes, PathLike, BinaryIO)
    ):
        fonts = [ImageFont.truetype(x, font_size) for x in font]  # pyright: ignore[reportArgumentType]
    else:
        fonts = [ImageFont.truetype(font, font_size)]
    char_im_list = list()
    # 替换 random.randint/randint/uniform 为 numpy 的随机函数
    for c in s:
        # draw the char, with random rotate
        fnt = choice(fonts)
        bbox = fnt.getbbox(c)
        text_width = bbox[2] - bbox[0] + font_size / 4
        text_height = bbox[3] - bbox[1] + font_size / 4
        im = Image.new("RGBA", (int(text_width), int(text_height)))
        draw = ImageDraw.Draw(im)
        draw.text(
            (0, 0),
            c,
            fill=(*random_color(32, 127), np.random.randint(128, 255)),
            font=fnt,
            align=choice(["left", "center", "right"]),
        )
        if rotate:
            im = im.rotate(np.random.uniform(-rotate, rotate), expand=rotate_expand)
        char_im_list.append(im)

    # Size of the main image object
    captcha_width = sum(im.width + padding * 2 for im in char_im_list)
    captcha_height = max(im.height for im in char_im_list) + padding * 2

    # Create the main image object
    # with noisy dots in NdArray
    rgb_array = np.random.randint(
        64, 255, size=(captcha_height, captcha_width, 3), dtype=np.uint8
    )
    alpha_array = np.full((captcha_height, captcha_width), 255, dtype=np.uint8)
    captcha_im = Image.fromarray(np.dstack((rgb_array, alpha_array)))

    # Paste chars on image
    w = 0
    for im in char_im_list:
        pos = (
            w + padding + np.random.randint(-padding, padding + 1),
            np.random.randint(0, padding * 2 + 1),
        )
        w += im.width + 2 * padding
        captcha_im.paste(im, pos, im.split()[-1])

    # Noise: random lines using numpy
    line_layer = Image.new("RGBA", captcha_im.size)
    num_lines = randint(len(s), len(s) * 4)
    line_coords = np.column_stack(
        [
            np.random.randint(0, captcha_im.width, size=num_lines * 2),
            np.random.randint(0, captcha_im.height, size=num_lines * 2),
        ]
    ).reshape(num_lines, 2, 2)  # (num_lines, 2 points per line, x/y)
    line_colors = np.random.randint(32, 127, size=(num_lines, 3), dtype=np.uint8)
    line_alphas = np.random.randint(96, 191, size=(num_lines, 1), dtype=np.uint8)
    line_widths = np.random.randint(1, 8, size=(num_lines,))

    # Convert to PIL ImageDraw commands
    draw_layer = ImageDraw.Draw(line_layer)
    for i in range(num_lines):
        coords = tuple(map(int, line_coords[i].flatten()))
        color = tuple(line_colors[i]) + tuple(line_alphas[i])
        width = int(line_widths[i])
        draw_layer.line(coords, fill=color, width=width)

    # composite alpha channel
    captcha_im = Image.alpha_composite(captcha_im, line_layer)

    # Noise: Fuzzy filter
    captcha_im = captcha_im.filter(ImageFilter.GaussianBlur(blur_radius))

    # Remove Alpha
    captcha_im = captcha_im.convert("RGB")

    # save
    if save:
        captcha_im.save(save, save_format, **(save_params or dict()))

    # return Pillow image object
    return captcha_im