captcha_solver/utils/image.py

"""Image / audio decoding helpers."""

from __future__ import annotations

import base64
import io
from typing import Optional

import numpy as np
from PIL import Image


def decode_base64_image(b64: str) -> bytes:
    """Decode a base64 string to raw image bytes.

    Strips data URL prefix if present.
    Raises ValueError on bad input.
    """
    if not b64:
        raise ValueError("empty base64 image")
    if "," in b64 and b64.startswith("data:"):
        b64 = b64.split(",", 1)[1]
    try:
        return base64.b64decode(b64, validate=True)
    except Exception as exc:
        raise ValueError(f"invalid base64: {exc}") from exc


def decode_base64_audio(b64: str) -> bytes:
    """Decode a base64 string to raw audio bytes."""
    if not b64:
        raise ValueError("empty base64 audio")
    if "," in b64 and b64.startswith("data:"):
        b64 = b64.split(",", 1)[1]
    try:
        return base64.b64decode(b64, validate=True)
    except Exception as exc:
        raise ValueError(f"invalid base64: {exc}") from exc


def image_to_pil(data: bytes) -> Image.Image:
    """Decode image bytes to a PIL Image. Raises ValueError on failure."""
    try:
        img = Image.open(io.BytesIO(data))
        img.load()
        if img.mode not in ("RGB", "RGBA", "L"):
            img = img.convert("RGB")
        return img
    except Exception as exc:
        raise ValueError(f"could not decode image: {exc}") from exc


def perceptual_hash(data: bytes, size: int = 16) -> str:
    """Compute a perceptual hash of an image for cache lookup.

    Returns a 64-char hex string. Robust to small JPEG re-encodes
    and minor color shifts.
    """
    img = image_to_pil(data).convert("L").resize((size + 1, size))
    pixels = np.asarray(img, dtype=np.float32)
    diff = pixels[:, 1:] > pixels[:, :-1]
    bits = diff.flatten()
    h = 0
    for bit in bits:
        h = (h << 1) | int(bit)
    return format(h, f"0{size*size // 4}x")


def audio_size_hash(data: bytes) -> str:
    """Cheap audio hash: size + first/last 1KB + sha-like prefix."""
    import hashlib
    h = hashlib.sha256()
    h.update(len(data).to_bytes(8, "big"))
    h.update(data[:1024])
    h.update(data[-1024:])
    return h.hexdigest()[:32]