File size: 6,047 Bytes
b4bedeb 0b2df65 b4bedeb 0b2df65 b4bedeb 0b2df65 b4bedeb 0b2df65 b4bedeb 0b2df65 b4bedeb 0b2df65 | 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 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | from __future__ import annotations
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import yaml
DEFAULT_CHECKPOINT = Path("checkpoints/detective_sam_v2.pth")
CHECKPOINT_ALIASES = {"detective_sam_v2": DEFAULT_CHECKPOINT}
V2_ADAPTER_TYPE = "spatial_cross_attention"
V2_MASK_ADAPTER_TYPE = "transformer"
@dataclass(frozen=True)
class InferenceConfig:
img_size: int
prompt_dim: int
downscale: int
dropout_rate: float
perturbation_type: str
perturbation_intensity: float
sam_config_file: str
sam_checkpoint: str
adapter_type: str = V2_ADAPTER_TYPE
mask_adapter_type: str = V2_MASK_ADAPTER_TYPE
@property
def max_streams(self) -> int:
return count_perturbation_streams(self.perturbation_type)
def resolve_checkpoint_path(checkpoint_value: str | Path | None, repo_root: str | Path) -> Path:
repo_root = Path(repo_root)
if checkpoint_value is None:
return repo_root / DEFAULT_CHECKPOINT
checkpoint_str = str(checkpoint_value)
if checkpoint_str in CHECKPOINT_ALIASES:
return repo_root / CHECKPOINT_ALIASES[checkpoint_str]
checkpoint_path = Path(checkpoint_value)
if checkpoint_path.is_absolute():
return checkpoint_path
if checkpoint_path.exists():
return checkpoint_path.resolve()
repo_candidate = repo_root / checkpoint_path
if repo_candidate.exists():
return repo_candidate
aliased_checkpoint = repo_root / "checkpoints" / f"{checkpoint_str}.pth"
if aliased_checkpoint.exists():
return aliased_checkpoint
return repo_candidate
def resolve_repo_path(path_value: str | Path, repo_root: str | Path) -> Path:
path = Path(path_value)
if path.is_absolute():
return path
repo_root = Path(repo_root)
direct = repo_root / path
if direct.exists():
return direct
sam_config = repo_root / "sam2configs" / path.name
if sam_config.exists():
return sam_config
return direct
def load_inference_config(checkpoint_path: str | Path) -> InferenceConfig:
params = _load_params_file(checkpoint_path)
config = InferenceConfig(
img_size=int(_resolve_param(params, "img_size", section="training_config", default=512)),
prompt_dim=int(
_resolve_param(
params,
"prompt_dim",
section="model_config",
default=_resolve_param(params, "prompt", section="model_config", default=128),
)
),
downscale=int(_resolve_param(params, "downscale", section="model_config", default=16)),
dropout_rate=float(
_resolve_param(
params,
"dropout_rate",
section="model_config",
default=_resolve_param(params, "dropout", section="model_config", default=0.1),
)
),
perturbation_type=str(_resolve_param(params, "perturbation_type", section="data_config", default="none")),
perturbation_intensity=float(
_resolve_param(params, "perturbation_intensity", section="data_config", default=0.0)
),
sam_config_file=str(
_resolve_param(params, "sam_config_file", section="sam_config", default="sam2.1_hiera_b+.yaml")
),
sam_checkpoint=str(
_resolve_param(
params,
"sam_checkpoint",
section="sam_config",
default="sam2configs/sam2.1_hiera_base_plus.pt",
)
),
adapter_type=str(_resolve_param(params, "adapter_type", section="model_config", default=V2_ADAPTER_TYPE)),
mask_adapter_type=str(
_resolve_param(params, "mask_adapter_type", section="model_config", default=V2_MASK_ADAPTER_TYPE)
),
)
_validate_v2_config(config)
return config
def count_perturbation_streams(perturbation_type: str) -> int:
if perturbation_type == "none":
return 0
if "+" in perturbation_type:
return len([item for item in perturbation_type.split("+") if item.strip()])
if "/" in perturbation_type:
return len([item for item in perturbation_type.split("/") if item.strip()])
return 1
def _load_params_file(checkpoint_path: str | Path) -> dict[str, Any]:
checkpoint_path = Path(checkpoint_path)
candidate_paths = [
checkpoint_path.with_name(f"{checkpoint_path.stem}_params.yaml"),
checkpoint_path.with_name(f"{checkpoint_path.stem}_params.yml"),
checkpoint_path.with_name(f"{checkpoint_path.stem}_params.json"),
checkpoint_path.parent / "model_params.yaml",
checkpoint_path.parent / "model_params.yml",
checkpoint_path.parent / "model_params.json",
]
for candidate in candidate_paths:
if candidate.exists():
with candidate.open("r", encoding="utf-8") as handle:
loaded = json.load(handle) if candidate.suffix == ".json" else yaml.safe_load(handle)
if not isinstance(loaded, dict):
raise ValueError(f"Checkpoint params file must deserialize to a mapping: {candidate}")
return loaded
raise FileNotFoundError(
f"Could not find a params file for checkpoint {checkpoint_path}. "
f"Checked: {', '.join(str(path) for path in candidate_paths)}"
)
def _resolve_param(
params: dict[str, Any],
key: str,
*,
section: str,
default: Any,
) -> Any:
if key in params:
return params[key]
return params.get(section, {}).get(key, default)
def _validate_v2_config(config: InferenceConfig) -> None:
if config.adapter_type != V2_ADAPTER_TYPE or config.mask_adapter_type != V2_MASK_ADAPTER_TYPE:
raise ValueError(
"This release is DetectiveSAMv2-only and requires "
f"adapter_type={V2_ADAPTER_TYPE!r}, mask_adapter_type={V2_MASK_ADAPTER_TYPE!r}. "
f"Got adapter_type={config.adapter_type!r}, mask_adapter_type={config.mask_adapter_type!r}."
)
|