detectivesam_inference/checkpoint.py

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}."
        )