mirror of
https://github.com/DS4SD/docling.git
synced 2025-08-01 15:02:21 +00:00
enable locks for threadsafe pdfium
Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>
This commit is contained in:
parent
1b0ead6907
commit
762a511d0a
@ -12,6 +12,7 @@ from pypdfium2 import PdfPage
|
||||
|
||||
from docling.backend.pdf_backend import PdfDocumentBackend, PdfPageBackend
|
||||
from docling.datamodel.base_models import Cell, Size
|
||||
from docling.utils.locks import pypdfium2_lock
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from docling.datamodel.document import InputDocument
|
||||
@ -182,6 +183,7 @@ class DoclingParseV2PageBackend(PdfPageBackend):
|
||||
padbox.r = page_size.width - padbox.r
|
||||
padbox.t = page_size.height - padbox.t
|
||||
|
||||
with pypdfium2_lock:
|
||||
image = (
|
||||
self._ppage.render(
|
||||
scale=scale * 1.5,
|
||||
@ -189,12 +191,15 @@ class DoclingParseV2PageBackend(PdfPageBackend):
|
||||
crop=padbox.as_tuple(),
|
||||
)
|
||||
.to_pil()
|
||||
.resize(size=(round(cropbox.width * scale), round(cropbox.height * scale)))
|
||||
.resize(
|
||||
size=(round(cropbox.width * scale), round(cropbox.height * scale))
|
||||
)
|
||||
) # We resize the image from 1.5x the given scale to make it sharper.
|
||||
|
||||
return image
|
||||
|
||||
def get_size(self) -> Size:
|
||||
with pypdfium2_lock:
|
||||
return Size(width=self._ppage.get_width(), height=self._ppage.get_height())
|
||||
|
||||
def unload(self):
|
||||
@ -206,6 +211,7 @@ class DoclingParseV2DocumentBackend(PdfDocumentBackend):
|
||||
def __init__(self, in_doc: "InputDocument", path_or_stream: Union[BytesIO, Path]):
|
||||
super().__init__(in_doc, path_or_stream)
|
||||
|
||||
with pypdfium2_lock:
|
||||
self._pdoc = pdfium.PdfDocument(self.path_or_stream)
|
||||
self.parser = pdf_parser_v2("fatal")
|
||||
|
||||
@ -236,6 +242,7 @@ class DoclingParseV2DocumentBackend(PdfDocumentBackend):
|
||||
return len_2
|
||||
|
||||
def load_page(self, page_no: int) -> DoclingParseV2PageBackend:
|
||||
with pypdfium2_lock:
|
||||
return DoclingParseV2PageBackend(
|
||||
self.parser, self.document_hash, page_no, self._pdoc[page_no]
|
||||
)
|
||||
@ -246,5 +253,6 @@ class DoclingParseV2DocumentBackend(PdfDocumentBackend):
|
||||
def unload(self):
|
||||
super().unload()
|
||||
self.parser.unload_document(self.document_hash)
|
||||
with pypdfium2_lock:
|
||||
self._pdoc.close()
|
||||
self._pdoc = None
|
||||
|
@ -13,6 +13,7 @@ from pypdfium2._helpers.misc import PdfiumError
|
||||
|
||||
from docling.backend.pdf_backend import PdfDocumentBackend, PdfPageBackend
|
||||
from docling.datamodel.base_models import Cell
|
||||
from docling.utils.locks import pypdfium2_lock
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from docling.datamodel.document import InputDocument
|
||||
@ -24,6 +25,7 @@ class PyPdfiumPageBackend(PdfPageBackend):
|
||||
def __init__(
|
||||
self, pdfium_doc: pdfium.PdfDocument, document_hash: str, page_no: int
|
||||
):
|
||||
# Note: lock applied by the caller
|
||||
self.valid = True # No better way to tell from pypdfium.
|
||||
try:
|
||||
self._ppage: pdfium.PdfPage = pdfium_doc[page_no]
|
||||
@ -40,6 +42,7 @@ class PyPdfiumPageBackend(PdfPageBackend):
|
||||
|
||||
def get_bitmap_rects(self, scale: float = 1) -> Iterable[BoundingBox]:
|
||||
AREA_THRESHOLD = 0 # 32 * 32
|
||||
with pypdfium2_lock:
|
||||
for obj in self._ppage.get_objects(filter=[pdfium_c.FPDF_PAGEOBJ_IMAGE]):
|
||||
pos = obj.get_pos()
|
||||
cropbox = BoundingBox.from_tuple(
|
||||
@ -52,6 +55,7 @@ class PyPdfiumPageBackend(PdfPageBackend):
|
||||
yield cropbox
|
||||
|
||||
def get_text_in_rect(self, bbox: BoundingBox) -> str:
|
||||
with pypdfium2_lock:
|
||||
if not self.text_page:
|
||||
self.text_page = self._ppage.get_textpage()
|
||||
|
||||
@ -63,6 +67,7 @@ class PyPdfiumPageBackend(PdfPageBackend):
|
||||
return text_piece
|
||||
|
||||
def get_text_cells(self) -> Iterable[Cell]:
|
||||
with pypdfium2_lock:
|
||||
if not self.text_page:
|
||||
self.text_page = self._ppage.get_textpage()
|
||||
|
||||
@ -214,6 +219,7 @@ class PyPdfiumPageBackend(PdfPageBackend):
|
||||
padbox.r = page_size.width - padbox.r
|
||||
padbox.t = page_size.height - padbox.t
|
||||
|
||||
with pypdfium2_lock:
|
||||
image = (
|
||||
self._ppage.render(
|
||||
scale=scale * 1.5,
|
||||
@ -221,12 +227,15 @@ class PyPdfiumPageBackend(PdfPageBackend):
|
||||
crop=padbox.as_tuple(),
|
||||
)
|
||||
.to_pil()
|
||||
.resize(size=(round(cropbox.width * scale), round(cropbox.height * scale)))
|
||||
.resize(
|
||||
size=(round(cropbox.width * scale), round(cropbox.height * scale))
|
||||
)
|
||||
) # We resize the image from 1.5x the given scale to make it sharper.
|
||||
|
||||
return image
|
||||
|
||||
def get_size(self) -> Size:
|
||||
with pypdfium2_lock:
|
||||
return Size(width=self._ppage.get_width(), height=self._ppage.get_height())
|
||||
|
||||
def unload(self):
|
||||
@ -239,6 +248,7 @@ class PyPdfiumDocumentBackend(PdfDocumentBackend):
|
||||
super().__init__(in_doc, path_or_stream)
|
||||
|
||||
try:
|
||||
with pypdfium2_lock:
|
||||
self._pdoc = pdfium.PdfDocument(self.path_or_stream)
|
||||
except PdfiumError as e:
|
||||
raise RuntimeError(
|
||||
@ -246,9 +256,11 @@ class PyPdfiumDocumentBackend(PdfDocumentBackend):
|
||||
) from e
|
||||
|
||||
def page_count(self) -> int:
|
||||
with pypdfium2_lock:
|
||||
return len(self._pdoc)
|
||||
|
||||
def load_page(self, page_no: int) -> PyPdfiumPageBackend:
|
||||
with pypdfium2_lock:
|
||||
return PyPdfiumPageBackend(self._pdoc, self.document_hash, page_no)
|
||||
|
||||
def is_valid(self) -> bool:
|
||||
@ -256,5 +268,6 @@ class PyPdfiumDocumentBackend(PdfDocumentBackend):
|
||||
|
||||
def unload(self):
|
||||
super().unload()
|
||||
with pypdfium2_lock:
|
||||
self._pdoc.close()
|
||||
self._pdoc = None
|
||||
|
3
docling/utils/locks.py
Normal file
3
docling/utils/locks.py
Normal file
@ -0,0 +1,3 @@
|
||||
import threading
|
||||
|
||||
pypdfium2_lock = threading.Lock()
|
Loading…
Reference in New Issue
Block a user