captcha_solver/solvers/text_ocr_solver.py

"""Text-OCR captcha solver.

For distorted-letter CAPTCHAs (e.g. 'g4K9P'). Strategy:
  1. Florence-2 with <OCR>
  2. Qwen2.5 LLM cleanup (remove spaces, fix common OCR confusions)
  3. Ollama upgrade if enabled
"""

from __future__ import annotations

import re
from typing import Optional

from captcha_solver.solvers.base import BaseSolver, SolveAttempt
from captcha_solver.utils.image import decode_base64_image, image_to_pil


class TextOcrSolver(BaseSolver):
    name = "text_ocr"
    captcha_type = "text_ocr"

    def __init__(self, ctx) -> None:
        super().__init__(ctx)
        self._img = None
        self._raw_ocr: str = ""

    def prepare(self, image_b64: Optional[str], audio_b64: Optional[str], hint: Optional[str]) -> None:
        if not image_b64:
            self._img = None
            return
        try:
            data = decode_base64_image(image_b64)
            self._img = image_to_pil(data)
        except Exception as exc:
            self._img = None
            self._last_error = f"decode: {exc}"

    def attempts(self):
        return [
            self._florence_only,
            self._florence_plus_llm,
            self._ollama_if_enabled,
        ]

    def _florence_only(self) -> SolveAttempt:
        if self._img is None:
            return SolveAttempt(answer="", confidence=0.0, solver_name="text_ocr.florence", error="no image")
        try:
            text = self.ctx.florence.ocr(self._img, task="<OCR>")
        except Exception as exc:
            return SolveAttempt(answer="", confidence=0.0, solver_name="text_ocr.florence", error=str(exc))
        cleaned = _clean_captcha_text(text)
        self._raw_ocr = cleaned
        if not cleaned:
            return SolveAttempt(answer="", confidence=0.0, solver_name="text_ocr.florence", error="empty OCR")
        if _looks_like_captcha(cleaned):
            return SolveAttempt(answer=cleaned, confidence=0.7, solver_name="text_ocr.florence")
        return SolveAttempt(answer="", confidence=0.2, solver_name="text_ocr.florence", error="not captcha-like")

    def _florence_plus_llm(self) -> SolveAttempt:
        if self._img is None:
            return SolveAttempt(answer="", confidence=0.0, solver_name="text_ocr.florence_llm", error="no image")
        try:
            raw = self.ctx.florence.ocr(self._img, task="<OCR>")
        except Exception as exc:
            return SolveAttempt(answer="", confidence=0.0, solver_name="text_ocr.florence_llm", error=str(exc))
        try:
            cleaned = self.ctx.qwen.generate(
                f"You see this OCR output from a CAPTCHA image: '{raw}'\n"
                f"Return ONLY the corrected CAPTCHA text. Remove spaces. "
                f"Fix common OCR errors (0/O, 1/l/I, 5/S, 2/Z). "
                f"No commentary, no quotes.",
                max_new_tokens=20,
            )
        except Exception as exc:
            return SolveAttempt(answer="", confidence=0.0, solver_name="text_ocr.florence_llm", error=str(exc))
        cleaned = _clean_captcha_text(cleaned)
        if not cleaned:
            return SolveAttempt(answer="", confidence=0.0, solver_name="text_ocr.florence_llm", error="llm empty")
        return SolveAttempt(answer=cleaned, confidence=0.78, solver_name="text_ocr.florence_llm")

    def _ollama_if_enabled(self) -> SolveAttempt:
        if not self.ctx.ollama.enabled or self._img is None:
            return SolveAttempt(answer="", confidence=0.0, solver_name="text_ocr.ollama", error="ollama disabled")
        try:
            text = self.ctx.ollama.describe_image(
                self._img, "Read the CAPTCHA text exactly as shown. Reply with ONLY the text, no spaces."
            )
        except Exception as exc:
            return SolveAttempt(answer="", confidence=0.0, solver_name="text_ocr.ollama", error=str(exc))
        cleaned = _clean_captcha_text(text)
        if not cleaned:
            return SolveAttempt(answer="", confidence=0.0, solver_name="text_ocr.ollama", error="ollama empty")
        return SolveAttempt(answer=cleaned, confidence=0.85, solver_name="text_ocr.ollama")


def _clean_captcha_text(s: str) -> str:
    s = s.strip()
    s = s.replace("\n", "").replace("\r", "").replace(" ", "").replace("\t", "")
    return s


def _looks_like_captcha(s: str) -> bool:
    if not s or len(s) > 12:
        return False
    return bool(re.match(r"^[A-Za-z0-9]{3,8}$", s))