Merge pull request #252 from funstory-ai/dev

fix: add threading locks for concurrent access to pixmap generation

Author: awwaawwa

babeldoc/__init__.py

@@ -1 +1 @@
-__version__ = "0.3.15"
+__version__ = "0.3.16"

babeldoc/const.py

@@ -3,7 +3,7 @@
 import subprocess
 from pathlib import Path
 
-__version__ = "0.3.15"
+__version__ = "0.3.16"
 
 CACHE_FOLDER = Path.home() / ".cache" / "babeldoc"
 

babeldoc/docvision/doclayout.py

@@ -3,6 +3,7 @@
 import logging
 import platform
 import re
+import threading
 from collections.abc import Generator
 
 import cv2
@@ -121,6 +122,7 @@ def __init__(self, model_path: str):
             model.SerializeToString(),
             providers=providers,
         )
+        self.lock = threading.Lock()
 
     @staticmethod
     def from_pretrained():
@@ -279,7 +281,8 @@ def handle_document(
     ]:
         for page in pages:
             translate_config.raise_if_cancelled()
-            pix = mupdf_doc[page.page_number].get_pixmap(dpi=72)
+            with self.lock:
+                pix = mupdf_doc[page.page_number].get_pixmap(dpi=72)
             image = np.fromstring(pix.samples, np.uint8).reshape(
                 pix.height,
                 pix.width,

babeldoc/docvision/rpc_doclayout.py

@@ -1,4 +1,5 @@
 import logging
+import threading
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 
@@ -125,6 +126,7 @@ def __init__(self, host: str = "http://localhost:8000"):
         self.host = host
         self._stride = 32  # Default stride value
         self._names = ["text", "title", "list", "table", "figure"]
+        self.lock = threading.Lock()
 
     @property
     def stride(self) -> int:
@@ -253,7 +255,8 @@ def predict_page(
         self, page, mupdf_doc: pymupdf.Document, translate_config, save_debug_image
     ):
         translate_config.raise_if_cancelled()
-        pix = mupdf_doc[page.page_number].get_pixmap(dpi=72)
+        with self.lock:
+            pix = mupdf_doc[page.page_number].get_pixmap(dpi=72)
         image = np.fromstring(pix.samples, np.uint8).reshape(
             pix.height,
             pix.width,

babeldoc/docvision/table_detection/rapidocr.py

@@ -1,5 +1,6 @@
 import logging
 import re
+import threading
 from collections.abc import Generator
 
 import cv2
@@ -97,6 +98,7 @@ def __init__(self):
             det_use_dml=self.use_dml,
         )
         self.names = {0: "table_text"}
+        self.lock = threading.Lock()
 
     @property
     def stride(self):
@@ -233,7 +235,8 @@ def handle_document(
     ]:
         for page in pages:
             translate_config.raise_if_cancelled()
-            pix = mupdf_doc[page.page_number].get_pixmap(dpi=72)
+            with self.lock:
+                pix = mupdf_doc[page.page_number].get_pixmap(dpi=72)
             image = np.fromstring(pix.samples, np.uint8).reshape(
                 pix.height,
                 pix.width,

babeldoc/main.py

@@ -23,7 +23,7 @@
 from babeldoc.translation_config import WatermarkOutputMode
 
 logger = logging.getLogger(__name__)
-__version__ = "0.3.15"
+__version__ = "0.3.16"
 
 
 def create_parser():

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "BabelDOC"
-version = "0.3.15"
+version = "0.3.16"
 description = "Yet Another Document Translator"
 license = "AGPL-3.0"
 readme = "README.md"
@@ -148,7 +148,7 @@ pythonpath = [".", "src"]
 testpaths = ["tests"]
 
 [bumpver]
-current_version = "0.3.15"
+current_version = "0.3.16"
 version_pattern = "MAJOR.MINOR.PATCH[.PYTAGNUM]"
 
 [bumpver.file_patterns]