BabelDOC 0.5.11__tar.gz → 0.5.12__tar.gz

{babeldoc-0.5.11 → babeldoc-0.5.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: BabelDOC
-Version: 0.5.11
+Version: 0.5.12
 Summary: Yet Another Document Translator
 Project-URL: Homepage, https://github.com/funstory-ai/BabelDOC
 Project-URL: Issues, https://github.com/funstory-ai/BabelDOC/issues

babeldoc-0.5.12/babeldoc/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.5.12"

{babeldoc-0.5.11 → babeldoc-0.5.12}/babeldoc/const.py

@@ -1,9 +1,12 @@
+import itertools
+import multiprocessing as mp
 import os
 import shutil
 import subprocess
+import threading
 from pathlib import Path
 
-__version__ = "0.5.11"
+__version__ = "0.5.12"
 
 CACHE_FOLDER = Path.home() / ".cache" / "babeldoc"
 
@@ -39,3 +42,43 @@ except (OSError, FileNotFoundError, subprocess.CalledProcessError):
 TIKTOKEN_CACHE_FOLDER = CACHE_FOLDER / "tiktoken"
 TIKTOKEN_CACHE_FOLDER.mkdir(parents=True, exist_ok=True)
 os.environ["TIKTOKEN_CACHE_DIR"] = str(TIKTOKEN_CACHE_FOLDER)
+
+
+_process_pool = None
+_process_pool_lock = threading.Lock()
+_ENABLE_PROCESS_POOL = False
+
+
+def enable_process_pool():
+    # Development and Testing ONLY API
+    global _ENABLE_PROCESS_POOL
+    _ENABLE_PROCESS_POOL = True
+
+
+# macos & windows use spawn mode
+# linux use forkserver mode
+
+
+def get_process_pool():
+    if not _ENABLE_PROCESS_POOL:
+        return None
+    global _process_pool
+    with _process_pool_lock:
+        if _process_pool is None:
+            # Create pool only in main process
+            if mp.current_process().name != "MainProcess":
+                return None
+
+            _process_pool = mp.Pool()
+        return _process_pool
+
+
+def batched(iterable, n, *, strict=False):
+    # batched('ABCDEFG', 3) → ABC DEF G
+    if n < 1:
+        raise ValueError("n must be at least one")
+    iterator = iter(iterable)
+    while batch := tuple(itertools.islice(iterator, n)):
+        if strict and len(batch) != n:
+            raise ValueError("batched(): incomplete batch")
+        yield batch

{babeldoc-0.5.11 → babeldoc-0.5.12}/babeldoc/docvision/rpc_doclayout6.py

@@ -28,7 +28,9 @@ from babeldoc.format.pdf.document_il.utils.extract_char import (
 )
 from babeldoc.format.pdf.document_il.utils.fontmap import FontMapper
 from babeldoc.format.pdf.document_il.utils.layout_helper import SPACE_REGEX
-from babeldoc.format.pdf.document_il.utils.mupdf_helper import get_no_rotation_img
+from babeldoc.format.pdf.document_il.utils.mupdf_helper import (
+    get_no_rotation_img_multiprocess,
+)
 
 logger = logging.getLogger(__name__)
 DPI = 150
@@ -141,7 +143,7 @@ def predict_layout(
         f"{host}/inference",
         json=request_data,
         headers={"Accept": "application/json", "Content-Type": "application/json"},
-        timeout=1800,
+        timeout=30,
         follow_redirects=True,
     )
 
@@ -233,7 +235,7 @@ def predict_layout2(
             "Content-Type": "application/msgpack",
             "Accept": "application/msgpack",
         },
-        timeout=480,
+        timeout=30,
         follow_redirects=True,
     )
 
@@ -569,18 +571,19 @@ class RpcDocLayoutModel(DocLayoutModel):
 
         return results
 
-    def predict_page(
-        self, page, mupdf_doc: pymupdf.Document, translate_config, save_debug_image
-    ):
+    def predict_page(self, page, pdf_bytes: Path, translate_config, save_debug_image):
         translate_config.raise_if_cancelled()
-        with self.lock:
-            # pix = mupdf_doc[page.page_number].get_pixmap(dpi=72)
-            pix = get_no_rotation_img(mupdf_doc[page.page_number], dpi=DPI)
-        image = np.frombuffer(pix.samples, np.uint8).reshape(
-            pix.height,
-            pix.width,
-            3,
-        )[:, :, ::-1]
+        # doc = pymupdf.open(io.BytesIO(pdf_bytes))
+        # with self.lock:
+        # pix = mupdf_doc[page.page_number].get_pixmap(dpi=72)
+        image = get_no_rotation_img_multiprocess(
+            pdf_bytes.as_posix(), page.page_number, dpi=DPI
+        )
+        # image = np.frombuffer(pix.samples, np.uint8).reshape(
+        #     pix.height,
+        #     pix.width,
+        #     3,
+        # )[:, :, ::-1]
         char_boxes = convert_page_to_char_boxes(page)
         lines = process_page_chars_to_lines(char_boxes)
         predict_result = self.predict_image(image, 800, lines)
@@ -594,11 +597,13 @@ class RpcDocLayoutModel(DocLayoutModel):
         translate_config,
         save_debug_image,
     ):
+        layout_temp_path = translate_config.get_working_file_path("layout.temp.pdf")
+        mupdf_doc.save(layout_temp_path.as_posix())
         with ThreadPoolExecutor(max_workers=32) as executor:
             yield from executor.map(
                 self.predict_page,
                 pages,
-                (mupdf_doc for _ in range(len(pages))),
+                (layout_temp_path for _ in range(len(pages))),
                 (translate_config for _ in range(len(pages))),
                 (save_debug_image for _ in range(len(pages))),
             )

{babeldoc-0.5.11 → babeldoc-0.5.12}/babeldoc/format/pdf/document_il/midend/detect_scanned_file.py

@@ -9,6 +9,7 @@ from skimage.metrics import structural_similarity
 from babeldoc.babeldoc_exception.BabelDOCException import ScannedPDFError
 from babeldoc.format.pdf.document_il import il_version_1
 from babeldoc.format.pdf.document_il.backend.pdf_creater import PDFCreater
+from babeldoc.format.pdf.document_il.utils.style_helper import BLACK
 from babeldoc.format.pdf.document_il.utils.style_helper import GREEN
 from babeldoc.format.pdf.translation_config import TranslationConfig
 
@@ -143,6 +144,8 @@ class DetectScannedFile:
         for page in docs.page:
             for char in page.pdf_character:
                 char.render_order = None
+                if not char.debug_info:
+                    char.pdf_style.graphic_state = BLACK
 
     def detect_page_is_scanned(
         self, page: il_version_1.Page, pdf: pymupdf.Document, pdf_creater: PDFCreater

{babeldoc-0.5.11 → babeldoc-0.5.12}/babeldoc/format/pdf/document_il/midend/layout_parser.py

@@ -1,4 +1,7 @@
 import logging
+import math
+import os
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 
 import cv2
@@ -7,7 +10,6 @@ from pymupdf import Document
 
 import babeldoc.format.pdf.document_il.utils.extract_char
 from babeldoc.format.pdf.document_il import il_version_1
-from babeldoc.format.pdf.document_il.utils.mupdf_helper import get_no_rotation_img
 from babeldoc.format.pdf.document_il.utils.style_helper import GREEN
 from babeldoc.format.pdf.translation_config import TranslationConfig
 
@@ -120,7 +122,7 @@ class LayoutParser:
         total = len(docs.page)
         with self.translation_config.progress_monitor.stage_start(
             self.stage_name,
-            total,
+            total * 2,
         ) as progress:
             # Process predictions for each page
             for page, layouts in self.model.handle_document(
@@ -134,9 +136,15 @@ class LayoutParser:
                     # Convert coordinate system from picture to il
                     # system to the il coordinate system
                     x0, y0, x1, y1 = layout.xyxy
-                    pix = get_no_rotation_img(mupdf_doc[page.page_number])
+                    # pix = get_no_rotation_img(mupdf_doc[page.page_number])
                     # pix = mupdf_doc[page.page_number].get_pixmap()
-                    h, w = pix.height, pix.width
+                    # h, w = pix.height, pix.width
+                    box = mupdf_doc[page.page_number].mediabox_size
+                    b_h = math.ceil(box.y)
+                    b_w = math.ceil(box.x)
+                    # if b_h != h or b_w != w:
+                    #     logger.warning(f"page {page.page_number} mediabox is not correct, b_h: {b_h}, h: {h}, b_w: {b_w}, w: {w}")
+                    h, w = b_h, b_w
                     x0, y0, x1, y1 = (
                         np.clip(int(x0 - 1), 0, w - 1),
                         np.clip(int(h - y1 - 1), 0, h - 1),
@@ -157,39 +165,47 @@ class LayoutParser:
                     page_layouts.append(page_layout)
 
                 page.page_layout = page_layouts
-                self.generate_fallback_line_layout_for_page(page)
-                self._save_debug_box_to_page(page)
+                # self.generate_fallback_line_layout_for_page(page)
+                # self._save_debug_box_to_page(page)
                 progress.advance(1)
-
+            with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
+                for page in docs.page:
+                    executor.submit(
+                        self.generate_fallback_line_layout_for_page, page, progress
+                    )
         return docs
 
-    def generate_fallback_line_layout_for_page(self, page: il_version_1.Page):
-        exists_page_layouts = page.page_layout
-        char_boxes = babeldoc.format.pdf.document_il.utils.extract_char.convert_page_to_char_boxes(
-            page
-        )
-        if not char_boxes:
-            return
+    def generate_fallback_line_layout_for_page(self, page: il_version_1.Page, progress):
+        try:
+            exists_page_layouts = page.page_layout
+            char_boxes = babeldoc.format.pdf.document_il.utils.extract_char.convert_page_to_char_boxes(
+                page
+            )
+            if not char_boxes:
+                return
 
-        clusters = babeldoc.format.pdf.document_il.utils.extract_char.process_page_chars_to_lines(
-            char_boxes
-        )
-        for cluster in clusters:
-            boxes = [c[0] for c in cluster.chars]
-            min_x = min(b.x for b in boxes)
-            max_x = max(b.x2 for b in boxes)
-            min_y = min(b.y for b in boxes)
-            max_y = max(b.y2 for b in boxes)
-            cluster.chars = il_version_1.Box(min_x, min_y, max_x, max_y)
-            page_layout = il_version_1.PageLayout(
-                id=len(exists_page_layouts) + 1,
-                box=il_version_1.Box(
-                    min_x,
-                    min_y,
-                    max_x,
-                    max_y,
-                ),
-                conf=1,
-                class_name="fallback_line",
+            clusters = babeldoc.format.pdf.document_il.utils.extract_char.process_page_chars_to_lines(
+                char_boxes
             )
-            exists_page_layouts.append(page_layout)
+            for cluster in clusters:
+                boxes = [c[0] for c in cluster.chars]
+                min_x = min(b.x for b in boxes)
+                max_x = max(b.x2 for b in boxes)
+                min_y = min(b.y for b in boxes)
+                max_y = max(b.y2 for b in boxes)
+                cluster.chars = il_version_1.Box(min_x, min_y, max_x, max_y)
+                page_layout = il_version_1.PageLayout(
+                    id=len(exists_page_layouts) + 1,
+                    box=il_version_1.Box(
+                        min_x,
+                        min_y,
+                        max_x,
+                        max_y,
+                    ),
+                    conf=1,
+                    class_name="fallback_line",
+                )
+                exists_page_layouts.append(page_layout)
+            self._save_debug_box_to_page(page)
+        finally:
+            progress.advance(1)

{babeldoc-0.5.11 → babeldoc-0.5.12}/babeldoc/format/pdf/document_il/midend/styles_and_formulas.py

@@ -462,11 +462,17 @@ class StylesAndFormulas:
                 line.pdf_character[i + 1] if i < len(line.pdf_character) - 1 else None
             )
             isspace = char.char_unicode.isspace() if char.char_unicode else False
+            prev_is_space = (
+                previous_char.char_unicode.isspace()
+                if previous_char and previous_char.char_unicode
+                else False
+            )
 
             is_corner_mark = (
                 (
                     previous_char is not None
                     and not isspace
+                    and not prev_is_space
                     and not first_is_bullet
                     # 角标字体，有 0.76 的角标和 0.799 的大写，这里用 0.79 取中，同时考虑首字母放大的情况
                     and char.pdf_style.font_size
@@ -476,6 +482,7 @@ class StylesAndFormulas:
                 or (
                     previous_char is not None
                     and not isspace
+                    and not prev_is_space
                     and not first_is_bullet
                     # 角标字体，有 0.76 的角标和 0.799 的大写，这里用 0.79 取中，同时考虑首字母放大的情况
                     and char.pdf_style.font_size
@@ -487,6 +494,7 @@ class StylesAndFormulas:
                     previous_char is None
                     and next_char is not None
                     and not isspace
+                    and not prev_is_space
                     and not first_is_bullet
                     # 当前字符字体大小明显小于下一个字符，判定为角标
                     and char.pdf_style.font_size < next_char.pdf_style.font_size * 0.79

{babeldoc-0.5.11 → babeldoc-0.5.12}/babeldoc/format/pdf/document_il/utils/extract_char.py

@@ -11,6 +11,7 @@ from sklearn.cluster import DBSCAN
 
 import babeldoc.format.pdf.high_level
 import babeldoc.format.pdf.translation_config
+from babeldoc.const import get_process_pool
 from babeldoc.format.pdf.document_il import il_version_1
 
 logger = logging.getLogger(__name__)
@@ -571,6 +572,15 @@ def _merge_lines_on_page(page_lines: list[Line]) -> list[Line]:
 
 def process_page_chars_to_lines(
     chars: list[tuple[il_version_1.Box, str, bool]],
+) -> list[Line]:
+    pool = get_process_pool()
+    if pool is None:
+        return process_page_chars_to_lines_internal(chars)
+    return pool.apply(process_page_chars_to_lines_internal, (chars,))
+
+
+def process_page_chars_to_lines_internal(
+    chars: list[tuple[il_version_1.Box, str, bool]],
 ) -> list[Line]:
     """
     Process characters on a single page to cluster them into lines.

babeldoc-0.5.12/babeldoc/format/pdf/document_il/utils/mupdf_helper.py

@@ -0,0 +1,42 @@
+import numpy as np
+import pymupdf
+
+from babeldoc.const import get_process_pool
+
+
+def get_no_rotation_img(page: pymupdf.Page, dpi: int = 72) -> pymupdf.Pixmap:
+    # return page.get_pixmap(dpi=72)
+    original_rotation = page.rotation
+    page.set_rotation(0)
+    pix = page.get_pixmap(dpi=dpi)
+    page.set_rotation(original_rotation)
+    return pix
+
+
+def get_no_rotation_img_multiprocess_internal(
+    pdf_bytes: str, pagenum: int, dpi: int = 72
+) -> np.ndarray:
+    # return page.get_pixmap(dpi=72)
+    doc = pymupdf.open(pdf_bytes)
+    try:
+        page = doc[pagenum]
+        original_rotation = page.rotation
+        page.set_rotation(0)
+        pix = page.get_pixmap(dpi=dpi)
+        page.set_rotation(original_rotation)
+        return np.frombuffer(pix.samples, np.uint8).reshape(
+            pix.height,
+            pix.width,
+            3,
+        )[:, :, ::-1]
+    finally:
+        doc.close()
+
+
+def get_no_rotation_img_multiprocess(pdf_bytes: str, pagenum: int, dpi: int = 72):
+    pool = get_process_pool()
+    if pool is None:
+        return get_no_rotation_img_multiprocess_internal(pdf_bytes, pagenum, dpi)
+    return pool.apply(
+        get_no_rotation_img_multiprocess_internal, (pdf_bytes, pagenum, dpi)
+    )

{babeldoc-0.5.11 → babeldoc-0.5.12}/babeldoc/main.py

@@ -1,6 +1,9 @@
 import asyncio
 import logging
+import multiprocessing as mp
 import queue
+import random
+import sys
 from pathlib import Path
 from typing import Any
 
@@ -15,6 +18,7 @@ from rich.progress import TimeRemainingColumn
 
 import babeldoc.assets.assets
 import babeldoc.format.pdf.high_level
+from babeldoc.const import enable_process_pool
 from babeldoc.format.pdf.translation_config import TranslationConfig
 from babeldoc.format.pdf.translation_config import WatermarkOutputMode
 from babeldoc.glossary import Glossary
@@ -22,7 +26,7 @@ from babeldoc.translator.translator import OpenAITranslator
 from babeldoc.translator.translator import set_translate_rate_limiter
 
 logger = logging.getLogger(__name__)
-__version__ = "0.5.11"
+__version__ = "0.5.12"
 
 
 def create_parser():
@@ -103,6 +107,11 @@ def create_parser():
         default=None,
         help="Extra data for metadata",
     )
+    parser.add_argument(
+        "--enable-process-pool",
+        action="store_true",
+        help="DEBUG ONLY",
+    )
     # translation option argument group
     translation_group = parser.add_argument_group(
         "Translation",
@@ -460,6 +469,9 @@ async def main():
     if args.openai and not args.openai_api_key:
         parser.error("使用 OpenAI 服务时必须提供 API key")
 
+    if args.enable_process_pool:
+        enable_process_pool()
+
     # 实例化翻译器
     if args.openai:
         translator = OpenAITranslator(
@@ -688,7 +700,9 @@ async def main():
 
         getattr(doc_layout_model, "init_font_mapper", nop)(config)
         # Create progress handler
-        progress_context, progress_handler = create_progress_handler(config)
+        progress_context, progress_handler = create_progress_handler(
+            config, show_log=False
+        )
 
         # 开始翻译
         with progress_context:
@@ -725,7 +739,9 @@ async def main():
         )
 
 
-def create_progress_handler(translation_config: TranslationConfig):
+def create_progress_handler(
+    translation_config: TranslationConfig, show_log: bool = False
+):
     """Create a progress handler function based on the configuration.
 
     Args:
@@ -748,6 +764,8 @@ def create_progress_handler(translation_config: TranslationConfig):
         stage_tasks = {}
 
         def progress_handler(event):
+            if show_log and random.random() <= 0.1:  # noqa: S311
+                logger.info(event)
             if event["type"] == "progress_start":
                 if event["stage"] not in stage_tasks:
                     stage_tasks[event["stage"]] = progress.add_task(
@@ -876,4 +894,8 @@ def cli():
 
 
 if __name__ == "__main__":
+    if sys.platform == "darwin" or sys.platform == "win32":
+        mp.set_start_method("spawn")
+    else:
+        mp.set_start_method("forkserver")
     cli()

{babeldoc-0.5.11 → babeldoc-0.5.12}/pyproject.toml

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

babeldoc-0.5.11/babeldoc/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.5.11"

babeldoc-0.5.11/babeldoc/format/pdf/document_il/utils/mupdf_helper.py

@@ -1,10 +0,0 @@
-import pymupdf
-
-
-def get_no_rotation_img(page: pymupdf.Page, dpi: int = 72) -> pymupdf.Pixmap:
-    # return page.get_pixmap(dpi=72)
-    original_rotation = page.rotation
-    page.set_rotation(0)
-    pix = page.get_pixmap(dpi=dpi)
-    page.set_rotation(original_rotation)
-    return pix