diff --git a/packages/ai/python/ocr.py b/packages/ai/python/ocr.py index d79313ca..ad089cd1 100644 --- a/packages/ai/python/ocr.py +++ b/packages/ai/python/ocr.py @@ -1,6 +1,7 @@ """Text extraction from images using Tesseract or PaddleOCR.""" import sys import json +import os def emit_progress(percent, stage): @@ -8,6 +9,49 @@ def emit_progress(percent, stage): print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True) +def run_tesseract(input_path, language): + """Run Tesseract OCR.""" + import subprocess + + lang_map = {"en": "eng", "de": "deu", "fr": "fra", "es": "spa", "zh": "chi_sim", "ja": "jpn", "ko": "kor"} + tess_lang = lang_map.get(language, "eng") + + emit_progress(30, "Scanning") + result = subprocess.run( + ["tesseract", input_path, "stdout", "-l", tess_lang], + capture_output=True, + text=True, + timeout=120, + ) + emit_progress(70, "Extracting text") + text = result.stdout.strip() + if result.returncode != 0 and not text: + raise RuntimeError(result.stderr.strip() or "Tesseract failed") + return text, "tesseract" + + +def run_paddleocr(input_path, language): + """Run PaddleOCR.""" + os.environ["PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK"] = "True" + from paddleocr import PaddleOCR + + emit_progress(20, "Loading") + ocr = PaddleOCR(lang=language) + emit_progress(30, "Scanning") + result = ocr.ocr(input_path) + emit_progress(70, "Extracting text") + text = "\n".join( + [ + line[1][0] + for res in result + if res + for line in res + if line and line[1] + ] + ) + return text, "paddleocr" + + def main(): input_path = sys.argv[1] settings = json.loads(sys.argv[2]) if len(sys.argv) > 2 else {} @@ -16,74 +60,48 @@ def main(): language = settings.get("language", "en") try: - emit_progress(10, "Loading OCR engine") + emit_progress(10, "Preparing") + if engine == "paddleocr": try: - from paddleocr import PaddleOCR - - ocr = PaddleOCR(use_angle_cls=True, lang=language) - emit_progress(30, "Analyzing text regions") - result = ocr.ocr(input_path, cls=True) - emit_progress(70, "Extracting text") - text = "\n".join( - [ - line[1][0] - for res in result - if res - for line in res - if line and line[1] - ] - ) - emit_progress(95, "Formatting results") - print( - json.dumps({"success": True, "text": text, "engine": "paddleocr"}) - ) + text, used_engine = run_paddleocr(input_path, language) except ImportError: print( json.dumps( { "success": False, - "error": "PaddleOCR is not installed. Install with: pip install paddleocr paddlepaddle", + "error": "PaddleOCR is not installed", } ) ) sys.exit(1) - else: - # Tesseract via subprocess - import subprocess - - lang_map = {"en": "eng", "de": "deu", "fr": "fra", "es": "spa", "zh": "chi_sim", "ja": "jpn", "ko": "kor"} - tess_lang = lang_map.get(language, "eng") - - try: - emit_progress(30, "Running Tesseract") - result = subprocess.run( - ["tesseract", input_path, "stdout", "-l", tess_lang], - capture_output=True, - text=True, - timeout=120, - ) - emit_progress(70, "Extracting text") - text = result.stdout.strip() - if result.returncode != 0 and not text: - raise RuntimeError(result.stderr.strip() or "Tesseract failed") - emit_progress(95, "Formatting results") - print( - json.dumps( - {"success": True, "text": text, "engine": "tesseract"} + except Exception: + # PaddleOCR failed at runtime — fall back to Tesseract + emit_progress(25, "Falling back") + try: + text, used_engine = run_tesseract(input_path, language) + except FileNotFoundError: + print( + json.dumps({"success": False, "error": "OCR engines unavailable"}) ) - ) + sys.exit(1) + else: + try: + text, used_engine = run_tesseract(input_path, language) except FileNotFoundError: print( json.dumps( { "success": False, - "error": "Tesseract is not installed. Install with: apt-get install tesseract-ocr", + "error": "Tesseract is not installed", } ) ) sys.exit(1) + emit_progress(95, "Done") + print(json.dumps({"success": True, "text": text, "engine": used_engine})) + except Exception as e: print(json.dumps({"success": False, "error": str(e)})) sys.exit(1)