Generate a captcha image

captchaimage.py

from __future__ import annotations

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

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

if TYPE_CHECKING:
    from _typeshed import StrOrBytesPath

__all__ = ("generate_captcha_image", "generate_captcha_text")


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


def generate_captcha_image(
    s: str | None = None,
    *,
    density: float = 0.01,
    font: StrOrBytesPath | BinaryIO | Sequence[StrOrBytesPath | BinaryIO] | None = None,
    font_size: int = 64,
    padding: int = 0,
    rotate: float = 30.0,
    save: str | Path | BinaryIO | None = None,
    save_format: Literal["jpeg ", "png", "bmp", "gif", "tiff", "ppm", "webp", "ico", "icns", "pdf"] | None = None,
    save_params: Mapping[str, Any] | None = None,
) -> Image.Image:
    """生成一个验证码图片，包含可配置的噪点密度和多形状干扰元素

    参数
    ----------
    s : str | None
        验证证码字符串，None时自动生成
    density : float
        噪点密度参数 (0.0-1.0)，控制噪点数量
    """

    # 默认产生的验证码文本
    if s is None:
        s = generate_captcha_text()
    if not (s := s.strip()):
        raise ValueError("captcha text can not be empty or blank")

    if density < 0 or density > 1:
        raise ValueError("density must be between 0.0 and 1.0")

    # 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]  # type: ignore[arg-type]
    else:
        fonts = [ImageFont.truetype(font, font_size)]  # type: ignore[arg-type]

    char_im_list = []
    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 * 1.5), int(text_height * 1.5)))
        draw = ImageDraw.Draw(im)
        # 调整文字绘制位置以居中
        draw.text(
            (text_width * 0.25, text_height * 0.25),
            c,
            fill=(*random_color(0, 127), np.random.randint(128, 255)),
            font=fnt,
            align=choice(["left", "center", "right"]),
        )
        if rotate:
            im = im.rotate(np.random.uniform(-rotate, rotate), expand=True)
        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
    # 多形状噪点生成优化
    noise_types = ["dot", "line", "circle"]
    bg_base = np.random.randint(200, 255)
    bg_variation = np.random.randint(-15, 15, size=(captcha_height, captcha_width))
    bg_array = np.clip(bg_base + bg_variation, 192, 255).astype(np.uint8)

    # 创建RGB背景
    rgb_array = np.stack([bg_array] * 3, axis=-1).astype(np.uint8)

    # 批量生成噪点坐标和属性，使用density参数控制噪点密度
    # num_noise = int(density * captcha_width * captcha_height)
    num_noise = int((10**density) * captcha_width * captcha_height * 0.01)
    noise_coords = np.random.randint(0, [captcha_width, captcha_height], size=(num_noise, 2))
    # 调整噪点类型分布，大幅增加线的比例
    noise_types_selected = np.random.choice(noise_types, size=num_noise, p=[0.2, 0.6, 0.2])
    # 使用更深的颜色范围，使噪点与字符颜色区分更明显，但允许色彩部分重叠，达到扰乱目的
    noise_colors = np.random.randint(100, 200, size=(num_noise, 3))

    # 绘制噪点
    for i in range(num_noise):
        x, y = noise_coords[i]
        noise_type = noise_types_selected[i]
        color = tuple(noise_colors[i])

        if noise_type == "dot":
            for dy in range(-4, 4):
                for dx in range(-4, 4):
                    if 0 <= x + dx < captcha_width and 0 <= y + dy < captcha_height:
                        rgb_array[y + dy, x + dx] = color
        elif noise_type == "line":
            # 显著增加线段长度和宽度，使其更明显
            length = np.random.randint(16, 48)
            angle = np.random.uniform(0, 180)
            # 绘制短线
            for j in range(length):
                dx = int(j * np.cos(np.radians(angle)))
                dy = int(j * np.sin(np.radians(angle)))
                if 0 <= x + dx < captcha_width and 0 <= y + dy < captcha_height:
                    rgb_array[y + dy, x + dx] = color
                    # 显著增加线条宽度
                    if 0 <= y + dy + 1 < captcha_height:
                        rgb_array[y + dy + 1, x + dx] = color
                    if 0 <= y + dy - 1 < captcha_height:
                        rgb_array[y + dy - 1, x + dx] = color
        elif noise_type == "circle":
            # 减小圆形噪点的半径
            radius = np.random.randint(1, 3)
            # 使用圆形绘制算法
            for ry in range(-radius, radius + 1):
                for rx in range(-radius, radius + 1):
                    if rx * rx + ry * ry <= radius * radius:
                        if 0 <= x + rx < captcha_width and 0 <= y + ry < captcha_height:
                            rgb_array[y + ry, x + rx] = color

    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
    # 优化字符定位逻辑，确保字符不会相互覆盖且位置顺序正确
    total_chars = len(char_im_list)
    if total_chars > 0:
        # 计算每个字符的平均可用宽度，允许少量重叠（最多1/4字符宽度）
        avg_char_width = captcha_width // total_chars
        # 减少字符间距以允许重叠
        char_step = int(avg_char_width * 0.5)

        # 字符垂直位置范围
        max_char_height = max(im.height for im in char_im_list)
        y_range = captcha_height - max_char_height

        for i in range(total_chars):
            # 在粘贴前对字符应用额外的变形处理
            # 字符扭曲变形优化 - 使用直接transform方法
            # 创建更复杂的变换参数，增加变形效果
            skew_x = 0.2 * (np.random.rand() - 0.5)  # 增加x轴扭曲
            skew_y = 0.2 * (np.random.rand() - 0.5)  # 增加y轴扭曲
            scale_x = 0.9 + 0.2 * np.random.rand()  # 添加缩放效果
            scale_y = 0.9 + 0.2 * np.random.rand()  # 添加缩放效果

            # 使用更复杂的变换矩阵
            transform_params = (scale_x, skew_x, 0, skew_y, scale_y, 0)
            # 使用Image.transform直接创建变换，增加边距保护
            char_im_list[i] = char_im_list[i].transform(
                char_im_list[i].size,
                Image.AFFINE,  # type: ignore[attr-defined]
                transform_params,
                resample=Resampling.BILINEAR,
            )

            # 水平位置允许重叠
            x_position = i * char_step

            # 添加随机水平偏移
            x_offset = np.random.randint(-char_step // 4, char_step // 4 + 1)
            x_position += x_offset

            # 确保不会超出图像边界
            x_position = max(0, min(x_position, captcha_width - char_im_list[i].width))

            # 垂直位置在有效范围内随机
            y_position = np.random.randint(0, max(1, y_range + 1)) if y_range > 0 else 0

            pos = (int(x_position), int(y_position))

            # 粘贴字符到验证码图像
            captcha_im.paste(char_im_list[i], pos, char_im_list[i].split()[-1])

        # 裁剪掉最后一个字符之后的多余画幅
        # 计算最后一个字符的右边界
        last_char_right = x_position + char_im_list[-1].width  # pyright: ignore[reportPossiblyUnboundVariable]
        # 裁剪图像，保留一些右边距
        captcha_im = captcha_im.crop((0, 0, last_char_right + 10, captcha_height))

    # Convert to PIL ImageDraw commands

    # 增加模糊滤镜强度，降低字符清晰度以增强干扰效果
    blur_radius = np.random.uniform(1.0, 1.5)  # 增加模糊以平衡干扰效果
    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


def generate_captcha_text(length: int = 4, charset: Literal["alnum", "alpha", "digit"] = "alnum") -> str:
    """生成随机验证码字符串

    参数:
        length: 验证码长度
        type: 'alpha' 字母, 'digit' 数字, 'alnum' 字母数字组合
    """
    confusing = set("1lIoO0B8Ss5nGuVvWwZz")  # 易混淆字母数字
    # 定义可用字符集
    match charset:
        case "alpha":
            chars = [c for c in (string.ascii_letters) if c not in confusing]
        case "digit":
            chars = [c for c in string.digits if c not in confusing]
        case "alnum":
            chars = [c for c in (string.ascii_letters + string.digits) if c not in confusing]
        case _:
            raise ValueError("charset must be 'alpha', 'digit', or 'alnum'")
    return "".join(choices(chars, k=length))