yuze.zyz 1 tahun lalu
induk
melakukan
4559181f84
1 mengubah file dengan 6 tambahan dan 2 penghapusan
  1. 6 2
      pdf2zh/doclayout.py

+ 6 - 2
pdf2zh/doclayout.py

@@ -1,4 +1,6 @@
 import abc
+import os.path
+
 import cv2
 import numpy as np
 import ast
@@ -11,7 +13,7 @@ class DocLayoutModel(abc.ABC):
     @staticmethod
     def load_onnx():
         model = OnnxModel.from_pretrained(
-            repo_id="wybxc/DocLayout-YOLO-DocStructBench-onnx",
+            repo_id='AI-ModelScope/DocLayout-YOLO-DocStructBench-onnx',
             filename="doclayout_yolo_docstructbench_imgsz1024.onnx",
         )
         return model
@@ -70,7 +72,9 @@ class OnnxModel(DocLayoutModel):
 
     @staticmethod
     def from_pretrained(repo_id: str, filename: str):
-        pth = hf_hub_download(repo_id=repo_id, filename=filename, etag_timeout=1)
+        from modelscope import snapshot_download
+        model_dir = snapshot_download(repo_id)
+        pth = os.path.join(model_dir, filename)
         return OnnxModel(pth)
 
     @property