Like 0
Fork 0

Last active 1 day ago

atareao's Avatar atareao revised this gist 1 day ago. Go to revision

1 file changed, 446 insertions

ocr_ai.py(file created)

@@ -0,0 +1,446 @@
1 + #!/usr/bin/env python3
2 + """
3 + ocr-ai.py — OCR inteligente con detección automática de contenido
4 + y post-procesado con IA (Ollama).
5 +
6 + Captura una región de la pantalla o procesa una imagen existente,
7 + detecta automáticamente el tipo de contenido (código, tabla, prosa),
8 + ajusta los parámetros de Tesseract, y limpia errores con Ollama.
9 +
10 + Uso:
11 + python3 ocr-ai.py # Captura + OCR + IA
12 + python3 ocr-ai.py -f captura.png # Procesa imagen existente
13 + python3 ocr-ai.py -f captura.png --raw # Solo OCR, sin IA
14 + python3 ocr-ai.py -f captura.png --copy # OCR + IA + al portapapeles
15 +
16 + Dependencias:
17 + - Python 3.10+
18 + - ImageMagick (magick)
19 + - Tesseract (tesseract-ocr + tesseract-ocr-spa + tesseract-ocr-eng)
20 + - wl-clipboard (Wayland) o xclip (X11) para --copy
21 + - Ollama (ollama serve corriendo, modelo recomendado: llama3.2)
22 + """
23 +
24 + import argparse
25 + import json
26 + import os
27 + import re
28 + import subprocess
29 + import sys
30 + import tempfile
31 + from pathlib import Path
32 +
33 + import requests
34 +
35 +
36 + # ─── Configuración ─────────────────────────────────────────────────────────────
37 +
38 + OLLAMA_HOST = "http://localhost:11434"
39 + OLLAMA_MODEL = "llama3.2"
40 + OLLAMA_TIMEOUT = 60
41 +
42 + # Pipeline de ImageMagick por defecto
43 + DEFAULT_PIPELINE = [
44 + "-colorspace", "Gray",
45 + "-normalize",
46 + "-threshold", "60%",
47 + "-deskew", "40%",
48 + "-sharpen", "0x1",
49 + ]
50 +
51 + # Pipeline más agresivo para código (más contraste, menos ruido)
52 + CODE_PIPELINE = [
53 + "-colorspace", "Gray",
54 + "-normalize",
55 + "-contrast-stretch", "2%",
56 + "-threshold", "55%",
57 + "-deskew", "40%",
58 + "-sharpen", "0x2",
59 + "-morphology", "Close", "Octagon",
60 + ]
61 +
62 + # Pipeline para tablas (preserva líneas)
63 + TABLE_PIPELINE = [
64 + "-colorspace", "Gray",
65 + "-normalize",
66 + "-threshold", "50%",
67 + "-deskew", "30%",
68 + "-sharpen", "0x1",
69 + "-morphology", "Erode", "Octagon",
70 + ]
71 +
72 +
73 + # ─── Detección de tipo de contenido ───────────────────────────────────────────
74 +
75 + def detectar_tipo(texto: str) -> str:
76 + """
77 + Analiza el texto extraído por OCR y detecta el tipo de contenido.
78 + Devuelve: 'codigo', 'tabla', 'prosa' o 'desconocido'.
79 + """
80 + if not texto or len(texto.strip()) < 10:
81 + return "desconocido"
82 +
83 + lineas = texto.strip().split("\n")
84 + num_lineas = len(lineas)
85 + palabras = texto.split()
86 + num_palabras = len(palabras)
87 +
88 + # Heurísticas para código
89 + patrones_codigo = [
90 + 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',
91 + r'[{}()\[\];]',
92 + r'==|!=|<=|>=|->|=>|\+\+|--',
93 + r'^\s*[#//]\s',
94 + r'\.\w+\(',
95 + r'["\'].*["\']\s*[;,]',
96 + ]
97 + puntuacion_codigo = 0
98 + for patron in patrones_codigo:
99 + if re.search(patron, texto, re.IGNORECASE | re.MULTILINE):
100 + puntuacion_codigo += 2
101 +
102 + # Más peso si hay indentación consistente (código suele estar indentado)
103 + indentadas = sum(1 for l in lineas if l.startswith(" ") or l.startswith("\t"))
104 + if indentadas > num_lineas * 0.3:
105 + puntuacion_codigo += 3
106 +
107 + # Puntuación por densidad de palabras (código tiene menos variedad léxica)
108 + if num_palabras > 0 and len(set(palabras)) / num_palabras < 0.4:
109 + puntuacion_codigo += 2
110 +
111 + # Heurísticas para tabla (más restrictivas)
112 + patrones_tabla = [
113 + r'\|.*\|.*\|',
114 + r'\+[-+]+\+',
115 + r'(Nº|#|Item|Producto|Precio|Cantidad|Total|Unidades)\s*\||\|.*\|',
116 + r'^\s*\d+\..*\d+[.,]\d+',
117 + ]
118 + puntuacion_tabla = 0
119 + for patron in patrones_tabla:
120 + if re.search(patron, texto, re.IGNORECASE | re.MULTILINE):
121 + puntuacion_tabla += 2
122 +
123 + # Tablas suelen tener muchas líneas cortas con estructura repetitiva
124 + lineas_cortas = sum(1 for l in lineas if 5 < len(l.strip()) < 40)
125 + if num_lineas > 3 and lineas_cortas > num_lineas * 0.5:
126 + puntuacion_tabla += 2
127 +
128 + # Decisión
129 + if puntuacion_codigo >= puntuacion_tabla and puntuacion_codigo >= 4:
130 + return "codigo"
131 + elif puntuacion_tabla > puntuacion_codigo and puntuacion_tabla >= 3:
132 + return "tabla"
133 + elif num_palabras > 20 and puntuacion_codigo < 4 and puntuacion_tabla < 3:
134 + return "prosa"
135 + else:
136 + return "desconocido"
137 +
138 +
139 + def psm_para_tipo(tipo: str) -> int:
140 + """Devuelve el PSM de Tesseract óptimo para cada tipo."""
141 + return {
142 + "codigo": 6, # Bloque de texto uniforme
143 + "tabla": 6, # Bloque uniforme (mejor para TSV)
144 + "prosa": 4, # Texto a una columna
145 + "desconocido": 3, # Detección automática
146 + }.get(tipo, 3)
147 +
148 +
149 + def pipeline_para_tipo(tipo: str) -> list:
150 + """Devuelve el pipeline de ImageMagick óptimo para cada tipo."""
151 + return {
152 + "codigo": CODE_PIPELINE,
153 + "tabla": TABLE_PIPELINE,
154 + "prosa": DEFAULT_PIPELINE,
155 + "desconocido": DEFAULT_PIPELINE,
156 + }.get(tipo, DEFAULT_PIPELINE)
157 +
158 +
159 + # ─── Herramientas del sistema ─────────────────────────────────────────────────
160 +
161 + def check_deps():
162 + """Verifica que las dependencias del sistema están instaladas."""
163 + faltan = []
164 + for cmd, nombre in [("convert", "ImageMagick"), ("tesseract", "Tesseract")]:
165 + if subprocess.run(["which", cmd], capture_output=True).returncode != 0:
166 + faltan.append(nombre)
167 + if faltan:
168 + print(f"❌ Faltan dependencias: {', '.join(faltan)}", file=sys.stderr)
169 + print(f" Instala: sudo apt install imagemagick tesseract-ocr tesseract-ocr-spa tesseract-ocr-eng", file=sys.stderr)
170 + sys.exit(1)
171 +
172 +
173 + def capturar_pantalla(output_path: str):
174 + """Captura una región de la pantalla con import (ImageMagick)."""
175 + print("🔍 Selecciona una región con el ratón...", file=sys.stderr)
176 + result = subprocess.run(
177 + ["import", output_path],
178 + capture_output=True, text=True,
179 + )
180 + if result.returncode != 0:
181 + print("❌ Captura cancelada o fallida", file=sys.stderr)
182 + sys.exit(1)
183 + if not Path(output_path).exists() or Path(output_path).stat().st_size == 0:
184 + print("❌ No se generó la captura", file=sys.stderr)
185 + sys.exit(1)
186 + print(f" ✅ Captura guardada: {output_path}", file=sys.stderr)
187 +
188 +
189 + def preprocesar(input_path: str, output_path: str, pipeline: list):
190 + """Aplica el pipeline de ImageMagick a la imagen."""
191 + cmd = ["convert", input_path] + pipeline + [output_path]
192 + result = subprocess.run(cmd, capture_output=True, text=True)
193 + if result.returncode != 0:
194 + print(f"❌ Error en preprocesado: {result.stderr}", file=sys.stderr)
195 + sys.exit(1)
196 +
197 +
198 + def ocr(image_path: str, idiomas: str, psm: int) -> str:
199 + """Ejecuta Tesseract OCR y devuelve el texto extraído."""
200 + cmd = [
201 + "tesseract", image_path, "stdout",
202 + "-l", idiomas,
203 + "--psm", str(psm),
204 + "--oem", "1",
205 + ]
206 + result = subprocess.run(cmd, capture_output=True, text=True)
207 + if result.returncode != 0:
208 + print(f"❌ Error en OCR: {result.stderr}", file=sys.stderr)
209 + return ""
210 + return result.stdout.strip()
211 +
212 +
213 + def copiar_portapapeles(texto: str):
214 + """Copia texto al portapapeles (Wayland o X11)."""
215 + if subprocess.run(["which", "wl-copy"], capture_output=True).returncode == 0:
216 + subprocess.run(["wl-copy"], input=texto, text=True)
217 + elif subprocess.run(["which", "xclip"], capture_output=True).returncode == 0:
218 + subprocess.run(["xclip", "-selection", "clipboard"], input=texto, text=True)
219 + else:
220 + print("⚠️ Ni wl-copy ni xclip disponibles. Texto solo por stdout.", file=sys.stderr)
221 + return
222 + print(" 📋 Copiado al portapapeles", file=sys.stderr)
223 +
224 +
225 + # ─── Post-procesado con Ollama ────────────────────────────────────────────────
226 +
227 + def llamar_ollama(system_prompt: str, user_prompt: str, model: str = OLLAMA_MODEL) -> str | None:
228 + """Llama a Ollama y devuelve el texto generado."""
229 + payload = {
230 + "model": model,
231 + "messages": [
232 + {"role": "system", "content": system_prompt},
233 + {"role": "user", "content": user_prompt},
234 + ],
235 + "stream": False,
236 + "options": {"temperature": 0.1, "num_predict": 2048},
237 + }
238 +
239 + try:
240 + resp = requests.post(
241 + f"{OLLAMA_HOST}/api/chat",
242 + json=payload,
243 + timeout=OLLAMA_TIMEOUT,
244 + )
245 + resp.raise_for_status()
246 + return resp.json().get("message", {}).get("content", "").strip()
247 + except requests.exceptions.ConnectionError:
248 + print(" ⚠️ Ollama no está disponible. Omitiendo corrección con IA.", file=sys.stderr)
249 + return None
250 + except Exception as e:
251 + print(f" ⚠️ Error en Ollama: {e}. Omitiendo corrección.", file=sys.stderr)
252 + return None
253 +
254 +
255 + def corregir_con_ia(texto: str, tipo: str, model: str = OLLAMA_MODEL) -> str:
256 + """Corrige errores de OCR usando Ollama, adaptado al tipo de contenido."""
257 + prompts = {
258 + "codigo": {
259 + "system": (
260 + "Eres un corrector de OCR especializado en código fuente. "
261 + "Tu tarea es corregir errores de reconocimiento óptico en código.\n\n"
262 + "REGLAS:\n"
263 + "1. Corrige caracteres mal reconocidos (0 por O, l por 1, etc.)\n"
264 + "2. Restaura indentación y formato del código original\n"
265 + "3. NO cambies nombres de variables, funciones o lógica\n"
266 + "4. NO inventes líneas de código que no están\n"
267 + "5. Mantén comentarios exactamente como están\n"
268 + "6. Preserva la sintaxis del lenguaje\n"
269 + "7. Devuelve SOLO el código corregido, sin explicaciones"
270 + ),
271 + "user": f"Corrige los errores de OCR en este código:\n\n```\n{texto}\n```",
272 + },
273 + "tabla": {
274 + "system": (
275 + "Eres un corrector de OCR especializado en tablas y datos numéricos. "
276 + "Tu tarea es corregir errores de reconocimiento óptico en tablas.\n\n"
277 + "REGLAS:\n"
278 + "1. Corrige números mal reconocidos (1.000 por 1000, etc.)\n"
279 + "2. Corrige símbolos monetarios (€ por EUR o similar)\n"
280 + "3. Preserva la estructura de filas y columnas\n"
281 + "4. NO inventes datos que no están en el original\n"
282 + "5. NO rellenes celdas vacías con datos inventados\n"
283 + "6. Si no estás seguro de un valor, déjalo como está\n"
284 + "7. Devuelve SOLO la tabla corregida, sin explicaciones"
285 + ),
286 + "user": f"Corrige los errores de OCR en esta tabla:\n\n```\n{texto}\n```",
287 + },
288 + "prosa": {
289 + "system": (
290 + "Eres un corrector de OCR especializado en texto en español. "
291 + "Tu tarea es corregir errores de reconocimiento óptico.\n\n"
292 + "REGLAS:\n"
293 + "1. Corrige caracteres mal reconocidos\n"
294 + "2. Corrige espacios perdidos o añadidos entre palabras\n"
295 + "3. Corrige puntuación y acentos cuando sea evidente\n"
296 + "4. NO cambies el significado del texto\n"
297 + "5. NO añadas información que no está\n"
298 + "6. Preserva el estilo y tono original\n"
299 + "7. Si una palabra no la reconoces bien, déjala como está\n"
300 + "8. Devuelve SOLO el texto corregido, sin explicaciones"
301 + ),
302 + "user": f"Corrige los errores de OCR en este texto:\n\n```\n{texto}\n```",
303 + },
304 + "desconocido": {
305 + "system": (
306 + "Eres un corrector de OCR. Tu tarea es corregir errores "
307 + "de reconocimiento óptico en el texto que se te pasa.\n\n"
308 + "REGLAS:\n"
309 + "1. Corrige solo errores evidentes de OCR\n"
310 + "2. NO inventes contenido\n"
311 + "3. NO cambies el formato o estructura\n"
312 + "4. Devuelve SOLO el texto corregido, sin explicaciones"
313 + ),
314 + "user": f"Corrige los errores de OCR en este texto:\n\n```\n{texto}\n```",
315 + },
316 + }
317 +
318 + p = prompts.get(tipo, prompts["desconocido"])
319 + resultado = llamar_ollama(p["system"], p["user"], model)
320 + return resultado if resultado else texto
321 +
322 +
323 + # ─── Main ─────────────────────────────────────────────────────────────────────
324 +
325 + def main():
326 +
327 + parser = argparse.ArgumentParser(
328 + description="ocr-ai: OCR inteligente con detección de contenido y corrección con IA",
329 + formatter_class=argparse.RawDescriptionHelpFormatter,
330 + epilog="""\
331 + Ejemplos:
332 + python3 ocr-ai.py # Captura + OCR + IA
333 + python3 ocr-ai.py -f captura.png # Procesa imagen existente
334 + python3 ocr-ai.py -f captura.png --raw # Solo OCR, sin IA
335 + python3 ocr-ai.py -f captura.png --copy --lang eng # Código al portapapeles
336 + """,
337 + )
338 + parser.add_argument("-f", "--file", help="Archivo de imagen a procesar (si no, captura interactiva)")
339 + parser.add_argument("--lang", default="spa+eng", help="Idiomas para Tesseract (default: spa+eng)")
340 + parser.add_argument("--psm", type=int, default=0, help="PSM de Tesseract (0 = auto-detect)")
341 + parser.add_argument("--raw", action="store_true", help="OCR sin corrección con IA")
342 + parser.add_argument("--copy", action="store_true", help="Copiar resultado al portapapeles")
343 + parser.add_argument("--model", default=OLLAMA_MODEL, help=f"Modelo Ollama (default: {OLLAMA_MODEL})")
344 + parser.add_argument("--detect-only", action="store_true", help="Solo detectar tipo de contenido y salir")
345 + parser.add_argument("--debug", action="store_true", help="Modo verbose con información de depuración")
346 +
347 + args = parser.parse_args()
348 +
349 + # Verificar dependencias del sistema
350 + check_deps()
351 +
352 + # Configurar modelo Ollama
353 + ollama_model = args.model
354 +
355 + with tempfile.TemporaryDirectory(prefix="ocr-ai-") as tmpdir:
356 + if args.file:
357 + img_path = args.file
358 + if not Path(img_path).exists():
359 + print(f"❌ Archivo no encontrado: {img_path}", file=sys.stderr)
360 + sys.exit(1)
361 + print(f"📄 Procesando: {img_path}", file=sys.stderr)
362 + else:
363 + img_path = os.path.join(tmpdir, "captura.png")
364 + capturar_pantalla(img_path)
365 +
366 + # ─── 2. Preprocesar con pipeline por defecto ─────────────────────────
367 + proc_path = os.path.join(tmpdir, "procesada.png")
368 + print(f"🖼️ Preprocesando con ImageMagick...", file=sys.stderr)
369 +
370 + # Primera OCR para detectar tipo (con pipeline por defecto)
371 + first_ocr_path = os.path.join(tmpdir, "first_ocr.png")
372 + preprocesar(img_path, first_ocr_path, DEFAULT_PIPELINE)
373 + texto_bruto = ocr(first_ocr_path, args.lang, args.psm if args.psm > 0 else 3)
374 +
375 + if not texto_bruto:
376 + print("⚠️ No se detectó texto en la imagen", file=sys.stderr)
377 + sys.exit(0)
378 +
379 + # ─── 3. Detectar tipo de contenido ──────────────────────────────────
380 + tipo = detectar_tipo(texto_bruto)
381 + print(f" 📋 Tipo detectado: {tipo.upper()}", file=sys.stderr)
382 +
383 + if args.detect_only:
384 + print(f"\nTexto OCR (primeros 500 chars):\n{texto_bruto[:500]}")
385 + return
386 +
387 + # ─── 4. Reprocesar con pipeline óptimo para el tipo ─────────────────
388 + if args.psm == 0:
389 + psm_optimo = psm_para_tipo(tipo)
390 + else:
391 + psm_optimo = args.psm
392 +
393 + pipeline_optimo = pipeline_para_tipo(tipo)
394 + if args.debug:
395 + print(f" 🔧 Pipeline: {' '.join(pipeline_optimo)}", file=sys.stderr)
396 + print(f" 🔧 PSM: {psm_optimo}", file=sys.stderr)
397 +
398 + # Solo reprocesar si no era el pipeline por defecto
399 + if pipeline_optimo != DEFAULT_PIPELINE:
400 + print(f" 🔄 Reprocesando con pipeline optimizado para {tipo}...", file=sys.stderr)
401 + preprocesar(img_path, proc_path, pipeline_optimo)
402 + else:
403 + proc_path = first_ocr_path
404 +
405 + texto_ocr = ocr(proc_path, args.lang, psm_optimo)
406 +
407 + if not texto_ocr:
408 + print("⚠️ No se detectó texto tras reprocesado", file=sys.stderr)
409 + texto_ocr = texto_bruto # Fallback al primer intento
410 +
411 + # ─── 5. Post-procesar con IA (opcional) ─────────────────────────────
412 + if args.raw:
413 + resultado = texto_ocr
414 + else:
415 + print(f" 🧠 Corrigiendo con {ollama_model}...", file=sys.stderr)
416 + resultado = corregir_con_ia(texto_ocr, tipo, ollama_model)
417 + if not resultado:
418 + resultado = texto_ocr # Fallback si Ollama falla
419 +
420 + # ─── 6. Mostrar resultado ───────────────────────────────────────────
421 + print()
422 + print("╔══════════════════════════════════════════════════════════════╗")
423 + print(f"║ 📝 RESULTADO ({tipo.upper()}){' ' * (27 - len(tipo))}║")
424 + print("╚══════════════════════════════════════════════════════════════╝")
425 + print()
426 + print(resultado)
427 + print()
428 +
429 + # ─── 7. Copiar al portapapeles ──────────────────────────────────────
430 + if args.copy:
431 + copiar_portapapeles(resultado)
432 +
433 + # ─── 8. Estadísticas ────────────────────────────────────────────────
434 + if args.debug:
435 + print(f"📊 Estadísticas:", file=sys.stderr)
436 + print(f" Tipo detectado: {tipo}", file=sys.stderr)
437 + print(f" PSM usado: {psm_optimo}", file=sys.stderr)
438 + print(f" Caracteres OCR: {len(texto_ocr)}", file=sys.stderr)
439 + print(f" Caracteres tras IA: {len(resultado)}", file=sys.stderr)
440 + if not args.raw and texto_ocr != resultado:
441 + # Mostrar cambios principales
442 + print(f" ✏️ Correcciones aplicadas por IA", file=sys.stderr)
443 +
444 +
445 + if __name__ == "__main__":
446 + main()
Newer Older

Powered by Opengist ⋅ Load: 108ms⋅

English