doc-page-extractor 0.1.1__py3-none-any.whl → 1.0.2__py3-none-any.whl

doc_page_extractor/redacter.py

@@ -0,0 +1,111 @@
+from typing import Any, Generator, Iterable, cast
+
+from PIL import Image, ImageDraw
+
+
+def redact(
+    image: Image.Image,
+    fill_color: tuple[int, int, int],
+    rectangles: Iterable[tuple[int, int, int, int]],
+) -> Image.Image:
+    draw = ImageDraw.Draw(image)
+    for x1, y1, x2, y2 in rectangles:
+        draw.rectangle((x1, y1, x2, y2), fill=fill_color)
+    return image
+
+
+class _AveragingColor:
+    def __init__(self) -> None:
+        self._r: float = 0.0
+        self._g: float = 0.0
+        self._b: float = 0.0
+        self._a: float = 0.0
+        self._count: int = 0
+
+    @property
+    def count(self) -> int:
+        return self._count
+
+    @property
+    def average(self) -> tuple[float, float, float, float]:
+        if self._count == 0:
+            return 1.0, 1.0, 1.0, 1.0
+        return (
+            self._r / self._count,
+            self._g / self._count,
+            self._b / self._count,
+            self._a / self._count,
+        )
+
+    def add_color(self, r: float, g: float, b: float, a: float) -> None:
+        self._r += r
+        self._g += g
+        self._b += b
+        self._a += a
+        self._count += 1
+
+
+def background_color(image: Image.Image) -> tuple[int, int, int]:
+    """将像素颜色按灰度排序，取中位颜色。此颜色与纸张的颜色相同，可做背景色"""
+    pixels_count = image.width * image.height
+    if pixels_count == 0:
+        return 255, 255, 255
+
+    bucket: list[_AveragingColor | None] = [None] * 256
+    for r, g, b, a in _iter_pixels(image):
+        gray = round(255 * _gray(r, g, b, a))
+        colors = bucket[gray]
+        if colors is None:
+            colors = _AveragingColor()
+            bucket[gray] = colors
+        colors.add_color(r, g, b, a)
+
+    offset: int = 0
+    found_colors: _AveragingColor | None = None
+
+    for colors in bucket:
+        if not colors:
+            continue
+        offset += colors.count
+        if offset > pixels_count // 2:
+            found_colors = colors
+            break
+
+    assert found_colors is not None
+    r, g, b, a = found_colors.average
+
+    # 背景色为白色
+    r = r * a + 1.0 * (1.0 - a)
+    g = g * a + 1.0 * (1.0 - a)
+    b = b * a + 1.0 * (1.0 - a)
+
+    return round(r * 255), round(g * 255), round(b * 255)
+
+
+def _gray(r: float, g: float, b: float, a: float) -> float:
+    # ITU-R BT.601 https://en.wikipedia.org/wiki/Rec._601
+    gray = 0.299 * r + 0.587 * g + 0.114 * b
+    return gray * a
+
+
+def _iter_pixels(
+    image: Image.Image,
+) -> Generator[tuple[float, float, float, float], None, None]:
+    for pixel in cast(Any, image.getdata()):
+        pixel_len = len(cast(tuple, pixel)) if isinstance(pixel, tuple) else 1
+        if pixel_len == 4:
+            # RGBA 格式
+            r, g, b, a = cast(tuple[int, int, int, int], pixel)
+        elif pixel_len == 3:
+            # RGB 格式
+            r, g, b = cast(tuple[int, int, int], pixel)
+            a = 255
+        elif pixel_len == 2:
+            # LA 格式 (灰度 + alpha)
+            l, a = cast(tuple[int, int], pixel)
+            r = g = b = l
+        else:
+            # L 格式 (灰度)
+            r = g = b = cast(int, pixel)
+            a = 255
+        yield (r / 255.0, g / 255.0, b / 255.0, a / 255.0)

doc_page_extractor-1.0.2.dist-info/METADATA

@@ -0,0 +1,120 @@
+Metadata-Version: 2.4
+Name: doc-page-extractor
+Version: 1.0.2
+Summary: Document page extraction tool powered by DeepSeek-OCR
+License: MIT
+License-File: LICENSE
+Author: Tao Zeyu
+Author-email: i@taozeyu.com
+Maintainer: Tao Zeyu
+Maintainer-email: i@taozeyu.com
+Requires-Python: >=3.10,<3.14
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: addict (>=2.4.0)
+Requires-Dist: easydict (>=1.13)
+Requires-Dist: einops (>=0.8.0)
+Requires-Dist: transformers (>=4.46.0,<4.48.0)
+Project-URL: Repository, https://github.com/moskize91/doc-page-extractor
+Description-Content-Type: text/markdown
+
+# doc-page-extractor
+
+Document page extraction tool powered by DeepSeek-OCR.
+
+## Installation
+
+> **⚠️ Important:** This package requires PyTorch with CUDA support (GPU Required). PyTorch is NOT automatically installed - you must install it manually first.
+
+### Step 1: Install PyTorch with CUDA
+
+Choose the command that matches your CUDA version:
+
+```bash
+# For CUDA 12.1 (recommended for most users)
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
+
+# For CUDA 11.8
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
+
+# For CUDA 12.6
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cu126
+```
+
+> 💡 **Don't know your CUDA version?** Run `nvidia-smi` to check, or just try CUDA 12.1 (works with most recent drivers).
+
+### Step 2: Install doc-page-extractor
+
+```bash
+pip install doc-page-extractor
+```
+
+### Verify Installation
+
+Check if everything is working:
+
+```bash
+python -c "import doc_page_extractor; import torch; print('✓ Installation successful!'); print('✓ CUDA available:', torch.cuda.is_available())"
+```
+
+Expected output:
+```
+✓ Installation successful!
+✓ CUDA available: True
+```
+
+If CUDA shows `False`, see the troubleshooting section below.
+
+## Usage
+
+```python
+from doc_page_extractor import PageExtractor
+
+# Your code here
+```
+
+## Troubleshooting
+
+### "PyTorch is required but not installed!"
+
+Install PyTorch first:
+```bash
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
+```
+
+### "CUDA is not available!"
+
+**Check your GPU driver:**
+```bash
+nvidia-smi
+```
+
+**If the command fails**, you need to install NVIDIA drivers:
+- Download from: https://www.nvidia.com/download/index.aspx
+
+**If it succeeds**, you might have CPU-only PyTorch. Reinstall with CUDA:
+```bash
+pip uninstall torch torchvision
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
+```
+
+## Requirements
+
+- Python >= 3.10, < 3.14
+- **NVIDIA GPU with CUDA 11.8 or 12.1 support (Required)**
+- Sufficient GPU memory (recommended: 4GB+ VRAM)
+
+## Development
+
+For contributors and developers, see [Development Guide](docs/DEVELOPMENT.md) for:
+- Running tests
+- Running lint checks
+- Building the package
+
+

doc_page_extractor-1.0.2.dist-info/RECORD

@@ -0,0 +1,11 @@
+doc_page_extractor/__init__.py,sha256=BCLTWrjj0r8HJGoUzrY4T630WiN-di33NWwYOLF7YXc,191
+doc_page_extractor/check_env.py,sha256=pYk_58eqhSbe3GB0INYli6mATjCvZtHUTo3QX-ZExzw,1460
+doc_page_extractor/extractor.py,sha256=XSP_LlIKjOgKS-25fxPouCJnWmh2KNF6Z79oc3b3QGs,3096
+doc_page_extractor/model.py,sha256=I5RLt5GsoWeux5QHrKM80uVuN5kB6drsQx-Gp8X3wEk,3111
+doc_page_extractor/parser.py,sha256=1PdDKQ6SOftoklVH5DnvJYUhJPHtVr0hclGxgBIj2LE,1652
+doc_page_extractor/plot.py,sha256=3ZD-rw__7pu7EPMnxwHpHkhLolbOPJbdDK4949XsKKA,1647
+doc_page_extractor/redacter.py,sha256=jVfH-XWmuq2IYn4g1tGSnZc6gUXrRhn7roSCPtoYbHQ,3227
+doc_page_extractor-1.0.2.dist-info/METADATA,sha256=MAXKv7u4f4MNCak7-PArUTzUClTDCbyL34nn1iFAlMY,3182
+doc_page_extractor-1.0.2.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+doc_page_extractor-1.0.2.dist-info/licenses/LICENSE,sha256=1Kv5XShR6SbZVHr1Z_2tBC8oFk_rfO6CBtmmygj3Jlo,1074
+doc_page_extractor-1.0.2.dist-info/RECORD,,

{doc_page_extractor-0.1.1.dist-info → doc_page_extractor-1.0.2.dist-info}/WHEEL

@@ -1,5 +1,4 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: poetry-core 2.2.1
 Root-Is-Purelib: true
 Tag: py3-none-any
-

doc_page_extractor-1.0.2.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2025 Tao Zeyu
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

doc_page_extractor/clipper.py

@@ -1,119 +0,0 @@
-import numpy as np
-
-from math import pi, ceil, sin, cos, sqrt
-from PIL.Image import Image, Transform
-from .types import Layout, ExtractedResult
-from .rectangle import Rectangle
-from .rotation import calculate_rotation_with_rect, normal_vertical_rotation
-
-
-def clip(
-    extracted_result: ExtractedResult,
-    layout: Layout,
-    wrapped_width: float = 0.0,
-    wrapped_height: float = 0.0,
-  ) -> Image:
-  image: Image
-  if extracted_result.adjusted_image is None:
-    image = extracted_result.extracted_image
-  else:
-    image = extracted_result.adjusted_image
-
-  return clip_from_image(
-    image, layout.rect,
-    wrapped_width, wrapped_height,
-  )
-
-def clip_from_image(
-    image: Image,
-    rect: Rectangle,
-    wrapped_width: float = 0.0,
-    wrapped_height: float = 0.0,
-  ) -> Image:
-  horizontal_rotation, vertical_rotation = calculate_rotation_with_rect(rect)
-  image = image.copy()
-  matrix_move = np.array(_get_move_matrix(rect.lt[0], rect.lt[1])).reshape(3, 3)
-  matrix_rotate = np.array(_get_rotate_matrix(-horizontal_rotation)).reshape(3, 3)
-  matrix = np.dot(matrix_move, matrix_rotate)
-
-  y_axis_rotation = normal_vertical_rotation(vertical_rotation - horizontal_rotation)
-
-  if abs(y_axis_rotation - 0.25 * pi) > 0.0:
-    x = cos(y_axis_rotation)
-    y = sin(y_axis_rotation)
-    matrix_shear = np.array(_get_shear_matrix(x, y)).reshape(3, 3)
-    matrix = np.dot(matrix, matrix_shear)
-
-  width, height, max_width, max_height = _size_and_wrapper(rect)
-  max_width += wrapped_width
-  max_height += wrapped_height
-
-  if max_width != width or max_height != height:
-    dx = (max_width - width) / 2.0
-    dy = (max_height - height) / 2.0
-    matrix_move = np.array(_get_move_matrix(-dx, -dy)).reshape(3, 3)
-    matrix = np.dot(matrix, matrix_move)
-
-  return image.transform(
-    size=(ceil(max_width), ceil(max_height)),
-    method=Transform.AFFINE,
-    data=_to_pillow_matrix(matrix),
-  )
-
-def _size_and_wrapper(rect: Rectangle):
-  widths: list[float] = []
-  heights: list[float] = []
-
-  for i, (p1, p2) in enumerate(rect.segments):
-    dx = p2[0] - p1[0]
-    dy = p2[1] - p1[1]
-    distance = sqrt(dx*dx + dy*dy)
-    if i % 2 == 0:
-      heights.append(distance)
-    else:
-      widths.append(distance)
-
-  if len(widths) == 0 and len(heights) == 0:
-    return 0.0, 0.0, 0.0, 0.0
-
-  width: float = sum(widths) / len(widths)
-  height: float = sum(heights) / len(heights)
-  max_width: float = width
-  max_height: float = height
-
-  for width in widths:
-    if width > max_width:
-      max_width = width
-
-  for height in heights:
-    if height > max_height:
-      max_height = height
-
-  return width, height, max_width, max_height
-
-def _to_pillow_matrix(matrix: np.array):
-  return (
-    matrix[0][0], matrix[0][1], matrix[0][2],
-    matrix[1][0], matrix[1][1], matrix[1][2],
-  )
-
-def _get_move_matrix(dx: float, dy: float):
-  return (
-    1.0, 0.0, dx,
-    0.0, 1.0, dy,
-    0.0, 0.0, 1.0,
-  )
-
-def _get_rotate_matrix(rotation: float):
-  return (
-    cos(rotation),  sin(rotation),  0.0,
-    -sin(rotation), cos(rotation),  0.0,
-    0.0,            0.0,            1.0
-  )
-
-def _get_shear_matrix(x0: float, y0: float):
-  return (
-    1.0, 0.0, 0.0,
-    x0,  y0,  0.0,
-    0.0, 0.0, 1.0,
-  )

doc_page_extractor/downloader.py

@@ -1,16 +0,0 @@
-import os
-import requests
-from pathlib import Path
-
-
-def download(url: str, file_path: Path):
-  response = requests.get(url, stream=True, timeout=60)
-  if response.status_code != 200:
-    raise FileNotFoundError(f"Failed to download file from {url}: {response.status_code}")
-  try:
-    with open(file_path, "wb") as file:
-      file.write(response.content)
-  except Exception as e:
-    if os.path.exists(file_path):
-      os.remove(file_path)
-    raise e

doc_page_extractor/latex.py

@@ -1,57 +0,0 @@
-import os
-import torch
-import requests
-
-from munch import Munch
-from pix2tex.cli import LatexOCR
-from PIL.Image import Image
-from .utils import expand_image
-
-
-class LaTeX:
-  def __init__(self, model_path: str):
-    self._model_path: str = model_path
-    self._model: LatexOCR | None = None
-
-  def extract(self, image: Image) -> str | None:
-    image = expand_image(image, 0.1) # 添加边缘提高识别准确率
-    model = self._get_model()
-    with torch.no_grad():
-      return model(image)
-
-  def _get_model(self) -> LatexOCR:
-    if self._model is None:
-      if not os.path.exists(self._model_path):
-        self._download_model()
-
-      self._model = LatexOCR(Munch({
-        "config": os.path.join("settings", "config.yaml"),
-        "checkpoint": os.path.join(self._model_path, "weights.pth"),
-        "no_cuda": True,
-        "no_resize": False,
-      }))
-    return self._model
-
-  # from https://github.com/lukas-blecher/LaTeX-OCR/blob/5c1ac929bd19a7ecf86d5fb8d94771c8969fcb80/pix2tex/model/checkpoints/get_latest_checkpoint.py#L37-L45
-  def _download_model(self):
-    os.makedirs(self._model_path, exist_ok=True)
-    tag = "v0.0.1"
-    files: list[tuple[str, str]] = (
-      ("weights.pth", f"https://github.com/lukas-blecher/LaTeX-OCR/releases/download/{tag}/weights.pth"),
-      ("image_resizer.pth", f"https://github.com/lukas-blecher/LaTeX-OCR/releases/download/{tag}/image_resizer.pth")
-    )
-    for file_name, url in files:
-      file_path = os.path.join(self._model_path, file_name)
-      try:
-        with open(file_path, "wb") as file:
-          response = requests.get(url, stream=True, timeout=15)
-          response.raise_for_status()
-          for chunk in response.iter_content(chunk_size=8192):
-            if chunk:  # 过滤掉保持连接的新块
-              file.write(chunk)
-          file.flush()
-
-      except BaseException as e:
-        if os.path.exists(file_path):
-          os.remove(file_path)
-        raise e