File size: 3,846 Bytes
5652fbb fd515c9 5652fbb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | from __future__ import annotations
import json
import os
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
@dataclass(frozen=True)
class LoadedModelBundle:
model: object
input_shape: tuple[int, int, int, int]
config_path: Path
weights_path: Path
model_dir: Path
@dataclass(frozen=True)
class ModelArtifactPaths:
model_dir: Path
config_path: Path
weights_path: Path
def _choose_latest(paths: list[Path]) -> Path:
if not paths:
raise FileNotFoundError("No candidate model files were found.")
return sorted(paths, key=lambda p: p.stat().st_mtime_ns)[-1]
def resolve_model_artifacts(model_dir: Path) -> ModelArtifactPaths:
model_dir = Path(model_dir).expanduser().resolve()
if not model_dir.exists():
raise FileNotFoundError(f"Model directory not found: {model_dir}")
config_path = model_dir / "config.json"
if not config_path.exists():
raise FileNotFoundError(f"Missing model config: {config_path}")
candidates: list[Path] = []
candidates.extend(sorted(model_dir.glob("*.final.weights.h5")))
named_best = model_dir / "best_model_dynamic.weights.h5"
if named_best.exists():
candidates.append(named_best)
candidates.extend(sorted(model_dir.glob("*.weights.h5")))
# Deduplicate while preserving order.
seen = set()
deduped: list[Path] = []
for path in candidates:
key = str(path.resolve())
if key in seen:
continue
seen.add(key)
deduped.append(path)
if not deduped:
raise FileNotFoundError(
"No weights file found in model directory. "
"Expected something like '*.final.weights.h5' or '*.weights.h5'."
)
weights_path = _choose_latest(deduped)
return ModelArtifactPaths(
model_dir=model_dir,
config_path=config_path,
weights_path=weights_path,
)
def _read_model_config(config_path: Path) -> dict:
payload = json.loads(config_path.read_text())
input_shape = payload.get("INPUT_SHAPE")
if not input_shape:
raise RuntimeError(f"INPUT_SHAPE missing in model config: {config_path}")
return payload
def _should_force_cpu() -> bool:
raw = os.environ.get("INFER_FORCE_CPU", "0").strip().lower()
return raw in {"1", "true", "yes", "on"}
@lru_cache(maxsize=4)
def _load_model_cached(
model_dir_str: str,
config_mtime_ns: int,
weights_mtime_ns: int,
) -> LoadedModelBundle:
del config_mtime_ns, weights_mtime_ns
if _should_force_cpu():
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "-1")
from tensorflow.keras import mixed_precision
from .model_defs import ModelBuildConfig, build_dynamic_model
artifacts = resolve_model_artifacts(Path(model_dir_str))
payload = _read_model_config(artifacts.config_path)
# Force float32 inference graph for CPU compatibility.
mixed_precision.set_global_policy("float32")
cfg = ModelBuildConfig(
input_shape=tuple(int(x) for x in payload["INPUT_SHAPE"]),
base_filters=int(payload.get("BASE_FILTERS", 8)),
sam_heads=int(payload.get("SAM_HEADS", 2)),
l2_reg=float(payload.get("L2_REG", 0.0)),
)
model = build_dynamic_model(cfg)
model.load_weights(str(artifacts.weights_path))
return LoadedModelBundle(
model=model,
input_shape=cfg.input_shape,
config_path=artifacts.config_path,
weights_path=artifacts.weights_path,
model_dir=artifacts.model_dir,
)
def load_model_bundle(model_dir: Path) -> LoadedModelBundle:
artifacts = resolve_model_artifacts(model_dir)
return _load_model_cached(
str(artifacts.model_dir),
artifacts.config_path.stat().st_mtime_ns,
artifacts.weights_path.stat().st_mtime_ns,
)
|