BabelDOC 0.1.26__tar.gz → 0.1.27__tar.gz

{babeldoc-0.1.26 → babeldoc-0.1.27}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: BabelDOC
-Version: 0.1.26
+Version: 0.1.27
 Summary: Yet Another Document Translator
 Project-URL: Homepage, https://github.com/funstory-ai/BabelDOC
 Project-URL: Issues, https://github.com/funstory-ai/BabelDOC/issues
@@ -65,7 +65,8 @@ Description-Content-Type: text/markdown
 
 PDF scientific paper translation and bilingual comparison library.
 
-- Beta version launched [Immersive Translate - BabelDOC](https://app.immersivetranslate.com/babel-doc/) 1000 free pages per month.
+- **Online Service**: Beta version launched [Immersive Translate - BabelDOC](https://app.immersivetranslate.com/babel-doc/) 1000 free pages per month.
+- **Self-deployment**: [PDFMathTranslate](https://github.com/Byaidu/PDFMathTranslate) 1.9.3+ Experimental support for BabelDOC, available for self-deployment + WebUI with more translation services.
 - Provides a simple [command line interface](#getting-started).
 - Provides a [Python API](#python-api).
 - Mainly designed to be embedded into other programs, but can also be used directly for simple translation tasks.

{babeldoc-0.1.26 → babeldoc-0.1.27}/README.md

@@ -30,7 +30,8 @@
 
 PDF scientific paper translation and bilingual comparison library.
 
-- Beta version launched [Immersive Translate - BabelDOC](https://app.immersivetranslate.com/babel-doc/) 1000 free pages per month.
+- **Online Service**: Beta version launched [Immersive Translate - BabelDOC](https://app.immersivetranslate.com/babel-doc/) 1000 free pages per month.
+- **Self-deployment**: [PDFMathTranslate](https://github.com/Byaidu/PDFMathTranslate) 1.9.3+ Experimental support for BabelDOC, available for self-deployment + WebUI with more translation services.
 - Provides a simple [command line interface](#getting-started).
 - Provides a [Python API](#python-api).
 - Mainly designed to be embedded into other programs, but can also be used directly for simple translation tasks.

babeldoc-0.1.27/babeldoc/__init__.py

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

{babeldoc-0.1.26 → babeldoc-0.1.27}/babeldoc/const.py

@@ -2,7 +2,7 @@ import shutil
 import subprocess
 from pathlib import Path
 
-__version__ = "0.1.26"
+__version__ = "0.1.27"
 
 CACHE_FOLDER = Path.home() / ".cache" / "babeldoc"
 

{babeldoc-0.1.26 → babeldoc-0.1.27}/babeldoc/document_il/midend/il_translator.py

@@ -1,6 +1,7 @@
 import concurrent.futures
 import json
 import logging
+import re
 from pathlib import Path
 
 from tqdm import tqdm
@@ -555,6 +556,7 @@ class ILTranslator:
                     return
 
                 translated_text = self.translate_engine.translate(text)
+                translated_text = re.sub(r"[. 。…]{20,}", ".", translated_text)
 
                 tracker.set_output(translated_text)
 

{babeldoc-0.1.26 → babeldoc-0.1.27}/babeldoc/document_il/midend/layout_parser.py

@@ -110,71 +110,43 @@ class LayoutParser:
     def process(self, docs: il_version_1.Document, mupdf_doc: Document):
         """Generate layouts for all pages that need to be translated."""
         # Get pages that need to be translated
-        pages_to_translate = [
-            page
-            for page in docs.page
-            if self.translation_config.should_translate_page(page.page_number + 1)
-        ]
-        total = len(pages_to_translate)
+        total = len(docs.page)
         with self.translation_config.progress_monitor.stage_start(
             self.stage_name,
             total,
         ) as progress:
-            # Process pages in batches
-            batch_size = 1
-            for i in range(0, total, batch_size):
-                self.translation_config.raise_if_cancelled()
-                batch_pages = pages_to_translate[i : i + batch_size]
-
-                # Prepare batch images
-                batch_images = []
-                for page in batch_pages:
-                    pix = mupdf_doc[page.page_number].get_pixmap(dpi=72)
-                    image = np.fromstring(pix.samples, np.uint8).reshape(
-                        pix.height,
-                        pix.width,
-                        3,
-                    )[:, :, ::-1]
-                    batch_images.append(image)
-
-                # Get predictions for the batch
-                layouts_batch = self.model.predict(batch_images, batch_size=batch_size)
-
-                # Process predictions for each page
-                for page, layouts in zip(batch_pages, layouts_batch, strict=False):
-                    page_layouts = []
-                    self._save_debug_image(
-                        batch_images[batch_pages.index(page)],
-                        layouts,
-                        page.page_number + 1,
+            # Process predictions for each page
+            for page, layouts in self.model.handle_document(
+                docs.page, mupdf_doc, self.translation_config, self._save_debug_image
+            ):
+                page_layouts = []
+                for layout in layouts.boxes:
+                    # Convert coordinate system from picture to il
+                    # system to the il coordinate system
+                    x0, y0, x1, y1 = layout.xyxy
+                    pix = mupdf_doc[page.page_number].get_pixmap()
+                    h, w = pix.height, pix.width
+                    x0, y0, x1, y1 = (
+                        np.clip(int(x0 - 1), 0, w - 1),
+                        np.clip(int(h - y1 - 1), 0, h - 1),
+                        np.clip(int(x1 + 1), 0, w - 1),
+                        np.clip(int(h - y0 + 1), 0, h - 1),
                     )
-                    for layout in layouts.boxes:
-                        # Convert coordinate system from picture to il
-                        # system to the il coordinate system
-                        x0, y0, x1, y1 = layout.xyxy
-                        pix = mupdf_doc[page.page_number].get_pixmap()
-                        h, w = pix.height, pix.width
-                        x0, y0, x1, y1 = (
-                            np.clip(int(x0 - 1), 0, w - 1),
-                            np.clip(int(h - y1 - 1), 0, h - 1),
-                            np.clip(int(x1 + 1), 0, w - 1),
-                            np.clip(int(h - y0 + 1), 0, h - 1),
-                        )
-                        page_layout = il_version_1.PageLayout(
-                            id=len(page_layouts) + 1,
-                            box=il_version_1.Box(
-                                x0.item(),
-                                y0.item(),
-                                x1.item(),
-                                y1.item(),
-                            ),
-                            conf=layout.conf.item(),
-                            class_name=layouts.names[layout.cls],
-                        )
-                        page_layouts.append(page_layout)
+                    page_layout = il_version_1.PageLayout(
+                        id=len(page_layouts) + 1,
+                        box=il_version_1.Box(
+                            x0.item(),
+                            y0.item(),
+                            x1.item(),
+                            y1.item(),
+                        ),
+                        conf=layout.conf.item(),
+                        class_name=layouts.names[layout.cls],
+                    )
+                    page_layouts.append(page_layout)
 
-                    page.page_layout = page_layouts
-                    self._save_debug_box_to_page(page)
-                    progress.advance(1)
+                page.page_layout = page_layouts
+                self._save_debug_box_to_page(page)
+                progress.advance(1)
 
-            return docs
+        return docs

{babeldoc-0.1.26 → babeldoc-0.1.27}/babeldoc/document_il/midend/paragraph_finder.py

@@ -244,6 +244,7 @@ class ParagraphFinder:
     def is_text_layout(self, layout: Layout):
         return layout is not None and layout.name in [
             "plain text",
+            "tiny text",
             "title",
             "abandon",
             "figure_caption",
@@ -312,6 +313,7 @@ class ParagraphFinder:
             "figure_caption",
             "abandon",
             "plain text",
+            "tiny text",
             "title",
         ]
         char_box = char.box

{babeldoc-0.1.26 → babeldoc-0.1.27}/babeldoc/document_il/translator/translator.py

@@ -8,6 +8,7 @@ import unicodedata
 from abc import ABC
 from abc import abstractmethod
 
+import httpx
 import openai
 import requests
 
@@ -271,3 +272,52 @@ class OpenAITranslator(BaseTranslator):
 
     def get_rich_text_right_placeholder(self, placeholder_id: int):
         return self.get_formular_placeholder(placeholder_id + 1)
+
+
+class TranslateTranslator(BaseTranslator):
+    # https://github.com/openai/openai-python
+    name = "openai"
+
+    def __init__(
+        self,
+        lang_in,
+        lang_out,
+        url=None,
+        ignore_cache=False,
+    ):
+        super().__init__(lang_in, lang_out, ignore_cache)
+        self.client = httpx.Client()
+        self.url = url
+
+    def do_translate(self, text) -> str:
+        response = self.client.post(
+            self.url,
+            json={
+                "text": [text],
+                "src": "Englih",
+                "tgt": "Simplifed Chinese",
+            },
+        )
+        return response.json()["text"][0]
+
+    def prompt(self, text):
+        return [
+            {
+                "role": "system",
+                "content": "You are a professional,authentic machine translation engine.",
+            },
+            {
+                "role": "user",
+                "content": f";; Treat next line as plain text input and translate it into {self.lang_out}, output translation ONLY. If translation is unnecessary (e.g. proper nouns, codes, {'{{1}}, etc. '}), return the original text. NO explanations. NO notes. Input:\n\n{text}",
+            },
+        ]
+
+    def get_formular_placeholder(self, placeholder_id: int):
+        return "{{v" + str(placeholder_id) + "}}"
+        return "{{" + str(placeholder_id) + "}}"
+
+    def get_rich_text_left_placeholder(self, placeholder_id: int):
+        return self.get_formular_placeholder(placeholder_id)
+
+    def get_rich_text_right_placeholder(self, placeholder_id: int):
+        return self.get_formular_placeholder(placeholder_id + 1)

{babeldoc-0.1.26 → babeldoc-0.1.27}/babeldoc/docvision/doclayout.py

@@ -2,12 +2,15 @@ import abc
 import ast
 import logging
 import platform
+from collections.abc import Generator
 
 import cv2
 import numpy as np
 import onnx
 import onnxruntime
+import pymupdf
 
+import babeldoc.document_il.il_version_1
 from babeldoc.assets.assets import get_doclayout_onnx_model_path
 
 # from huggingface_hub import hf_hub_download
@@ -15,6 +18,19 @@ from babeldoc.assets.assets import get_doclayout_onnx_model_path
 logger = logging.getLogger(__name__)
 
 
+class YoloResult:
+    """Helper class to store detection results from ONNX model."""
+
+    def __init__(self, names, boxes=None, boxes_data=None):
+        if boxes is not None:
+            self.boxes = boxes
+        else:
+            assert boxes_data is not None
+            self.boxes = [YoloBox(data=d) for d in boxes_data]
+        self.boxes.sort(key=lambda x: x.conf, reverse=True)
+        self.names = names
+
+
 class DocLayoutModel(abc.ABC):
     @staticmethod
     def load_onnx():
@@ -42,18 +58,19 @@ class DocLayoutModel(abc.ABC):
             **kwargs: Additional arguments.
         """
 
-
-class YoloResult:
-    """Helper class to store detection results from ONNX model."""
-
-    def __init__(self, names, boxes=None, boxes_data=None):
-        if boxes is not None:
-            self.boxes = boxes
-        else:
-            assert boxes_data is not None
-            self.boxes = [YoloBox(data=d) for d in boxes_data]
-        self.boxes.sort(key=lambda x: x.conf, reverse=True)
-        self.names = names
+    @abc.abstractmethod
+    def handle_document(
+        self,
+        pages: list[babeldoc.document_il.il_version_1.Page],
+        mupdf_doc: pymupdf.Document,
+        translate_config,
+        save_debug_image,
+    ) -> Generator[
+        tuple[babeldoc.document_il.il_version_1.Page, YoloResult], None, None
+    ]:
+        """
+        Handle a document.
+        """
 
 
 class YoloBox:
@@ -254,3 +271,28 @@ class OnnxModel(DocLayoutModel):
                 results.append(YoloResult(boxes_data=preds, names=self._names))
 
         return results
+
+    def handle_document(
+        self,
+        pages: list[babeldoc.document_il.il_version_1.Page],
+        mupdf_doc: pymupdf.Document,
+        translate_config,
+        save_debug_image,
+    ) -> Generator[
+        tuple[babeldoc.document_il.il_version_1.Page, YoloResult], None, None
+    ]:
+        for page in pages:
+            translate_config.raise_if_cancelled()
+            pix = mupdf_doc[page.page_number].get_pixmap(dpi=72)
+            image = np.fromstring(pix.samples, np.uint8).reshape(
+                pix.height,
+                pix.width,
+                3,
+            )[:, :, ::-1]
+            predict_result = self.predict(image)[0]
+            save_debug_image(
+                image,
+                predict_result,
+                page.page_number + 1,
+            )
+            yield page, predict_result

babeldoc-0.1.27/babeldoc/docvision/rpc_doclayout.py

@@ -0,0 +1,304 @@
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import cv2
+import httpx
+import msgpack
+import numpy as np
+import pymupdf
+from tenacity import retry
+from tenacity import retry_if_exception_type
+from tenacity import stop_after_attempt
+from tenacity import wait_exponential
+
+import babeldoc
+from babeldoc.docvision.doclayout import DocLayoutModel
+from babeldoc.docvision.doclayout import YoloBox
+from babeldoc.docvision.doclayout import YoloResult
+
+logger = logging.getLogger(__name__)
+
+
+def encode_image(image) -> bytes:
+    """Read and encode image to bytes
+
+    Args:
+        image: Can be either a file path (str) or numpy array
+    """
+    if isinstance(image, str):
+        if not Path(image).exists():
+            raise FileNotFoundError(f"Image file not found: {image}")
+        img = cv2.imread(image)
+        if img is None:
+            raise ValueError(f"Failed to read image: {image}")
+    else:
+        img = image
+
+    # logger.debug(f"Image shape: {img.shape}")
+    encoded = cv2.imencode(".jpg", img)[1].tobytes()
+    # logger.debug(f"Encoded image size: {len(encoded)} bytes")
+    return encoded
+
+
+@retry(
+    stop=stop_after_attempt(3),  # 最多重试3次
+    wait=wait_exponential(
+        multiplier=1, min=1, max=10
+    ),  # 指数退避策略，初始1秒，最大10秒
+    retry=retry_if_exception_type((httpx.HTTPError, Exception)),  # 针对哪些异常重试
+    before_sleep=lambda retry_state: logger.warning(
+        f"Request failed, retrying in {retry_state.next_action.sleep} seconds... "
+        f"(Attempt {retry_state.attempt_number}/3)"
+    ),
+)
+def predict_layout(
+    image,
+    host: str = "http://localhost:8000",
+    imgsz: int = 1024,
+):
+    """
+    Predict document layout using the MOSEC service
+
+    Args:
+        image: Can be either a file path (str) or numpy array
+        host: Service host URL
+        imgsz: Image size for model input
+
+    Returns:
+        List of predictions containing bounding boxes and classes
+    """
+    # Prepare request data
+    if not isinstance(image, list):
+        image = [image]
+    image_data = [encode_image(image) for image in image]
+    data = {
+        "image": image_data,
+        "imgsz": imgsz,
+    }
+
+    # Pack data using msgpack
+    packed_data = msgpack.packb(data, use_bin_type=True)
+    logger.debug(f"Packed data size: {len(packed_data)} bytes")
+
+    # Send request
+    logger.debug(f"Sending request to {host}/inference")
+    response = httpx.post(
+        f"{host}/inference",
+        data=packed_data,
+        headers={
+            "Content-Type": "application/msgpack",
+            "Accept": "application/msgpack",
+        },
+        timeout=300,
+        follow_redirects=True,
+    )
+
+    logger.debug(f"Response status: {response.status_code}")
+    logger.debug(f"Response headers: {response.headers}")
+
+    if response.status_code == 200:
+        try:
+            result = msgpack.unpackb(response.content, raw=False)
+            return result
+        except Exception as e:
+            logger.exception(f"Failed to unpack response: {e!s}")
+            raise
+    else:
+        logger.error(f"Request failed with status {response.status_code}")
+        logger.error(f"Response content: {response.content}")
+        raise Exception(
+            f"Request failed with status {response.status_code}: {response.text}",
+        )
+
+
+class ResultContainer:
+    def __init__(self):
+        self.result = YoloResult(boxes_data=np.array([]), names=[])
+
+
+class RpcDocLayoutModel(DocLayoutModel):
+    """DocLayoutModel implementation that uses RPC service."""
+
+    def __init__(self, host: str = "http://localhost:8000"):
+        """Initialize RPC model with host address."""
+        self.host = host
+        self._stride = 32  # Default stride value
+        self._names = ["text", "title", "list", "table", "figure"]
+
+    @property
+    def stride(self) -> int:
+        """Stride of the model input."""
+        return self._stride
+
+    def resize_and_pad_image(self, image, new_shape):
+        """
+        Resize and pad the image to the specified size,
+        ensuring dimensions are multiples of stride.
+
+        Parameters:
+        - image: Input image
+        - new_shape: Target size (integer or (height, width) tuple)
+        - stride: Padding alignment stride, default 32
+
+        Returns:
+        - Processed image
+        """
+        if isinstance(new_shape, int):
+            new_shape = (new_shape, new_shape)
+
+        h, w = image.shape[:2]
+        new_h, new_w = new_shape
+
+        # Calculate scaling ratio
+        r = min(new_h / h, new_w / w)
+        resized_h, resized_w = int(round(h * r)), int(round(w * r))
+
+        # Resize image
+        image = cv2.resize(
+            image, (resized_w, resized_h), interpolation=cv2.INTER_LINEAR
+        )
+
+        # Calculate padding size
+        pad_h = new_h - resized_h
+        pad_w = new_w - resized_w
+        top, bottom = pad_h // 2, pad_h - pad_h // 2
+        left, right = pad_w // 2, pad_w - pad_w // 2
+
+        # Add padding
+        image = cv2.copyMakeBorder(
+            image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114)
+        )
+
+        return image
+
+    def scale_boxes(self, img1_shape, boxes, img0_shape):
+        """
+        Rescales bounding boxes (in the format of xyxy by default) from the shape of the image they were originally
+        specified in (img1_shape) to the shape of a different image (img0_shape).
+
+        Args:
+            img1_shape (tuple): The shape of the image that the bounding boxes are for,
+                in the format of (height, width).
+            boxes (torch.Tensor): the bounding boxes of the objects in the image, in the format of (x1, y1, x2, y2)
+            img0_shape (tuple): the shape of the target image, in the format of (height, width).
+
+        Returns:
+            boxes (torch.Tensor): The scaled bounding boxes, in the format of (x1, y1, x2, y2)
+        """
+
+        # Calculate scaling ratio
+        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])
+
+        # Calculate padding size
+        pad_x = round((img1_shape[1] - img0_shape[1] * gain) / 2 - 0.1)
+        pad_y = round((img1_shape[0] - img0_shape[0] * gain) / 2 - 0.1)
+
+        # Remove padding and scale boxes
+        boxes = (boxes - [pad_x, pad_y, pad_x, pad_y]) / gain
+        return boxes
+
+    def predict_image(
+        self,
+        image,
+        host: str = None,
+        result_container: ResultContainer | None = None,
+        imgsz: int = 1024,
+    ) -> ResultContainer:
+        """Predict the layout of document pages using RPC service."""
+        if result_container is None:
+            result_container = ResultContainer()
+        target_imgsz = (800, 800)
+        orig_h, orig_w = image.shape[:2]
+        if image.shape[0] != target_imgsz[0] or image.shape[1] != target_imgsz[1]:
+            image = self.resize_and_pad_image(image, new_shape=target_imgsz)
+        preds = predict_layout([image], host=self.host, imgsz=800)
+
+        if len(preds) > 0:
+            for pred in preds:
+                boxes = [
+                    YoloBox(
+                        None,
+                        self.scale_boxes(
+                            (800, 800), np.array(x["xyxy"]), (orig_h, orig_w)
+                        ),
+                        np.array(x["conf"]),
+                        x["cls"],
+                    )
+                    for x in pred["boxes"]
+                ]
+                result_container.result = YoloResult(
+                    boxes=boxes,
+                    names={int(k): v for k, v in pred["names"].items()},
+                )
+        return result_container.result
+
+    def predict(self, image, imgsz=1024, **kwargs) -> list[YoloResult]:
+        """Predict the layout of document pages using RPC service."""
+        # Handle single image input
+        if isinstance(image, np.ndarray) and len(image.shape) == 3:
+            image = [image]
+
+        result_containers = [ResultContainer() for _ in image]
+        predict_thread = ThreadPoolExecutor(max_workers=len(image))
+        for img, result_container in zip(image, result_containers, strict=True):
+            predict_thread.submit(
+                self.predict_image, img, self.host, result_container, 800
+            )
+        predict_thread.shutdown(wait=True)
+        result = [result_container.result for result_container in result_containers]
+        return result
+
+    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)
+        image = np.fromstring(pix.samples, np.uint8).reshape(
+            pix.height,
+            pix.width,
+            3,
+        )[:, :, ::-1]
+        predict_result = self.predict_image(image, self.host, None, 800)
+        save_debug_image(image, predict_result, page.page_number + 1)
+        return page, predict_result
+
+    def handle_document(
+        self,
+        pages: list[babeldoc.document_il.il_version_1.Page],
+        mupdf_doc: pymupdf.Document,
+        translate_config,
+        save_debug_image,
+    ):
+        with ThreadPoolExecutor(max_workers=16) as executor:
+            yield from executor.map(
+                self.predict_page,
+                pages,
+                (mupdf_doc for _ in range(len(pages))),
+                (translate_config for _ in range(len(pages))),
+                (save_debug_image for _ in range(len(pages))),
+            )
+
+    @staticmethod
+    def from_host(host: str) -> "RpcDocLayoutModel":
+        """Create RpcDocLayoutModel from host address."""
+        return RpcDocLayoutModel(host=host)
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.DEBUG)
+    # Test the service
+    try:
+        # Use a default test image if example/1.png doesn't exist
+        image_path = "example/1.png"
+        if not Path(image_path).exists():
+            print(f"Warning: {image_path} not found.")
+            print("Please provide the path to a test image:")
+            image_path = input("> ")
+
+        logger.info(f"Processing image: {image_path}")
+        result = predict_layout(image_path)
+        print("Prediction results:")
+        print(result)
+    except Exception as e:
+        print(f"Error: {e!s}")

{babeldoc-0.1.26 → babeldoc-0.1.27}/babeldoc/main.py

@@ -17,6 +17,7 @@ import babeldoc.high_level
 from babeldoc.document_il.translator.translator import BingTranslator
 from babeldoc.document_il.translator.translator import GoogleTranslator
 from babeldoc.document_il.translator.translator import OpenAITranslator
+from babeldoc.document_il.translator.translator import TranslateTranslator
 from babeldoc.document_il.translator.translator import set_translate_rate_limiter
 from babeldoc.docvision.doclayout import DocLayoutModel
 from babeldoc.docvision.rpc_doclayout import RpcDocLayoutModel
@@ -24,7 +25,7 @@ from babeldoc.translation_config import TranslationConfig
 from babeldoc.translation_config import WatermarkOutputMode
 
 logger = logging.getLogger(__name__)
-__version__ = "0.1.26"
+__version__ = "0.1.27"
 
 
 def create_parser():
@@ -200,6 +201,11 @@ def create_parser():
         action="store_true",
         help="Use Google translator.",
     )
+    service_group.add_argument(
+        "--translate",
+        action="store_true",
+        help="Use translate translator.",
+    )
     service_group.add_argument(
         "--bing",
         action="store_true",
@@ -223,6 +229,14 @@ def create_parser():
         "-k",
         help="The API key for the OpenAI API.",
     )
+    service_group = parser.add_argument_group(
+        "Translation - Translate Options",
+        description="Translate specific options",
+    )
+    service_group.add_argument(
+        "--translate-url",
+        help="The base URL for the Translation API.",
+    )
 
     return parser
 
@@ -254,8 +268,8 @@ async def main():
         return
 
     # 验证翻译服务选择
-    if not (args.openai or args.google or args.bing):
-        parser.error("必须选择一个翻译服务：--openai、--google 或 --bing")
+    if not (args.openai or args.google or args.bing or args.translate):
+        parser.error("必须选择一个翻译服务：--openai、--google、--bing 或 --translate")
 
     # 验证 OpenAI 参数
     if args.openai and not args.openai_api_key:
@@ -277,6 +291,13 @@ async def main():
             lang_out=args.lang_out,
             ignore_cache=args.ignore_cache,
         )
+    elif args.translate:
+        translator = TranslateTranslator(
+            lang_in=args.lang_in,
+            lang_out=args.lang_out,
+            ignore_cache=args.ignore_cache,
+            url=args.translate_url,
+        )
     else:
         translator = GoogleTranslator(
             lang_in=args.lang_in,
@@ -321,6 +342,16 @@ async def main():
     else:
         args.output = None
 
+    watermark_output_mode = WatermarkOutputMode.Watermarked
+    if args.no_watermark:
+        watermark_output_mode = WatermarkOutputMode.NoWatermark
+    elif args.watermark_output_mode == "both":
+        watermark_output_mode = WatermarkOutputMode.Both
+    elif args.watermark_output_mode == "watermarked":
+        watermark_output_mode = WatermarkOutputMode.Watermarked
+    elif args.watermark_output_mode == "no_watermark":
+        watermark_output_mode = WatermarkOutputMode.NoWatermark
+
     for file in pending_files:
         # 清理文件路径，去除两端的引号
         file = file.strip("\"'")
@@ -349,11 +380,7 @@ async def main():
             use_alternating_pages_dual=args.use_alternating_pages_dual,
             report_interval=args.report_interval,
             min_text_length=args.min_text_length,
-            watermark_output_mode=(
-                WatermarkOutputMode.NoWatermark
-                if args.no_watermark
-                else getattr(WatermarkOutputMode, args.watermark_output_mode.title())
-            ),
+            watermark_output_mode=watermark_output_mode,
         )
 
         # Create progress handler

{babeldoc-0.1.26 → babeldoc-0.1.27}/pyproject.toml

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

{babeldoc-0.1.26 → babeldoc-0.1.27}/uv.lock

@@ -48,7 +48,7 @@ wheels = [
 
 [[package]]
 name = "babeldoc"
-version = "0.1.26"
+version = "0.1.27"
 source = { editable = "." }
 dependencies = [
     { name = "bitstring" },

babeldoc-0.1.26/babeldoc/__init__.py

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

babeldoc-0.1.26/babeldoc/docvision/rpc_doclayout.py

@@ -1,176 +0,0 @@
-import logging
-from pathlib import Path
-
-import cv2
-import httpx
-import msgpack
-import numpy as np
-from tenacity import retry
-from tenacity import retry_if_exception_type
-from tenacity import stop_after_attempt
-from tenacity import wait_exponential
-
-from babeldoc.docvision.doclayout import DocLayoutModel
-from babeldoc.docvision.doclayout import YoloBox
-from babeldoc.docvision.doclayout import YoloResult
-
-logger = logging.getLogger(__name__)
-
-
-def encode_image(image) -> bytes:
-    """Read and encode image to bytes
-
-    Args:
-        image: Can be either a file path (str) or numpy array
-    """
-    if isinstance(image, str):
-        if not Path(image).exists():
-            raise FileNotFoundError(f"Image file not found: {image}")
-        img = cv2.imread(image)
-        if img is None:
-            raise ValueError(f"Failed to read image: {image}")
-    else:
-        img = image
-
-    # logger.debug(f"Image shape: {img.shape}")
-    encoded = cv2.imencode(".jpg", img)[1].tobytes()
-    # logger.debug(f"Encoded image size: {len(encoded)} bytes")
-    return encoded
-
-
-@retry(
-    stop=stop_after_attempt(3),  # 最多重试3次
-    wait=wait_exponential(
-        multiplier=1, min=1, max=10
-    ),  # 指数退避策略，初始1秒，最大10秒
-    retry=retry_if_exception_type((httpx.HTTPError, Exception)),  # 针对哪些异常重试
-    before_sleep=lambda retry_state: logger.warning(
-        f"Request failed, retrying in {retry_state.next_action.sleep} seconds... "
-        f"(Attempt {retry_state.attempt_number}/3)"
-    ),
-)
-def predict_layout(
-    image,
-    host: str = "http://localhost:8000",
-    imgsz: int = 1024,
-):
-    """
-    Predict document layout using the MOSEC service
-
-    Args:
-        image: Can be either a file path (str) or numpy array
-        host: Service host URL
-        imgsz: Image size for model input
-
-    Returns:
-        List of predictions containing bounding boxes and classes
-    """
-    # Prepare request data
-    if not isinstance(image, list):
-        image = [image]
-    image_data = [encode_image(image) for image in image]
-    data = {
-        "image": image_data,
-        "imgsz": imgsz,
-    }
-
-    # Pack data using msgpack
-    packed_data = msgpack.packb(data, use_bin_type=True)
-    logger.debug(f"Packed data size: {len(packed_data)} bytes")
-
-    # Send request
-    logger.debug(f"Sending request to {host}/inference")
-    response = httpx.post(
-        f"{host}/inference",
-        data=packed_data,
-        headers={
-            "Content-Type": "application/msgpack",
-            "Accept": "application/msgpack",
-        },
-        timeout=300,
-        follow_redirects=True,
-    )
-
-    logger.debug(f"Response status: {response.status_code}")
-    logger.debug(f"Response headers: {response.headers}")
-
-    if response.status_code == 200:
-        try:
-            result = msgpack.unpackb(response.content, raw=False)
-            return result
-        except Exception as e:
-            logger.exception(f"Failed to unpack response: {e!s}")
-            raise
-    else:
-        logger.error(f"Request failed with status {response.status_code}")
-        logger.error(f"Response content: {response.content}")
-        raise Exception(
-            f"Request failed with status {response.status_code}: {response.text}",
-        )
-
-
-class RpcDocLayoutModel(DocLayoutModel):
-    """DocLayoutModel implementation that uses RPC service."""
-
-    def __init__(self, host: str = "http://localhost:8000"):
-        """Initialize RPC model with host address."""
-        self.host = host
-        self._stride = 32  # Default stride value
-        self._names = ["text", "title", "list", "table", "figure"]
-
-    @property
-    def stride(self) -> int:
-        """Stride of the model input."""
-        return self._stride
-
-    def predict(self, image, imgsz=1024, **kwargs) -> list[YoloResult]:
-        """Predict the layout of document pages using RPC service."""
-        # Handle single image input
-        if isinstance(image, np.ndarray) and len(image.shape) == 3:
-            image = [image]
-
-        results = []
-        for img in image:
-            preds = predict_layout([img], host=self.host, imgsz=imgsz)
-            if len(preds) > 0:
-                for pred in preds:
-                    boxes = [
-                        YoloBox(
-                            None, np.array(x["xyxy"]), np.array(x["conf"]), x["cls"]
-                        )
-                        for x in pred["boxes"]
-                    ]
-                    results.append(
-                        YoloResult(
-                            boxes=boxes,
-                            names={int(k): v for k, v in pred["names"].items()},
-                        ),
-                    )
-            else:
-                results.append(YoloResult(boxes_data=np.array([]), names=[]))
-
-        return results
-
-    @staticmethod
-    def from_host(host: str) -> "RpcDocLayoutModel":
-        """Create RpcDocLayoutModel from host address."""
-        return RpcDocLayoutModel(host=host)
-
-
-if __name__ == "__main__":
-    logging.basicConfig(level=logging.DEBUG)
-    # Test the service
-    try:
-        # Use a default test image if example/1.png doesn't exist
-        image_path = "example/1.png"
-        if not Path(image_path).exists():
-            print(f"Warning: {image_path} not found.")
-            print("Please provide the path to a test image:")
-            image_path = input("> ")
-
-        logger.info(f"Processing image: {image_path}")
-        result = predict_layout(image_path)
-        print("Prediction results:")
-        print(result)
-    except Exception as e:
-        print(f"Error: {e!s}")