PixText-Studio / ocr_worker.py
Flourish's picture
Upload application
a73e62b verified
Raw History Blame Contribute Delete
761 Bytes
"""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()