ballagb19's picture
Upload captcha_solver/utils/tesseract_ocr.py with huggingface_hub
8f0eaf7 verified
Raw
History Blame Contribute Delete
3.43 kB
"""Tesseract OCR helper.
Pure-Python wrapper that uses pytesseract. Tesseract binary must be
installed separately (Windows installer from UB-Mannheim). The
captcha solver falls back to this if Florence-2 is unavailable.
Install: pip install pytesseract
Windows: https://github.com/UB-Mannheim/tesseract/wiki (use the .exe installer)
Linux: apt install tesseract-ocr
macOS: brew install tesseract
"""
from __future__ import annotations
import os
import shutil
from typing import Optional
import pytesseract
from PIL import Image
def is_available() -> bool:
"""Return True if tesseract binary is on PATH, in TESSERACT_CMD env, or in standard locations."""
if shutil.which("tesseract") or shutil.which("tesseract.exe"):
return True
if os.environ.get("TESSERACT_CMD") and os.path.exists(os.environ["TESSERACT_CMD"]):
return True
for p in (
r"C:\Program Files\Tesseract-OCR\tesseract.exe",
r"C:\Program Files (x86)\Tesseract-OCR\tesseract.exe",
"/usr/bin/tesseract",
"/usr/local/bin/tesseract",
):
if os.path.exists(p):
pytesseract.pytesseract.tesseract_cmd = p
return True
return False
def _preprocess_for_ocr(pil_image: Image.Image) -> Image.Image:
"""Upscale + add padding to help tesseract read cramped captcha images.
Tesseract needs significant whitespace between characters and at borders.
Small captchas (220x70) often bleed characters together, producing OCR errors
like "3+5=?" -> "34+5=?". Scaling 2-3x and adding 20px padding fixes most of these.
"""
w, h = pil_image.size
# Upscale if smaller than 400px wide
if w < 400:
scale = 3 if w < 300 else 2
new_size = (w * scale, h * scale)
pil_image = pil_image.resize(new_size, Image.LANCZOS)
# Add 20px white border all around
if pil_image.mode != "RGB":
pil_image = pil_image.convert("RGB")
bordered = Image.new("RGB", (pil_image.size[0] + 40, pil_image.size[1] + 40), (255, 255, 255))
bordered.paste(pil_image, (20, 20))
return bordered
def ocr_image(
pil_image: Image.Image,
psm: int = 7,
lang: str = "eng",
whitelist: Optional[str] = None,
) -> str:
"""Run Tesseract on a PIL image.
psm=7: single line
psm=8: single word
psm=6: assume uniform block
"""
config_parts = [f"--psm {psm}", f"-l {lang}"]
if whitelist:
config_parts.append(f"-c tessedit_char_whitelist={whitelist}")
return pytesseract.image_to_string(pil_image, config=" ".join(config_parts)).strip()
def ocr_captcha_math(pil_image: Image.Image) -> str:
"""OCR a math captcha. Upscales, then tries psm modes in priority order.
Returns the first non-empty result. PSM 7 (single line) is the most reliable
for math captchas after upscaling; falls back to others if empty.
"""
pil_image = _preprocess_for_ocr(pil_image)
for psm in (7, 6, 13, 8, 11):
try:
txt = ocr_image(pil_image, psm=psm, whitelist="0123456789+-*/=()xX? ").strip()
if txt:
return txt
except Exception:
continue
return ""
def ocr_captcha_text(pil_image: Image.Image) -> str:
"""OCR a text-only captcha (alphanumeric)."""
pil_image = _preprocess_for_ocr(pil_image)
return ocr_image(pil_image, psm=7, whitelist="ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789")