ocr_worker.py

"""Keep OCR's native CPU thread pools out of the ZeroGPU parent process."""
import json
import os
import sys
from pathlib import Path

os.environ["SPACES_ZERO_GPU"] = "0"
os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["PIXTEXT_CPU_THREADS"] = "1"

import torch
from PIL import Image
from pixtext_ocr import detect_words


def main():
    source, models, output, text_threshold, link_threshold = sys.argv[1:]
    torch.set_num_threads(1)
    with Image.open(source) as image:
        result = detect_words(image, Path(models), float(text_threshold), float(link_threshold))
    Path(output).write_text(json.dumps(result, ensure_ascii=False))


if __name__ == "__main__":
    main()