| |
| """ |
| Download the MediaPipe Tasks model bundles used by the extractor. |
| |
| MediaPipe >= 0.10.18 (and 1.x, which this project runs on) no longer ships the |
| legacy ``mediapipe.solutions`` API; landmark extraction therefore uses the Tasks |
| API with explicit ``.task`` bundles. |
| |
| Usage |
| ----- |
| python scripts/download_models.py # holistic (default) + pose + hand |
| python scripts/download_models.py --only holistic |
| python scripts/download_models.py --check # verify what is already present |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| import urllib.request |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
|
|
| import config |
|
|
| ASSETS_DIR = Path(config.MODELS_DIR) / "mediapipe_assets" |
|
|
| BUNDLES = { |
| "holistic": "https://storage.googleapis.com/mediapipe-models/holistic_landmarker/" |
| "holistic_landmarker/float16/1/holistic_landmarker.task", |
| "pose": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/" |
| "pose_landmarker_lite/float16/1/pose_landmarker_lite.task", |
| "hand": "https://storage.googleapis.com/mediapipe-models/hand_landmarker/" |
| "hand_landmarker/float16/1/hand_landmarker.task", |
| "pose_full": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/" |
| "pose_landmarker_full/float16/1/pose_landmarker_full.task", |
| } |
| MIN_BYTES = 1_000_000 |
|
|
|
|
| def fetch(name: str, url: str, force: bool = False) -> Path: |
| dest = ASSETS_DIR / Path(url).name |
| if dest.exists() and dest.stat().st_size > MIN_BYTES and not force: |
| print(f" present {dest.name} ({dest.stat().st_size/1e6:.1f} MB)") |
| return dest |
| dest.parent.mkdir(parents=True, exist_ok=True) |
| print(f" fetch {dest.name} <- {url}") |
| req = urllib.request.Request(url, headers={"User-Agent": "isl-recognition/1.0"}) |
| with urllib.request.urlopen(req, timeout=300) as r, open(dest, "wb") as f: |
| while True: |
| chunk = r.read(1 << 20) |
| if not chunk: |
| break |
| f.write(chunk) |
| size = dest.stat().st_size |
| if size < MIN_BYTES: |
| raise RuntimeError(f"{dest.name} looks truncated ({size} bytes)") |
| print(f" saved {dest.name} ({size/1e6:.1f} MB)") |
| return dest |
|
|
|
|
| def main(argv=None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--only", choices=list(BUNDLES), default=None) |
| ap.add_argument("--check", action="store_true") |
| ap.add_argument("--force", action="store_true") |
| args = ap.parse_args(argv) |
|
|
| targets = {args.only: BUNDLES[args.only]} if args.only else \ |
| {k: v for k, v in BUNDLES.items() if k != "pose_full"} |
|
|
| print(f"MediaPipe assets -> {ASSETS_DIR}") |
| missing = [] |
| for name, url in targets.items(): |
| path = ASSETS_DIR / Path(url).name |
| if args.check: |
| ok = path.exists() and path.stat().st_size > MIN_BYTES |
| print(f" {'OK ' if ok else 'MISSING'} {path.name}") |
| if not ok: |
| missing.append(name) |
| else: |
| fetch(name, url, force=args.force) |
|
|
| if args.check and missing: |
| print(f"\nRun `python scripts/download_models.py --only {missing[0]}` to fetch the missing bundle(s).") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
