detectivesam_inference/runtime.py

from __future__ import annotations

from contextlib import nullcontext
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import torch
from hydra import initialize_config_dir
from hydra.core.global_hydra import GlobalHydra

from detectivesam_inference.checkpoint import (
    InferenceConfig,
    load_inference_config,
    resolve_checkpoint_path,
    resolve_repo_path,
)
from detectivesam_inference.dataset import PreparedSample
from detectivesam_inference.models.forgerylocalizer import ForgeryLocalizer


@dataclass(frozen=True)
class PredictionResult:
    probability: np.ndarray
    pred_mask: np.ndarray


def get_repo_root() -> Path:
    return Path(__file__).resolve().parent.parent


def select_device(device: str | None = None) -> torch.device:
    if device is not None:
        return torch.device(device)
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def initialize_sam2_config(config_dir: str | Path) -> None:
    config_dir = str(Path(config_dir).resolve())
    hydra = GlobalHydra.instance()
    current_dir = getattr(initialize_sam2_config, "_current_dir", None)
    if hydra.is_initialized():
        if current_dir == config_dir:
            return
        hydra.clear()
    initialize_config_dir(config_dir=config_dir, version_base=None)
    initialize_sam2_config._current_dir = config_dir


class DetectiveSAMRunner:
    def __init__(
        self,
        checkpoint_path: str | Path | None = None,
        device: str | None = None,
    ) -> None:
        self.repo_root = get_repo_root()
        self.checkpoint_path = resolve_checkpoint_path(checkpoint_path, self.repo_root)
        self.device = select_device(device)
        self.config = load_inference_config(self.checkpoint_path)
        self.model = self._load_model()

    def _load_model(self) -> ForgeryLocalizer:
        sam_config_path = resolve_repo_path(self.config.sam_config_file, self.repo_root)
        sam_checkpoint_path = resolve_repo_path(self.config.sam_checkpoint, self.repo_root)
        initialize_sam2_config(sam_config_path.parent)

        model = ForgeryLocalizer(
            sam_config=sam_config_path.name,
            sam_checkpoint=str(sam_checkpoint_path),
            prompt_dim=self.config.prompt_dim,
            dropout_rate=self.config.dropout_rate,
            max_streams=self.config.max_streams,
            device=str(self.device),
        ).to(self.device)

        checkpoint = torch.load(self.checkpoint_path, map_location=self.device, weights_only=False)
        state_dict = self._select_state_dict(checkpoint)
        state_dict = self._clean_state_dict_prefixes(state_dict)
        try:
            model.load_state_dict(state_dict)
        except RuntimeError:
            self._load_compatible_state_dict(model, state_dict)
        model.eval()
        return model

    @staticmethod
    def _select_state_dict(checkpoint) -> dict[str, torch.Tensor]:
        if isinstance(checkpoint, dict):
            for key in ("model", "ema", "model_state_dict", "state_dict"):
                if key in checkpoint:
                    return checkpoint[key]
        return checkpoint

    @staticmethod
    def _clean_state_dict_prefixes(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        return {
            key.removeprefix("module.").removeprefix("_orig_mod."): value
            for key, value in state_dict.items()
        }

    @staticmethod
    def _load_compatible_state_dict(
        model: ForgeryLocalizer,
        state_dict: dict[str, torch.Tensor],
    ) -> None:
        current_state = model.state_dict()
        compatible = {
            key: value
            for key, value in state_dict.items()
            if key in current_state and current_state[key].shape == value.shape
        }
        if not compatible:
            raise RuntimeError("Checkpoint did not contain any tensors compatible with the selected architecture.")
        model.load_state_dict(compatible, strict=False)

    def autocast_context(self):
        if self.device.type == "cuda":
            return torch.amp.autocast(device_type="cuda")
        return nullcontext()

    def predict_sample(
        self,
        sample: PreparedSample,
        threshold: float = 0.5,
    ) -> PredictionResult:
        orig = sample.orig.unsqueeze(0).to(self.device)
        streams = [stream.unsqueeze(0).to(self.device) for stream in sample.streams]

        with torch.inference_mode():
            with self.autocast_context():
                logits = self.model(orig, streams)

        probability = torch.sigmoid(logits).squeeze().detach().cpu().numpy()
        pred_mask = (probability > threshold).astype("uint8")
        return PredictionResult(probability=probability, pred_mask=pred_mask)