CaptchBreaker/generators/rotate_solver_gen.py

"""
旋转验证码求解器数据生成器

生成旋转验证码训练数据：随机图案 (色块/渐变/几何图形)，随机旋转 0-359°。
裁剪为圆形 (黑色背景填充圆外区域)。

标签 = 旋转角度 (整数)
文件名格式: {angle}_{index:06d}.png
"""

import math
import random

from PIL import Image, ImageDraw, ImageFilter, ImageFont

from config import SOLVER_CONFIG
from generators.base import BaseCaptchaGenerator

_FONT_PATHS = [
    "/usr/share/fonts/TTF/DejaVuSans-Bold.ttf",
    "/usr/share/fonts/TTF/DejaVuSerif-Bold.ttf",
    "/usr/share/fonts/liberation/LiberationSans-Bold.ttf",
    "/usr/share/fonts/liberation/LiberationSerif-Bold.ttf",
    "/usr/share/fonts/gnu-free/FreeSansBold.otf",
]


class RotateSolverDataGenerator(BaseCaptchaGenerator):
    """旋转验证码求解器数据生成器。"""

    def __init__(self, seed: int | None = None):
        from config import RANDOM_SEED
        super().__init__(seed=seed if seed is not None else RANDOM_SEED)

        self.cfg = SOLVER_CONFIG["rotate"]
        self.height, self.width = self.cfg["input_size"]  # (H, W)

        self._fonts: list[str] = []
        for p in _FONT_PATHS:
            try:
                ImageFont.truetype(p, 20)
                self._fonts.append(p)
            except OSError:
                continue

    def generate(self, text: str | None = None) -> tuple[Image.Image, str]:
        rng = self.rng

        # 随机旋转角度 0-359
        angle = rng.randint(0, 359)
        if text is None:
            text = str(angle)

        size = self.width  # 正方形
        radius = size // 2

        # 1. 生成正向图案 (未旋转)
        content = self._random_pattern(rng, size)

        # 2. 旋转图案
        rotated = content.rotate(-angle, resample=Image.BICUBIC, expand=False)

        # 3. 裁剪为圆形 (黑色背景)
        result = Image.new("RGB", (size, size), (0, 0, 0))
        mask = Image.new("L", (size, size), 0)
        mask_draw = ImageDraw.Draw(mask)
        mask_draw.ellipse([0, 0, size - 1, size - 1], fill=255)

        result.paste(rotated, (0, 0), mask)

        # 4. 轻微模糊
        result = result.filter(ImageFilter.GaussianBlur(radius=0.5))

        return result, text

    def _random_pattern(self, rng: random.Random, size: int) -> Image.Image:
        """生成随机图案 (带明显方向性，便于模型学习旋转)。"""
        img = Image.new("RGB", (size, size))
        draw = ImageDraw.Draw(img)

        # 渐变背景
        base_r = rng.randint(100, 220)
        base_g = rng.randint(100, 220)
        base_b = rng.randint(100, 220)
        for y in range(size):
            ratio = y / max(size - 1, 1)
            r = int(base_r * (1 - ratio) + rng.randint(40, 120) * ratio)
            g = int(base_g * (1 - ratio) + rng.randint(40, 120) * ratio)
            b = int(base_b * (1 - ratio) + rng.randint(40, 120) * ratio)
            draw.line([(0, y), (size, y)], fill=(r, g, b))

        cx, cy = size // 2, size // 2

        # 添加不对称几何图形 (让模型能感知方向)
        pattern_type = rng.choice(["triangle", "arrow", "text", "shapes"])

        if pattern_type == "triangle":
            # 顶部三角形标记
            color = tuple(rng.randint(180, 255) for _ in range(3))
            ts = size // 4
            draw.polygon(
                [(cx, cy - ts), (cx - ts // 2, cy), (cx + ts // 2, cy)],
                fill=color,
            )
            # 底部小圆
            draw.ellipse(
                [cx - 8, cy + ts // 2, cx + 8, cy + ts // 2 + 16],
                fill=tuple(rng.randint(50, 150) for _ in range(3)),
            )

        elif pattern_type == "arrow":
            # 向上的箭头
            color = tuple(rng.randint(180, 255) for _ in range(3))
            arrow_len = size // 3
            draw.line([(cx, cy - arrow_len), (cx, cy + arrow_len // 2)], fill=color, width=4)
            draw.polygon(
                [(cx, cy - arrow_len - 5), (cx - 10, cy - arrow_len + 10), (cx + 10, cy - arrow_len + 10)],
                fill=color,
            )

        elif pattern_type == "text" and self._fonts:
            # 文字 (有天然方向性)
            font_path = rng.choice(self._fonts)
            font_size = size // 3
            try:
                font = ImageFont.truetype(font_path, font_size)
                ch = rng.choice("ABCDEFGHJKLMNPRSTUVWXYZ23456789")
                bbox = font.getbbox(ch)
                tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
                draw.text(
                    (cx - tw // 2 - bbox[0], cy - th // 2 - bbox[1]),
                    ch,
                    fill=tuple(rng.randint(0, 80) for _ in range(3)),
                    font=font,
                )
            except OSError:
                pass

        else:
            # 混合不对称形状
            # 上方矩形
            w, h = rng.randint(15, 30), rng.randint(10, 20)
            color = tuple(rng.randint(150, 255) for _ in range(3))
            draw.rectangle([cx - w, cy - size // 3, cx + w, cy - size // 3 + h], fill=color)
            # 右下小圆
            r = rng.randint(5, 12)
            color2 = tuple(rng.randint(50, 150) for _ in range(3))
            draw.ellipse([cx + size // 5, cy + size // 5, cx + size // 5 + r * 2, cy + size // 5 + r * 2], fill=color2)

        # 添加纹理噪声
        for _ in range(rng.randint(20, 60)):
            nx, ny = rng.randint(0, size - 1), rng.randint(0, size - 1)
            nc = tuple(rng.randint(80, 220) for _ in range(3))
            draw.point((nx, ny), fill=nc)

        return img