| |
| """Where the tests, the golden generator and the CPU reference find the SceneSeg weights without downloading anything. |
| |
| Order: ``$SCENESEG_WEIGHTS_DIR`` > the weights staged for the upload in this checkout (``weights_upload/weights``, |
| PLAN D10) > the workspace download (``assets/sceneseg/gdrive_SceneSeg_FP32_onnx`` and ``gdrive_SceneSeg_pth``, the |
| Google Drive originals, sha256 in ``config``) > the local HF cache snapshot of ``changh95/sceneseg-p150`` at the |
| pinned weights revision (``api.SceneSeg.DEFAULT_REVISION``, once it exists). A directory may hold the files directly |
| or under ``weights/`` (the layout of the HF repo). The Python API resolves weights itself |
| (``ttaw.api_base.resolve_weights``: the same env var, then the Hub at the pinned revision).""" |
| from __future__ import annotations |
|
|
| import os |
| from pathlib import Path |
| from typing import List, Optional |
|
|
| from . import config |
|
|
| __all__ = ["find_weights_dir", "onnx_path", "state_dict_path", "find_in", "workspace_root", "pinned_revision", |
| "test_weights_dir", "ENV_VAR", "TEST_WEIGHTS_ENV"] |
|
|
| ENV_VAR = "SCENESEG_WEIGHTS_DIR" |
| |
| |
| TEST_WEIGHTS_ENV = "SCENESEG_TEST_WEIGHTS" |
|
|
|
|
| def workspace_root() -> Optional[Path]: |
| """The tt-models workspace when this checkout lives in ``<ws>/bundles/sceneseg-p150``, else None.""" |
| root = Path(__file__).resolve().parents[5] |
| return root if (root / "research" / "sceneseg").is_dir() else None |
|
|
|
|
| def find_in(directory: Path, name: str) -> Optional[Path]: |
| """``<dir>/<name>`` or ``<dir>/weights/<name>``.""" |
| for cand in (directory / name, directory / config.WEIGHTS_SUBDIR / name): |
| if cand.is_file(): |
| return cand |
| return None |
|
|
|
|
| def _candidates() -> List[Path]: |
| out = [] |
| if os.environ.get(ENV_VAR): |
| out.append(Path(os.environ[ENV_VAR]).expanduser()) |
| bundle = Path(__file__).resolve().parents[3] |
| out.append(bundle / "weights_upload") |
| ws = workspace_root() |
| if ws is not None: |
| out += [ws / "assets" / "sceneseg" / "gdrive_SceneSeg_FP32_onnx", |
| ws / "assets" / "sceneseg" / "gdrive_SceneSeg_pth"] |
| revision = pinned_revision() |
| if revision: |
| try: |
| from huggingface_hub import snapshot_download |
|
|
| out.append(Path(snapshot_download(config.WEIGHTS_REPO, revision=revision, |
| allow_patterns=config.ALLOW_PATTERNS, local_files_only=True))) |
| except Exception: |
| pass |
| return out |
|
|
|
|
| def pinned_revision() -> Optional[str]: |
| """The weights commit ``api.SceneSeg.DEFAULT_REVISION`` pins, or None while it is still deferred.""" |
| import re |
|
|
| from ..api import SceneSeg |
|
|
| rev = SceneSeg.DEFAULT_REVISION or "" |
| return rev if re.fullmatch(r"[0-9a-f]{40}", rev) else None |
|
|
|
|
| def _find(names, required: bool) -> Optional[Path]: |
| for d in _candidates(): |
| for name in names: |
| p = find_in(d, name) |
| if p is not None: |
| return p |
| if required: |
| raise FileNotFoundError(f"none of {names} found: set {ENV_VAR} to a directory holding them (or their " |
| f"{config.WEIGHTS_SUBDIR}/ sub-directory)") |
| return None |
|
|
|
|
| def onnx_path(required: bool = True) -> Optional[Path]: |
| """The deployed ``SceneSeg_FP32.onnx``.""" |
| return _find((config.ONNX_FILE,), required) |
|
|
|
|
| def state_dict_path(required: bool = True) -> Optional[Path]: |
| """The upstream state_dict: the re-hosted ``.safetensors`` or the original ``.pth``.""" |
| return _find((config.SAFETENSORS_FILE, config.PTH_FILE), required) |
|
|
|
|
| def test_weights_dir() -> Optional[Path]: |
| """The weights directory of the device tests: ``None`` when ``$SCENESEG_TEST_WEIGHTS`` is ``hub`` (the caller |
| then lets ``SceneSeg.from_pretrained`` resolve the pinned Hub snapshot itself; ``$SCENESEG_WEIGHTS_DIR`` must be |
| unset), else :func:`find_weights_dir` (required).""" |
| mode = os.environ.get(TEST_WEIGHTS_ENV, "local") |
| if mode == "hub": |
| if os.environ.get(ENV_VAR): |
| raise RuntimeError(f"{TEST_WEIGHTS_ENV}=hub with {ENV_VAR} set: unset {ENV_VAR}") |
| return None |
| if mode != "local": |
| raise ValueError(f"{TEST_WEIGHTS_ENV}={mode!r}: expected 'hub' or 'local'") |
| return find_weights_dir(True) |
|
|
|
|
| def hub_onnx_path() -> Path: |
| """``weights/SceneSeg_FP32.onnx`` of the pinned snapshot, resolved like ``SceneSeg.from_pretrained``.""" |
| from ..api import SceneSeg |
|
|
| p = find_in(SceneSeg.resolve_weights(), config.ONNX_FILE) |
| if p is None: |
| raise FileNotFoundError(f"{config.ONNX_FILE} not in the pinned snapshot of {SceneSeg.DEFAULT_REPO}") |
| return p |
|
|
|
|
| def find_weights_dir(required: bool = False) -> Optional[Path]: |
| """The directory holding the ONNX (the ``weights/`` level when the files sit there).""" |
| p = onnx_path(required=required) |
| return None if p is None else p.parent |
|
|
