captcha_solver/utils/detector.py

"""Heuristic captcha type detector.

This is intentionally simple: it looks at image dimensions, presence of
many equal-sized tiles, color distribution, etc. For audio it returns
"audio" if audio was provided. For images it picks the most likely
captcha type. The user-provided `type` field always wins.
"""

from __future__ import annotations

from typing import Optional

import numpy as np
from PIL import Image

from captcha_solver.utils.image import decode_base64_image, image_to_pil


def detect_type(image_b64: Optional[str], audio_b64: Optional[str]) -> str:
    """Best-guess captcha type from inputs. Returns one of:
    'math', 'text_ocr', 'image_grid', 'audio'.

    'math' is the default for any non-grid text-like image. When in doubt
    we still pick a solver rather than return None.
    """
    if audio_b64:
        return "audio"
    if not image_b64:
        return "math"

    try:
        data = decode_base64_image(image_b64)
        img = image_to_pil(data)
    except Exception:
        return "math"

    w, h = img.size
    aspect = w / max(h, 1)
    arr = np.asarray(img.convert("L"))
    h_std = float(arr.std())
    edges = _edge_density(arr)

    if 0.7 <= aspect <= 1.4 and h_std < 60 and edges < 0.08 and w < 400 and h < 200:
        return "math"
    if 0.4 <= aspect <= 0.7 and edges > 0.12 and _has_tile_grid(arr):
        return "image_grid"
    if h_std > 40 and (w >= 200 or h >= 60):
        return "text_ocr"
    return "math"


def _edge_density(arr: np.ndarray) -> float:
    """Fraction of pixels that are 'edges' (Sobel-lite)."""
    gx = np.abs(np.diff(arr.astype(np.int16), axis=1))
    gy = np.abs(np.diff(arr.astype(np.int16), axis=0))
    e = (gx[:-1, :] > 30).sum() + (gy[:, :-1] > 30).sum()
    return e / max(arr.size, 1)


def _has_tile_grid(arr: np.ndarray) -> bool:
    """Check for a 3x3 grid pattern (9 tiles) by looking for vertical/horizontal dark seams."""
    h, w = arr.shape
    if h < 90 or w < 90:
        return False
    rows = [h // 3, 2 * h // 3]
    cols = [w // 3, 2 * w // 3]
    seam_strengths = []
    for r in rows:
        seam_strengths.append(arr[r - 2 : r + 3, :].mean())
    for c in cols:
        seam_strengths.append(arr[:, c - 2 : c + 3].mean())
    return min(seam_strengths) < 100