File size: 2,672 Bytes
84c3ae8 a973131 84c3ae8 a973131 84c3ae8 a973131 84c3ae8 a973131 84c3ae8 a973131 84c3ae8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import json
import os
import requests
from modelscope import snapshot_download
def download_json(url):
# 下载JSON文件
response = requests.get(url)
response.raise_for_status() # 检查请求是否成功
return response.json()
def download_and_modify_json(url, local_filename, modifications):
if os.path.exists(local_filename):
data = json.load(open(local_filename))
config_version = data.get('config_version', '0.0.0')
if config_version < '1.0.0':
data = download_json(url)
else:
data = download_json(url)
# 修改内容
for key, value in modifications.items():
data[key] = value
# 保存修改后的内容
with open(local_filename, 'w', encoding='utf-8') as f:
json.dump(data, f, ensure_ascii=False, indent=4)
if __name__ == '__main__':
mineru_patterns = [
"models/Layout/LayoutLMv3/*",
"models/Layout/YOLO/*",
"models/MFD/YOLO/*",
"models/MFR/unimernet_small/*",
"models/TabRec/TableMaster/*",
"models/TabRec/StructEqTable/*",
]
# 在 Docker 环境中使用 /app 目录
base_dir = "/app"
cache_dir = os.path.join(base_dir, "model_cache")
os.makedirs(cache_dir, exist_ok=True)
# 下载模型，使用与原代码相同的参数
try:
model_dir = snapshot_download('opendatalab/PDF-Extract-Kit-1.0',
cache_dir=cache_dir,
revision='master')
model_dir = os.path.join(model_dir, 'models')
except TypeError:
# 如果 allow_patterns 不支持，则下载完整模型
model_dir = snapshot_download('opendatalab/PDF-Extract-Kit-1.0',
cache_dir=cache_dir)
model_dir = os.path.join(model_dir, 'models')
layoutreader_model_dir = snapshot_download('ppaanngggg/layoutreader',
cache_dir=cache_dir)
print(f'model_dir is: {model_dir}')
print(f'layoutreader_model_dir is: {layoutreader_model_dir}')
# 配置文件保存到 Docker 容器的应用目录
json_url = 'https://gitee.com/myhloli/MinerU/raw/master/magic-pdf.template.json'
config_file_name = 'magic-pdf.json'
config_file = os.path.join(base_dir, config_file_name)
json_mods = {
'models-dir': model_dir,
'layoutreader-model-dir': layoutreader_model_dir,
}
download_and_modify_json(json_url, config_file, json_mods)
print(f'配置文件已成功配置，路径为: {config_file}') |