Generate a captcha image

captchaimage.py

from itertools import product
from pathlib import Path
from random import choice, randint, randrange, uniform
from typing import Any, BinaryIO, Mapping, Optional, Sequence, Union

from PIL import Image, ImageDraw, ImageFilter, ImageFont

__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: Union[None, str, bytes, Sequence[Union[str, bytes]]] = None,
    font_size: int = 32,
    padding: int = 4,
    rotate: float = 30.0,
    rotate_expand: bool = False,
    save: Union[str, Path, BinaryIO, None] = None,
    save_format: Optional[str] = None,
    save_params: Optional[Mapping[str, Any]] = 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`

    char_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
    """
    # verification code written on the picture
    # Create a font object
    if isinstance(font, Sequence) and not isinstance(font, (str, bytes)):
        fonts = [ImageFont.truetype(x, font_size) for x in font]
    else:
        fonts = [ImageFont.truetype(font, font_size)]
    char_im_list = list()
    for c in s:
        # draw the char, with random rotate
        fnt = choice(fonts)
        im = Image.new('RGBA', fnt.getsize(c))
        draw = ImageDraw.Draw(im)
        draw.text(
            (0, 0), c,
            fill=random_color(32, 127),
            font=fnt, align='center',
        )
        if rotate:
            im = im.rotate(uniform(-rotate, rotate), expand=rotate_expand)
        char_im_list.append(im)
    # Create 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
    captcha_im = Image.new('RGBA', (captcha_width, captcha_height))
    # Noise: random color fill each pixel for the background
    draw = ImageDraw.Draw(captcha_im)
    for x, y in product(range(captcha_width), range(captcha_height)):
        draw.point((x, y), fill=random_color(64, 255))
    # Paste chars on image
    w = 0
    for im in char_im_list:
        pos = (
            w + padding + randint(-padding, padding),
            randint(0, padding*2)
        )
        w += im.width+2*padding
        captcha_im.paste(im, pos, im.split()[-1])
    # Noise: random lines:
    for _ in range(max(len(s), randrange(len(s)*2))):
        draw.line(
            [
                (randrange(captcha_im.width), randrange(captcha_im.height))
                for _ in range(min(2, randrange(len(s)*2)))
            ],
            random_color(32, 127)
        )
    # 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