152 lines
6.4 KiB
Python
152 lines
6.4 KiB
Python
from __future__ import annotations
|
|
|
|
from textwrap import dedent
|
|
|
|
|
|
def offline_helper_script() -> str:
|
|
"""Return the self-contained helper served to an internet-connected computer."""
|
|
|
|
return dedent(
|
|
r'''#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import shutil
|
|
import tarfile
|
|
import urllib.request
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
|
|
IMAGE_REPO = "sentence-transformers/clip-ViT-B-32"
|
|
TEXT_REPO = "sentence-transformers/clip-ViT-B-32-multilingual-v1"
|
|
AUDIO_REPO = "OpenVINO/whisper-small-fp16-ov"
|
|
RAPID_ROOT = (
|
|
"https://raw.githubusercontent.com/RapidAI/RapidOCR/main/python/rapidocr_onnxruntime/"
|
|
"rapidocr_onnxruntime/models"
|
|
)
|
|
ZOO_ROOT = "https://storage.openvinotoolkit.org/repositories/open_model_zoo/2022.3/models_bin/1"
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
while chunk := handle.read(1024 * 1024):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def download(url: str, destination: Path) -> None:
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
partial = destination.with_name(destination.name + ".part")
|
|
existing = partial.stat().st_size if partial.exists() else 0
|
|
request = urllib.request.Request(url, headers={"Range": f"bytes={existing}-"} if existing else {})
|
|
with urllib.request.urlopen(request, timeout=60) as response:
|
|
append = existing > 0 and getattr(response, "status", 200) == 206
|
|
with partial.open("ab" if append else "wb") as output:
|
|
shutil.copyfileobj(response, output, 1024 * 1024)
|
|
if destination.suffix == ".xml":
|
|
prefix = partial.read_bytes()[:512].lstrip().lower()
|
|
if prefix.startswith((b"<!doctype html", b"<html")) or b"<net " not in partial.read_bytes()[:4096]:
|
|
partial.unlink(missing_ok=True)
|
|
raise RuntimeError(f"人物模型下载地址返回了无效 XML：{url}")
|
|
if destination.suffix == ".bin" and partial.stat().st_size < 64 * 1024:
|
|
partial.unlink(missing_ok=True)
|
|
raise RuntimeError(f"人物模型权重过小，下载地址可能返回了错误页面：{url}")
|
|
partial.replace(destination)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="下载 ImageFind 的画面、OCR、人物与音频模型，并生成一个可上传的离线包"
|
|
)
|
|
parser.add_argument("--hf-endpoint", default="https://hf-mirror.com", help="Hugging Face 官方站或镜像")
|
|
parser.add_argument("--cache-dir", type=Path, default=Path(".imagefind-model-cache"), help="断点缓存目录")
|
|
parser.add_argument("--output", type=Path, help="输出 .tar.gz；默认包含当前日期")
|
|
args = parser.parse_args()
|
|
endpoint = args.hf_endpoint.rstrip("/")
|
|
if not endpoint.startswith(("http://", "https://")) or "@" in endpoint.split("://", 1)[-1].split("/", 1)[0]:
|
|
parser.error("--hf-endpoint 必须是无内嵌账号密码的 HTTP/HTTPS 地址")
|
|
try:
|
|
from huggingface_hub import snapshot_download
|
|
except ImportError as exc:
|
|
raise SystemExit("缺少模型准备依赖；请先运行: python -m pip install huggingface_hub") from exc
|
|
|
|
cache = args.cache_dir.resolve()
|
|
stage = cache / "stage"
|
|
image_dir = stage / "visual" / "image"
|
|
text_dir = stage / "visual" / "text"
|
|
snapshot_download(IMAGE_REPO, endpoint=endpoint, cache_dir=cache / "hub", local_dir=image_dir)
|
|
snapshot_download(TEXT_REPO, endpoint=endpoint, cache_dir=cache / "hub", local_dir=text_dir)
|
|
audio_dir = stage / "audio"
|
|
snapshot_download(AUDIO_REPO, endpoint=endpoint, cache_dir=cache / "hub", local_dir=audio_dir)
|
|
|
|
rapid_files = {
|
|
"det.onnx": "ch_PP-OCRv4_det_infer.onnx",
|
|
"rec.onnx": "ch_PP-OCRv4_rec_infer.onnx",
|
|
"cls.onnx": "ch_ppocr_mobile_v2.0_cls_infer.onnx",
|
|
}
|
|
for local_name, remote_name in rapid_files.items():
|
|
download(f"{RAPID_ROOT}/{remote_name}", stage / "ocr" / local_name)
|
|
|
|
face_files = {
|
|
"detector.xml": "face-detection-retail-0004/FP16/face-detection-retail-0004.xml",
|
|
"detector.bin": "face-detection-retail-0004/FP16/face-detection-retail-0004.bin",
|
|
"reidentification.xml": (
|
|
"face-reidentification-retail-0095/FP16/face-reidentification-retail-0095.xml"
|
|
),
|
|
"reidentification.bin": (
|
|
"face-reidentification-retail-0095/FP16/face-reidentification-retail-0095.bin"
|
|
),
|
|
}
|
|
for local_name, remote_name in face_files.items():
|
|
download(f"{ZOO_ROOT}/{remote_name}", stage / "faces" / local_name)
|
|
|
|
created = datetime.now(UTC)
|
|
files = sorted(path for path in stage.rglob("*") if path.is_file() and path.name != "manifest.json")
|
|
sources = {
|
|
"visual_image": f"https://huggingface.co/{IMAGE_REPO}",
|
|
"visual_text": f"https://huggingface.co/{TEXT_REPO}",
|
|
"ocr": RAPID_ROOT,
|
|
"faces": ZOO_ROOT,
|
|
"audio": f"https://huggingface.co/{AUDIO_REPO}",
|
|
}
|
|
manifest = {
|
|
"format_version": 2,
|
|
"version": f"offline-{created.date().isoformat()}",
|
|
"created_at": created.isoformat(),
|
|
"source": "imagefind-offline-helper",
|
|
"sources": sources,
|
|
"components": {
|
|
"visual": {
|
|
"version": f"{IMAGE_REPO}+{TEXT_REPO}",
|
|
"sources": [sources["visual_image"], sources["visual_text"]],
|
|
},
|
|
"ocr": {"version": "RapidOCR-PP-OCRv4", "source": sources["ocr"]},
|
|
"faces": {"version": "open-model-zoo-2022.3", "source": sources["faces"]},
|
|
"audio": {"version": AUDIO_REPO, "source": sources["audio"]},
|
|
},
|
|
"files": {path.relative_to(stage).as_posix(): sha256(path) for path in files},
|
|
}
|
|
manifest_bytes = (json.dumps(manifest, ensure_ascii=False, indent=2) + "\n").encode()
|
|
output = (args.output or Path(f"imagefind-models-{created.date().isoformat()}.tar.gz")).resolve()
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
with tarfile.open(output, "w:gz") as archive:
|
|
for path in files:
|
|
archive.add(path, path.relative_to(stage).as_posix(), recursive=False)
|
|
info = tarfile.TarInfo("manifest.json")
|
|
info.size = len(manifest_bytes)
|
|
info.mode = 0o644
|
|
info.mtime = int(created.timestamp())
|
|
archive.addfile(info, io.BytesIO(manifest_bytes))
|
|
print(f"已生成 {output}")
|
|
print(f"SHA-256 {sha256(output)}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
'''
|
|
)
|