Spaces:
Sleeping
Sleeping
File size: 2,685 Bytes
7a11b03 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | """Subprocess worker for PaddleOCR.
Running PaddleOCR in a separate process protects the FastAPI server from
native-library crashes in PaddlePaddle.
"""
import json
import os
import sys
from pathlib import Path
def extract_text_items(ocr_result) -> list[str]:
"""Normalize PaddleOCR output from common v2 and v3 response formats."""
text_items = []
if not ocr_result:
return text_items
if isinstance(ocr_result, list):
for item in ocr_result:
text_items.extend(extract_text_items(item))
return text_items
if isinstance(ocr_result, dict):
for key in ("rec_texts", "texts"):
values = ocr_result.get(key)
if isinstance(values, list):
text_items.extend(str(value) for value in values if str(value).strip())
for key in ("res", "data"):
if key in ocr_result:
text_items.extend(extract_text_items(ocr_result[key]))
return text_items
if isinstance(ocr_result, tuple) and len(ocr_result) >= 2:
possible_text = ocr_result[1]
if isinstance(possible_text, tuple) and possible_text:
text_items.append(str(possible_text[0]))
elif isinstance(possible_text, str):
text_items.append(possible_text)
return text_items
return text_items
def run_ocr(image_path: Path) -> str:
"""Run PaddleOCR on one image and return cleaned text."""
from paddleocr import PaddleOCR
try:
ocr = PaddleOCR(
lang="en",
use_doc_orientation_classify=False,
use_doc_unwarping=False,
use_textline_orientation=False,
)
except ValueError:
ocr = PaddleOCR(use_angle_cls=True, lang="en")
if hasattr(ocr, "predict"):
result = ocr.predict(str(image_path))
else:
result = ocr.ocr(str(image_path), cls=True)
text_items = extract_text_items(result)
return " ".join(" ".join(text_items).split())
def main() -> int:
"""CLI entry point."""
if len(sys.argv) != 3:
print(json.dumps({"error": "Usage: python -m src.ocr_worker IMAGE_PATH CACHE_DIR"}))
return 2
image_path = Path(sys.argv[1])
cache_dir = Path(sys.argv[2])
cache_dir.mkdir(parents=True, exist_ok=True)
os.environ.setdefault("PADDLE_PDX_CACHE_HOME", str(cache_dir))
os.environ.setdefault("PADDLEOCR_DISABLE_AUTO_LOGGING_CONFIG", "1")
try:
text = run_ocr(image_path)
except Exception as error:
print(json.dumps({"error": str(error)}))
return 1
print(json.dumps({"text": text}))
return 0
if __name__ == "__main__":
raise SystemExit(main())
|