将onnx模型和字体打包进docker镜像

.dockerignore

@@ -0,0 +1,3 @@
+.venv
+.github
+docs

.gitignore

@@ -165,6 +165,7 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+.idea/
 .vscode
 .DS_Store
+uv.lock

Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.12
+FROM ghcr.io/astral-sh/uv:python3.12-bookworm-slim
 
 WORKDIR /app
 
@@ -7,9 +7,11 @@ COPY . .
 EXPOSE 7860
 
 ENV PYTHONUNBUFFERED=1
+ADD "https://github.com/satbyy/go-noto-universal/releases/download/v7.0/GoNotoKurrent-Regular.ttf" /app
+RUN apt-get update && \
+     apt-get install --no-install-recommends -y libgl1 && \
+     rm -rf /var/lib/apt/lists/*
 
-RUN apt-get update && apt-get install -y libgl1
-
-RUN pip install .
+RUN uv pip install --system --no-cache . && uv run pdf2zh/warmup.py
 
 CMD ["pdf2zh", "-i"]

pdf2zh/doclayout.py

@@ -2,6 +2,7 @@
 import os.path
 
 import cv2
+import huggingface_hub.utils
 import numpy as np
 import ast
 import onnx
@@ -82,7 +83,12 @@ def from_pretrained(repo_id: str, filename: str):
             model_dir = snapshot_download(repo_mapping[repo_id])
             pth = os.path.join(model_dir, filename)
         else:
-            pth = hf_hub_download(repo_id=repo_id, filename=filename, etag_timeout=1)
+            try:
+                pth = hf_hub_download(repo_id=repo_id, filename=filename, etag_timeout=1, local_files_only=True)
+                print("Using local DocLayout-YOLO-DocStructBench-onnx file", pth)
+            except huggingface_hub.utils.LocalEntryNotFoundError:
+                print("Downloading DocLayout-YOLO-DocStructBench-onnx from Huggingface Hub...")
+                pth = hf_hub_download(repo_id=repo_id, filename=filename, etag_timeout=1)
         return OnnxModel(pth)
 
     @property

pdf2zh/high_level.py

@@ -186,7 +186,10 @@ def translate_stream(
         font_list.append((resfont, None))
     elif lang_out.lower() in noto_list:  # noto
         resfont = "noto"
-        ttf_path = os.path.join(tempfile.gettempdir(), "GoNotoKurrent-Regular.ttf")
+        # docker
+        ttf_path = '/app/GoNotoKurrent-Regular.ttf'
+        if not os.path.exists(ttf_path):
+            ttf_path = os.path.join(tempfile.gettempdir(), "GoNotoKurrent-Regular.ttf")
         if not os.path.exists(ttf_path):
             print("Downloading Noto font...")
             urllib.request.urlretrieve(

pdf2zh/warmup.py

@@ -0,0 +1,10 @@
+from doclayout import DocLayoutModel
+
+def warmup():
+    print('Warming up the model...')
+    DocLayoutModel.load_available()
+    print('Warm-up completed!')
+
+
+if __name__ == "__main__":
+    warmup()