AntonioJun/workspace

1"""Build hypothesis spatial codes from existing native or combined geometry caches."""2 3from __future__ import annotations4 5import argparse6import gzip7import json8import os9from pathlib import Path10import pickle11import sys12import types13 14WORKSPACE_ROOT = Path(__file__).resolve().parent.parent15if str(WORKSPACE_ROOT) not in sys.path:16    sys.path.insert(0, str(WORKSPACE_ROOT))17 18from encoder import adapters  # noqa: E40219from encoder import config as encoder_config  # noqa: E40220from experiments import config, loader  # noqa: E40221from experiments import adapters as format_adapters  # noqa: E40222 23 24class CachedDA3Prediction:25    """Attribute container used only to deserialize cached DA3 Prediction objects."""26 27 28class CachedAddictDict(dict):29    """Minimal attribute-access dictionary required by cached DA3 metadata."""30 31    def __getattr__(self, name):32        if name.startswith("__") and name.endswith("__"):33            raise AttributeError(name)34        try:35            return self[name]36        except KeyError as exc:37            raise AttributeError(name) from exc38 39    def __setattr__(self, name, value):40        self[name] = value41 42 43def install_da3_pickle_compatibility() -> None:44    """Provide the exact pickle class when the original DA3 package is unavailable."""45    try:46        __import__("depth_anything_3.specs")47    except ModuleNotFoundError:48        package = types.ModuleType("depth_anything_3")49        package.__path__ = []50        specs = types.ModuleType("depth_anything_3.specs")51        CachedDA3Prediction.__module__ = "depth_anything_3.specs"52        CachedDA3Prediction.__name__ = "Prediction"53        CachedDA3Prediction.__qualname__ = "Prediction"54        specs.Prediction = CachedDA3Prediction55        package.specs = specs56        sys.modules["depth_anything_3"] = package57        sys.modules["depth_anything_3.specs"] = specs58    try:59        __import__("addict.addict")60    except ModuleNotFoundError:61        package = types.ModuleType("addict")62        package.__path__ = []63        module = types.ModuleType("addict.addict")64        CachedAddictDict.__module__ = "addict.addict"65        CachedAddictDict.__name__ = "Dict"66        CachedAddictDict.__qualname__ = "Dict"67        module.Dict = CachedAddictDict68        package.addict = module69        package.Dict = CachedAddictDict70        sys.modules["addict"] = package71        sys.modules["addict.addict"] = module72 73 74def configure_single_scene_threads() -> int:75    """Give a direct single-scene run every CPU available to the process."""76    try:77        count = max(1, len(os.sched_getaffinity(0)))78    except AttributeError:79        count = max(1, os.cpu_count() or 1)80    value = str(count)81    for variable in (82        "OMP_NUM_THREADS",83        "MKL_NUM_THREADS",84        "OPENBLAS_NUM_THREADS",85        "NUMEXPR_NUM_THREADS",86        "VSI_KD_WORKERS",87    ):88        os.environ.setdefault(variable, value)89    return count90 91 92def load_existing_geometry(scene, depth, input_selection, tracking, frame_count):93    """Read existing caches without ever creating or modifying source caches."""94    combined_path = Path(95        encoder_config.cache_file(scene, depth, input_selection, tracking, frame_count)96    )97    if combined_path.is_file():98        with gzip.open(combined_path, "rb") as stream:99            return adapters.validate(pickle.load(stream))100 101    da3_path = Path(102        encoder_config.da3_cache_file(scene, depth, input_selection, frame_count)103    )104    sam3_path = Path(105        encoder_config.sam3_cache_file(scene, input_selection, tracking, frame_count)106    )107    missing = [str(path) for path in (da3_path, sam3_path) if not path.is_file()]108    if missing:109        raise FileNotFoundError(110            "required native cache(s) are missing: " + ", ".join(missing)111        )112    install_da3_pickle_compatibility()113    geometry = adapters.adapt(114        encoder_config.MODEL,115        scene=scene,116        root=str(encoder_config.CACHE_ROOT),117        rebuild=False,118        da3_path=str(da3_path),119        sam3_path=str(sam3_path),120    )121    return adapters.validate(geometry)122 123 124def run_scene(125    scene,126    hypothesis,127    depth="metric",128    tracking="tracking",129    input_selection="uniform",130    frame_count=64,131    rebuild=False,132    spatial_code_format="explicit",133):134    output = config.spatial_code_path(135        scene,136        hypothesis,137        depth,138        tracking,139        input_selection,140        frame_count,141        spatial_code_format,142    )143    if output.is_file() and not rebuild:144        with output.open(encoding="utf-8") as stream:145            return json.load(stream), "loaded", output146    geometry = load_existing_geometry(147        scene, depth, input_selection, tracking, frame_count148    )149    geometry_math = loader.load_hypothesis(hypothesis)150    code, *_ = format_adapters.build(geometry_math, geometry, spatial_code_format)151    output.parent.mkdir(parents=True, exist_ok=True)152    geometry_math.dump_spatial_code(code, str(output))153    return code, "built", output154 155 156def main() -> None:157    parser = argparse.ArgumentParser()158    parser.add_argument("scene", nargs="?")159    parser.add_argument("--hypothesis")160    parser.add_argument(161        "--depth", default="metric", choices=encoder_config.DEPTH_VARIANTS162    )163    parser.add_argument(164        "--input",165        default="uniform",166        choices=encoder_config.INPUT_SELECTIONS,167        dest="input_selection",168    )169    parser.add_argument(170        "--tracking", default="tracking", choices=encoder_config.TRACKING_MODES171    )172    parser.add_argument("--frames", type=int, default=64)173    parser.add_argument(174        "--format",175        default="explicit",176        choices=config.SPATIAL_CODE_FORMATS,177        dest="spatial_code_format",178    )179    parser.add_argument("--rebuild", action="store_true")180    parser.add_argument("--list", action="store_true", dest="list_only")181    args = parser.parse_args()182    if args.list_only:183        print("\n".join(loader.list_hypotheses()))184        return185    if not args.scene:186        parser.error("scene is required unless --list is used")187    if not args.hypothesis:188        parser.error("--hypothesis is required unless --list is used")189    if args.frames < 1:190        parser.error("--frames must be positive")191    configure_single_scene_threads()192    _, status, path = run_scene(193        args.scene,194        args.hypothesis,195        args.depth,196        args.tracking,197        args.input_selection,198        args.frames,199        args.rebuild,200        args.spatial_code_format,201    )202    print(f"[{args.scene}] {status} -> {path}")203 204 205if __name__ == "__main__":206    main()207