chore(setup): add model pre-download utility and CLI flag; update README usage
Browse files- README.md +14 -0
- app.py +13 -2
- utils/model_setup.py +65 -0
README.md
CHANGED
|
@@ -40,6 +40,20 @@ python app.py
|
|
| 40 |
|
| 41 |
ブラウザで表示されるGradio UIからポーズ編集を行えます。
|
| 42 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 43 |
## 使い方(Hugging Face Spaces)
|
| 44 |
|
| 45 |
スペース版の基本フローは「入力 → 編集 → 出力」です。
|
|
|
|
| 40 |
|
| 41 |
ブラウザで表示されるGradio UIからポーズ編集を行えます。
|
| 42 |
|
| 43 |
+
### 初回セットアップ(モデルの事前ダウンロード)
|
| 44 |
+
|
| 45 |
+
初回起動時に DWPose の ONNX モデルを自動ダウンロードしますが、起動前に取得しておきたい場合は次のどちらかを実行してください。
|
| 46 |
+
|
| 47 |
+
```bash
|
| 48 |
+
# 1) app.pyのオプションでダウンロードだけ実施
|
| 49 |
+
python app.py --setup-models
|
| 50 |
+
|
| 51 |
+
# 2) ユーティリティを直接実行(同等の動作)
|
| 52 |
+
python -m utils.model_setup
|
| 53 |
+
```
|
| 54 |
+
|
| 55 |
+
デフォルトでは `./models` にキャッシュします。企業ネットワーク等でトークンが必要な場合は、先に `huggingface-cli login` を実施してください。
|
| 56 |
+
|
| 57 |
## 使い方(Hugging Face Spaces)
|
| 58 |
|
| 59 |
スペース版の基本フローは「入力 → 編集 → 出力」です。
|
app.py
CHANGED
|
@@ -6,6 +6,7 @@ from utils.notifications import notify_success, notify_error, NotificationMessag
|
|
| 6 |
from utils.coordinate_system import update_coordinate_system
|
| 7 |
from utils.image_processing import process_uploaded_image
|
| 8 |
from utils.export_utils import export_pose_as_image, export_pose_as_json, get_timestamp_filename
|
|
|
|
| 9 |
import json
|
| 10 |
import tempfile
|
| 11 |
import base64
|
|
@@ -959,5 +960,15 @@ def main():
|
|
| 959 |
return demo
|
| 960 |
|
| 961 |
if __name__ == "__main__":
|
| 962 |
-
|
| 963 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 6 |
from utils.coordinate_system import update_coordinate_system
|
| 7 |
from utils.image_processing import process_uploaded_image
|
| 8 |
from utils.export_utils import export_pose_as_image, export_pose_as_json, get_timestamp_filename
|
| 9 |
+
from utils.model_setup import download_models as _download_models_setup
|
| 10 |
import json
|
| 11 |
import tempfile
|
| 12 |
import base64
|
|
|
|
| 960 |
return demo
|
| 961 |
|
| 962 |
if __name__ == "__main__":
|
| 963 |
+
import argparse, sys
|
| 964 |
+
parser = argparse.ArgumentParser()
|
| 965 |
+
parser.add_argument("--setup-models", action="store_true", help="Download required models into ./models and exit")
|
| 966 |
+
args, _ = parser.parse_known_args()
|
| 967 |
+
|
| 968 |
+
if args.setup_models:
|
| 969 |
+
code = _download_models_setup()
|
| 970 |
+
# exit with code so CI/automation can check
|
| 971 |
+
sys.exit(code)
|
| 972 |
+
else:
|
| 973 |
+
demo = main()
|
| 974 |
+
demo.launch()
|
utils/model_setup.py
ADDED
|
@@ -0,0 +1,65 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Lightweight model setup utility for dwpose-editor.
|
| 3 |
+
|
| 4 |
+
Provides a CLI to pre-download required ONNX models into the local cache
|
| 5 |
+
directory used by the app (./models), without initializing ONNX sessions.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
from typing import List
|
| 9 |
+
import os
|
| 10 |
+
import sys
|
| 11 |
+
|
| 12 |
+
try:
|
| 13 |
+
from huggingface_hub import hf_hub_download
|
| 14 |
+
except Exception as e: # pragma: no cover
|
| 15 |
+
hf_hub_download = None
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
DEFAULT_REPO_ID = "yzd-v/DWPose"
|
| 19 |
+
DEFAULT_CACHE_DIR = "./models"
|
| 20 |
+
REQUIRED_FILES: List[str] = [
|
| 21 |
+
"yolox_l.onnx",
|
| 22 |
+
"dw-ll_ucoco_384.onnx",
|
| 23 |
+
]
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def ensure_dir(path: str) -> None:
|
| 27 |
+
os.makedirs(path, exist_ok=True)
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def download_models(repo_id: str = DEFAULT_REPO_ID, cache_dir: str = DEFAULT_CACHE_DIR) -> int:
|
| 31 |
+
"""Download required model files into cache_dir. Returns 0 on success, non-zero on failure."""
|
| 32 |
+
if hf_hub_download is None:
|
| 33 |
+
print("[ERROR] huggingface-hub is not installed. Run: pip install huggingface-hub")
|
| 34 |
+
return 2
|
| 35 |
+
|
| 36 |
+
ensure_dir(cache_dir)
|
| 37 |
+
|
| 38 |
+
ok = True
|
| 39 |
+
for filename in REQUIRED_FILES:
|
| 40 |
+
try:
|
| 41 |
+
print(f"[SETUP] Downloading {filename} from {repo_id} → {cache_dir} ...")
|
| 42 |
+
local_path = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=cache_dir)
|
| 43 |
+
print(f"[SETUP] OK: {local_path}")
|
| 44 |
+
except Exception as e:
|
| 45 |
+
ok = False
|
| 46 |
+
print(f"[SETUP] ERROR: failed to download {filename}: {e}")
|
| 47 |
+
|
| 48 |
+
return 0 if ok else 1
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
def main(argv: List[str] | None = None) -> int:
|
| 52 |
+
"""CLI entry point: python -m utils.model_setup [--repo REPO] [--cache-dir DIR]"""
|
| 53 |
+
import argparse
|
| 54 |
+
|
| 55 |
+
parser = argparse.ArgumentParser(description="Pre-download DWPose models into ./models")
|
| 56 |
+
parser.add_argument("--repo", default=DEFAULT_REPO_ID, help="Hugging Face repo id")
|
| 57 |
+
parser.add_argument("--cache-dir", default=DEFAULT_CACHE_DIR, help="Local cache directory")
|
| 58 |
+
args = parser.parse_args(argv)
|
| 59 |
+
|
| 60 |
+
return download_models(repo_id=args.repo, cache_dir=args.cache_dir)
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
if __name__ == "__main__": # pragma: no cover
|
| 64 |
+
sys.exit(main())
|
| 65 |
+
|