atareao / ocr_ai.py

#!/usr/bin/env python3 """ ocr-ai.py — OCR inteligente con detección automática de contenido y post-procesado con IA (Ollama). Captura una región de la pantalla o procesa una imagen existente, detecta automáticamente el tipo de contenido (código, tabla, prosa), ajusta los parámetros de Tesseract, y limpia errores con Ollama. Uso: python3 ocr-ai.py # Captura + OCR + IA python3 ocr-ai.py -f captura.png # Procesa imagen existente python3 ocr-ai.py -f captura.png --raw # Solo OCR, sin IA python3 ocr-ai.py -f captura.png --copy # OCR + IA + al portapapeles Dependencias: - Python 3.10+ - ImageMagick (magick) - Tesseract (tesseract-ocr + tesseract-ocr-spa + tesseract-ocr-eng) - wl-clipboard (Wayland) o xclip (X11) para --copy - Ollama (ollama serve corriendo, modelo recomendado: llama3.2) """ import argparse import json import os import re import subprocess import sys import tempfile from pathlib import Path import requests # ─── Configuración ───────────────────────────────────────────────────────────── OLLAMA_HOST = "http://localhost:11434" OLLAMA_MODEL = "llama3.2" OLLAMA_TIMEOUT = 60 # Pipeline de ImageMagick por defecto DEFAULT_PIPELINE = [ "-colorspace", "Gray", "-normalize", "-threshold", "60%", "-deskew", "40%", "-sharpen", "0x1", ] # Pipeline más agresivo para código (más contraste, menos ruido) CODE_PIPELINE = [ "-colorspace", "Gray", "-normalize", "-contrast-stretch", "2%", "-threshold", "55%", "-deskew", "40%", "-sharpen", "0x2", "-morphology", "Close", "Octagon", ] # Pipeline para tablas (preserva líneas) TABLE_PIPELINE = [ "-colorspace", "Gray", "-normalize", "-threshold", "50%", "-deskew", "30%", "-sharpen", "0x1", "-morphology", "Erode", "Octagon", ] # ─── Detección de tipo de contenido ─────────────────────────────────────────── def detectar_tipo(texto: str) -> str: """ Analiza el texto extraído por OCR y detecta el tipo de contenido. Devuelve: 'codigo', 'tabla', 'prosa' o 'desconocido'. """ if not texto or len(texto.strip()) < 10: return "desconocido" lineas = texto.strip().split("\n") num_lineas = len(lineas) palabras = texto.split() num_palabras = len(palabras) # Heurísticas para código patrones_codigo = [ r'\b(import|from|def |class |fn |let |var |const|function|int |float|void|return|if |else |for |while|switch|case |package|#include|using namespace|public |private |static)\b', r'[{}()\[\];]', r'==|!=|<=|>=|->|=>|\+\+|--', r'^\s*[#//]\s', r'\.\w+\(', r'["\'].*["\']\s*[;,]', ] puntuacion_codigo = 0 for patron in patrones_codigo: if re.search(patron, texto, re.IGNORECASE | re.MULTILINE): puntuacion_codigo += 2 # Más peso si hay indentación consistente (código suele estar indentado) indentadas = sum(1 for l in lineas if l.startswith(" ") or l.startswith("\t")) if indentadas > num_lineas * 0.3: puntuacion_codigo += 3 # Puntuación por densidad de palabras (código tiene menos variedad léxica) if num_palabras > 0 and len(set(palabras)) / num_palabras < 0.4: puntuacion_codigo += 2 # Heurísticas para tabla (más restrictivas) patrones_tabla = [ r'\|.*\|.*\|', r'\+[-+]+\+', r'(Nº|#|Item|Producto|Precio|Cantidad|Total|Unidades)\s*\||\|.*\|', r'^\s*\d+\..*\d+[.,]\d+', ] puntuacion_tabla = 0 for patron in patrones_tabla: if re.search(patron, texto, re.IGNORECASE | re.MULTILINE): puntuacion_tabla += 2 # Tablas suelen tener muchas líneas cortas con estructura repetitiva lineas_cortas = sum(1 for l in lineas if 5 < len(l.strip()) < 40) if num_lineas > 3 and lineas_cortas > num_lineas * 0.5: puntuacion_tabla += 2 # Decisión if puntuacion_codigo >= puntuacion_tabla and puntuacion_codigo >= 4: return "codigo" elif puntuacion_tabla > puntuacion_codigo and puntuacion_tabla >= 3: return "tabla" elif num_palabras > 20 and puntuacion_codigo < 4 and puntuacion_tabla < 3: return "prosa" else: return "desconocido" def psm_para_tipo(tipo: str) -> int: """Devuelve el PSM de Tesseract óptimo para cada tipo.""" return { "codigo": 6, # Bloque de texto uniforme "tabla": 6, # Bloque uniforme (mejor para TSV) "prosa": 4, # Texto a una columna "desconocido": 3, # Detección automática }.get(tipo, 3) def pipeline_para_tipo(tipo: str) -> list: """Devuelve el pipeline de ImageMagick óptimo para cada tipo.""" return { "codigo": CODE_PIPELINE, "tabla": TABLE_PIPELINE, "prosa": DEFAULT_PIPELINE, "desconocido": DEFAULT_PIPELINE, }.get(tipo, DEFAULT_PIPELINE) # ─── Herramientas del sistema ───────────────────────────────────────────────── def check_deps(): """Verifica que las dependencias del sistema están instaladas.""" faltan = [] for cmd, nombre in [("convert", "ImageMagick"), ("tesseract", "Tesseract")]: if subprocess.run(["which", cmd], capture_output=True).returncode != 0: faltan.append(nombre) if faltan: print(f"❌ Faltan dependencias: {', '.join(faltan)}", file=sys.stderr) print(f" Instala: sudo apt install imagemagick tesseract-ocr tesseract-ocr-spa tesseract-ocr-eng", file=sys.stderr) sys.exit(1) def capturar_pantalla(output_path: str): """Captura una región de la pantalla con import (ImageMagick).""" print("🔍 Selecciona una región con el ratón...", file=sys.stderr) result = subprocess.run( ["import", output_path], capture_output=True, text=True, ) if result.returncode != 0: print("❌ Captura cancelada o fallida", file=sys.stderr) sys.exit(1) if not Path(output_path).exists() or Path(output_path).stat().st_size == 0: print("❌ No se generó la captura", file=sys.stderr) sys.exit(1) print(f" ✅ Captura guardada: {output_path}", file=sys.stderr) def preprocesar(input_path: str, output_path: str, pipeline: list): """Aplica el pipeline de ImageMagick a la imagen.""" cmd = ["convert", input_path] + pipeline + [output_path] result = subprocess.run(cmd, capture_output=True, text=True) if result.returncode != 0: print(f"❌ Error en preprocesado: {result.stderr}", file=sys.stderr) sys.exit(1) def ocr(image_path: str, idiomas: str, psm: int) -> str: """Ejecuta Tesseract OCR y devuelve el texto extraído.""" cmd = [ "tesseract", image_path, "stdout", "-l", idiomas, "--psm", str(psm), "--oem", "1", ] result = subprocess.run(cmd, capture_output=True, text=True) if result.returncode != 0: print(f"❌ Error en OCR: {result.stderr}", file=sys.stderr) return "" return result.stdout.strip() def copiar_portapapeles(texto: str): """Copia texto al portapapeles (Wayland o X11).""" if subprocess.run(["which", "wl-copy"], capture_output=True).returncode == 0: subprocess.run(["wl-copy"], input=texto, text=True) elif subprocess.run(["which", "xclip"], capture_output=True).returncode == 0: subprocess.run(["xclip", "-selection", "clipboard"], input=texto, text=True) else: print("⚠️ Ni wl-copy ni xclip disponibles. Texto solo por stdout.", file=sys.stderr) return print(" 📋 Copiado al portapapeles", file=sys.stderr) # ─── Post-procesado con Ollama ──────────────────────────────────────────────── def llamar_ollama(system_prompt: str, user_prompt: str, model: str = OLLAMA_MODEL) -> str | None: """Llama a Ollama y devuelve el texto generado.""" payload = { "model": model, "messages": [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}, ], "stream": False, "options": {"temperature": 0.1, "num_predict": 2048}, } try: resp = requests.post( f"{OLLAMA_HOST}/api/chat", json=payload, timeout=OLLAMA_TIMEOUT, ) resp.raise_for_status() return resp.json().get("message", {}).get("content", "").strip() except requests.exceptions.ConnectionError: print(" ⚠️ Ollama no está disponible. Omitiendo corrección con IA.", file=sys.stderr) return None except Exception as e: print(f" ⚠️ Error en Ollama: {e}. Omitiendo corrección.", file=sys.stderr) return None def corregir_con_ia(texto: str, tipo: str, model: str = OLLAMA_MODEL) -> str: """Corrige errores de OCR usando Ollama, adaptado al tipo de contenido.""" prompts = { "codigo": { "system": ( "Eres un corrector de OCR especializado en código fuente. " "Tu tarea es corregir errores de reconocimiento óptico en código.\n\n" "REGLAS:\n" "1. Corrige caracteres mal reconocidos (0 por O, l por 1, etc.)\n" "2. Restaura indentación y formato del código original\n" "3. NO cambies nombres de variables, funciones o lógica\n" "4. NO inventes líneas de código que no están\n" "5. Mantén comentarios exactamente como están\n" "6. Preserva la sintaxis del lenguaje\n" "7. Devuelve SOLO el código corregido, sin explicaciones" ), "user": f"Corrige los errores de OCR en este código:\n\n```\n{texto}\n```", }, "tabla": { "system": ( "Eres un corrector de OCR especializado en tablas y datos numéricos. " "Tu tarea es corregir errores de reconocimiento óptico en tablas.\n\n" "REGLAS:\n" "1. Corrige números mal reconocidos (1.000 por 1000, etc.)\n" "2. Corrige símbolos monetarios (€ por EUR o similar)\n" "3. Preserva la estructura de filas y columnas\n" "4. NO inventes datos que no están en el original\n" "5. NO rellenes celdas vacías con datos inventados\n" "6. Si no estás seguro de un valor, déjalo como está\n" "7. Devuelve SOLO la tabla corregida, sin explicaciones" ), "user": f"Corrige los errores de OCR en esta tabla:\n\n```\n{texto}\n```", }, "prosa": { "system": ( "Eres un corrector de OCR especializado en texto en español. " "Tu tarea es corregir errores de reconocimiento óptico.\n\n" "REGLAS:\n" "1. Corrige caracteres mal reconocidos\n" "2. Corrige espacios perdidos o añadidos entre palabras\n" "3. Corrige puntuación y acentos cuando sea evidente\n" "4. NO cambies el significado del texto\n" "5. NO añadas información que no está\n" "6. Preserva el estilo y tono original\n" "7. Si una palabra no la reconoces bien, déjala como está\n" "8. Devuelve SOLO el texto corregido, sin explicaciones" ), "user": f"Corrige los errores de OCR en este texto:\n\n```\n{texto}\n```", }, "desconocido": { "system": ( "Eres un corrector de OCR. Tu tarea es corregir errores " "de reconocimiento óptico en el texto que se te pasa.\n\n" "REGLAS:\n" "1. Corrige solo errores evidentes de OCR\n" "2. NO inventes contenido\n" "3. NO cambies el formato o estructura\n" "4. Devuelve SOLO el texto corregido, sin explicaciones" ), "user": f"Corrige los errores de OCR en este texto:\n\n```\n{texto}\n```", }, } p = prompts.get(tipo, prompts["desconocido"]) resultado = llamar_ollama(p["system"], p["user"], model) return resultado if resultado else texto # ─── Main ───────────────────────────────────────────────────────────────────── def main(): parser = argparse.ArgumentParser( description="ocr-ai: OCR inteligente con detección de contenido y corrección con IA", formatter_class=argparse.RawDescriptionHelpFormatter, epilog="""\ Ejemplos: python3 ocr-ai.py # Captura + OCR + IA python3 ocr-ai.py -f captura.png # Procesa imagen existente python3 ocr-ai.py -f captura.png --raw # Solo OCR, sin IA python3 ocr-ai.py -f captura.png --copy --lang eng # Código al portapapeles """, ) parser.add_argument("-f", "--file", help="Archivo de imagen a procesar (si no, captura interactiva)") parser.add_argument("--lang", default="spa+eng", help="Idiomas para Tesseract (default: spa+eng)") parser.add_argument("--psm", type=int, default=0, help="PSM de Tesseract (0 = auto-detect)") parser.add_argument("--raw", action="store_true", help="OCR sin corrección con IA") parser.add_argument("--copy", action="store_true", help="Copiar resultado al portapapeles") parser.add_argument("--model", default=OLLAMA_MODEL, help=f"Modelo Ollama (default: {OLLAMA_MODEL})") parser.add_argument("--detect-only", action="store_true", help="Solo detectar tipo de contenido y salir") parser.add_argument("--debug", action="store_true", help="Modo verbose con información de depuración") args = parser.parse_args() # Verificar dependencias del sistema check_deps() # Configurar modelo Ollama ollama_model = args.model with tempfile.TemporaryDirectory(prefix="ocr-ai-") as tmpdir: if args.file: img_path = args.file if not Path(img_path).exists(): print(f"❌ Archivo no encontrado: {img_path}", file=sys.stderr) sys.exit(1) print(f"📄 Procesando: {img_path}", file=sys.stderr) else: img_path = os.path.join(tmpdir, "captura.png") capturar_pantalla(img_path) # ─── 2. Preprocesar con pipeline por defecto ───────────────────────── proc_path = os.path.join(tmpdir, "procesada.png") print(f"🖼️ Preprocesando con ImageMagick...", file=sys.stderr) # Primera OCR para detectar tipo (con pipeline por defecto) first_ocr_path = os.path.join(tmpdir, "first_ocr.png") preprocesar(img_path, first_ocr_path, DEFAULT_PIPELINE) texto_bruto = ocr(first_ocr_path, args.lang, args.psm if args.psm > 0 else 3) if not texto_bruto: print("⚠️ No se detectó texto en la imagen", file=sys.stderr) sys.exit(0) # ─── 3. Detectar tipo de contenido ────────────────────────────────── tipo = detectar_tipo(texto_bruto) print(f" 📋 Tipo detectado: {tipo.upper()}", file=sys.stderr) if args.detect_only: print(f"\nTexto OCR (primeros 500 chars):\n{texto_bruto[:500]}") return # ─── 4. Reprocesar con pipeline óptimo para el tipo ───────────────── if args.psm == 0: psm_optimo = psm_para_tipo(tipo) else: psm_optimo = args.psm pipeline_optimo = pipeline_para_tipo(tipo) if args.debug: print(f" 🔧 Pipeline: {' '.join(pipeline_optimo)}", file=sys.stderr) print(f" 🔧 PSM: {psm_optimo}", file=sys.stderr) # Solo reprocesar si no era el pipeline por defecto if pipeline_optimo != DEFAULT_PIPELINE: print(f" 🔄 Reprocesando con pipeline optimizado para {tipo}...", file=sys.stderr) preprocesar(img_path, proc_path, pipeline_optimo) else: proc_path = first_ocr_path texto_ocr = ocr(proc_path, args.lang, psm_optimo) if not texto_ocr: print("⚠️ No se detectó texto tras reprocesado", file=sys.stderr) texto_ocr = texto_bruto # Fallback al primer intento # ─── 5. Post-procesar con IA (opcional) ───────────────────────────── if args.raw: resultado = texto_ocr else: print(f" 🧠 Corrigiendo con {ollama_model}...", file=sys.stderr) resultado = corregir_con_ia(texto_ocr, tipo, ollama_model) if not resultado: resultado = texto_ocr # Fallback si Ollama falla # ─── 6. Mostrar resultado ─────────────────────────────────────────── print() print("╔══════════════════════════════════════════════════════════════╗") print(f"║ 📝 RESULTADO ({tipo.upper()}){' ' * (27 - len(tipo))}║") print("╚══════════════════════════════════════════════════════════════╝") print() print(resultado) print() # ─── 7. Copiar al portapapeles ────────────────────────────────────── if args.copy: copiar_portapapeles(resultado) # ─── 8. Estadísticas ──────────────────────────────────────────────── if args.debug: print(f"📊 Estadísticas:", file=sys.stderr) print(f" Tipo detectado: {tipo}", file=sys.stderr) print(f" PSM usado: {psm_optimo}", file=sys.stderr) print(f" Caracteres OCR: {len(texto_ocr)}", file=sys.stderr) print(f" Caracteres tras IA: {len(resultado)}", file=sys.stderr) if not args.raw and texto_ocr != resultado: # Mostrar cambios principales print(f" ✏️ Correcciones aplicadas por IA", file=sys.stderr) if __name__ == "__main__": main()