isaac/scripts/fetch_models.py

#!/usr/bin/env python3
"""Fetch the model weights the pipeline needs. They are far too big for git.

    python scripts/fetch_models.py

Downloads in small Range chunks with per-chunk retries: the CRE weights come off a slow
mirror and a plain single-stream GET truncates silently, leaving an unloadable file.
"""
from __future__ import annotations

import os
import sys
import time
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
MODELS = ROOT / "assets" / "models"
CHUNK = 1 << 21
TRIES = 12

SOURCES = {
    "FastSAM-s.pt":
        "https://huggingface.co/Ultralytics/FastSAM/resolve/main/FastSAM-s.pt?download=true",
    # CRE-Stereo ETH3D weights, from the megvii research release mirror
    "crestereo_eth3d.pth":
        "https://github.com/ibaiGorordo/CREStereo-Pytorch/releases/download/0.0.1/crestereo_eth3d.pth",
}

CRE_REPO = "https://github.com/ibaiGorordo/CREStereo-Pytorch"


def content_length(url):
    req = urllib.request.Request(url, method="HEAD")
    with urllib.request.urlopen(req, timeout=60) as r:
        if r.status in (301, 302) and r.headers.get("Location"):
            return content_length(r.headers["Location"])
        return int(r.headers.get("Content-Length", 0))


def get_range(url, a, b):
    req = urllib.request.Request(url)
    req.add_header("Range", f"bytes={a}-{b}")
    with urllib.request.urlopen(req, timeout=120) as r:
        return r.read()


def fetch(name, url, out_dir):
    out = out_dir / name
    total = content_length(url)
    have = out.stat().st_size if out.exists() else 0
    if total and have == total:
        print(f"  {name}: already complete ({total/1e6:.1f} MB)")
        return True
    if not total:                                  # server refuses HEAD: plain download
        print(f"  {name}: streaming (no content-length)")
        urllib.request.urlretrieve(url, out)
        return out.exists()
    if have > total:
        have = 0
    mode = "r+b" if have else "wb"
    pos = have
    with open(out, mode) as fh:
        fh.seek(pos)
        t0 = time.time()
        while pos < total:
            end = min(pos + CHUNK - 1, total - 1)
            for attempt in range(TRIES):
                try:
                    data = get_range(url, pos, end)
                    if len(data) != end - pos + 1:
                        raise IOError("short chunk")
                    fh.write(data)
                    fh.flush()
                    pos += len(data)
                    pct = 100.0 * pos / total
                    rate = (pos - have) / max(time.time() - t0, 1e-6) / 1024
                    print(f"\r  {name}: {pct:5.1f}%  {rate:6.0f} KB/s", end="", flush=True)
                    break
                except Exception as exc:
                    if attempt == TRIES - 1:
                        print(f"\n  {name}: FAILED at byte {pos}: {exc}")
                        return False
                    time.sleep(min(2 ** attempt, 20))
    print()
    ok = out.stat().st_size == total
    print(f"  {name}: {'ok' if ok else 'SIZE MISMATCH'} ({out.stat().st_size/1e6:.1f} MB)")
    return ok


def main():
    MODELS.mkdir(parents=True, exist_ok=True)
    print(f"fetching model weights into {MODELS}")
    results = {n: fetch(n, u, MODELS) for n, u in SOURCES.items()}

    cre_dir = MODELS / "crestereo"
    if not (cre_dir / "nets").exists():
        print(f"\nCRE-Stereo network code is not vendored here. Clone it next to the weights:")
        print(f"  git clone {CRE_REPO} {cre_dir}")
        print("  (only the `nets` package is imported)")

    missing = [n for n, ok in results.items() if not ok]
    if missing:
        print(f"\nincomplete: {missing} - re-run, downloads resume where they stopped")
        return 1
    print("\nall weights present")
    return 0


if __name__ == "__main__":
    sys.exit(main())