Publish DetectiveSAM inference bundle

.gitattributes

@@ -1,35 +1,16 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+sam2configs/*.pt filter=lfs diff=lfs merge=lfs -text
+checkpoints/*.pth filter=lfs diff=lfs merge=lfs -text
+demo/cocoglide/source/airplane_139871.png filter=lfs diff=lfs merge=lfs -text
+demo/cocoglide/source/banana_28809.png filter=lfs diff=lfs merge=lfs -text
+demo/cocoglide/source/giraffe_296969.png filter=lfs diff=lfs merge=lfs -text
+demo/cocoglide/source/train_221213.png filter=lfs diff=lfs merge=lfs -text
+demo/cocoglide/source/tv_453722.png filter=lfs diff=lfs merge=lfs -text
+demo/cocoglide/target/airplane_139871.png filter=lfs diff=lfs merge=lfs -text
+demo/cocoglide/target/banana_28809.png filter=lfs diff=lfs merge=lfs -text
+demo/cocoglide/target/giraffe_296969.png filter=lfs diff=lfs merge=lfs -text
+demo/cocoglide/target/train_221213.png filter=lfs diff=lfs merge=lfs -text
+demo/cocoglide/target/tv_453722.png filter=lfs diff=lfs merge=lfs -text
+demo/flux_test/source/548.png filter=lfs diff=lfs merge=lfs -text
+demo/flux_test/target/548.png filter=lfs diff=lfs merge=lfs -text
+demo/qwen_test/source/166.png filter=lfs diff=lfs merge=lfs -text
+demo/qwen_test/target/166.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.pyc
+.DS_Store
+.venv/
+venv/
+.pytest_cache/
+outputs/

README.md

@@ -0,0 +1,121 @@
+---
+license: other
+library_name: pytorch
+pipeline_tag: image-segmentation
+tags:
+  - image-forensics
+  - image-manipulation-detection
+  - image-segmentation
+  - sam2
+  - pytorch
+---
+
+# DetectiveSAM
+
+DetectiveSAM is an inference-only image forgery localization bundle built around SAM2. This release includes bundled checkpoints and a small set of ready-to-run examples for demos.
+
+## What is bundled
+
+- Inference checkpoints under `checkpoints/`
+- SAM2 config and weights under `sam2configs/`
+- Poster demo pairs under `demo/cocoglide/`, `demo/flux_test/`, and `demo/qwen_test/`
+- A drop-in single-image slot at `demo/user_image/demo_input.png`
+
+Built-in checkpoint aliases:
+
+- `detective_sam`
+- `detective_sam_sota`
+
+## Setup
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+```
+
+## Hugging Face Usage
+
+```bash
+git lfs install
+git clone https://huggingface.co/Gertlek/DetectiveSAM
+cd DetectiveSAM
+python -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+python -m detectivesam_inference.predict \
+  --checkpoint detective_sam \
+  --output-dir outputs/poster_baseline
+```
+
+## Poster Demo Flows
+
+### 1. Live single-image demo
+
+Place your image at `demo/user_image/demo_input.png`, then run:
+
+```bash
+python -m detectivesam_inference.predict \
+  --checkpoint detective_sam \
+  --output-dir outputs/poster_user_image
+```
+
+In this mode the CLI reuses the target image as its own source reference so the demo stays runnable with a single image.
+
+### 2. Bundled baseline example
+
+If `demo/user_image/demo_input.png` is absent, the default `predict` command falls back to the bundled CocoGlide sample `banana_28809`.
+
+```bash
+python -m detectivesam_inference.predict \
+  --checkpoint detective_sam \
+  --output-dir outputs/poster_baseline
+```
+
+### 3. Bundled SOTA examples
+
+Flux example:
+
+```bash
+python -m detectivesam_inference.predict \
+  --checkpoint detective_sam_sota \
+  --source demo/flux_test/source/548.png \
+  --target demo/flux_test/target/548.png \
+  --mask demo/flux_test/mask/548.png \
+  --output-dir outputs/poster_flux
+```
+
+Qwen example:
+
+```bash
+python -m detectivesam_inference.predict \
+  --checkpoint detective_sam_sota \
+  --source demo/qwen_test/source/166.png \
+  --target demo/qwen_test/target/166.png \
+  --mask demo/qwen_test/mask/166.png \
+  --output-dir outputs/poster_qwen
+```
+
+## Outputs
+
+Each `predict` run writes a compact set of visual artifacts plus a JSON summary:
+
+- `<name>_comparison.png`
+- `<name>_probability.png`
+- `<name>_pred_mask.png`
+- `<name>_pred_overlay.png`
+- `<name>_summary.json`
+
+If a ground-truth mask is provided, the run also saves:
+
+- `<name>_gt_mask.png`
+- `<name>_gt_overlay.png`
+
+The `evaluate` command writes `summary.json` plus a few visualization examples under `visualizations/`.
+
+## Notes
+
+- The runtime selects `cuda` automatically when available and otherwise runs on CPU.
+- Checkpoint settings come from the YAML sidecars in `checkpoints/`; you only need the alias or checkpoint path.
+- This repo does not include training code or training-only dependencies.
+- License metadata is currently marked `other`: the bundled SAM2 components are Apache-2.0, while DetectiveSAM release terms should be finalized before broader redistribution.

checkpoints/detective_sam_sota.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0c7c33c4c4ddb82459bc2b6bbfdfd1003561e2714d02163b3daffb25ede4e1cf
+size 297395157

checkpoints/detective_sam_sota_params.yaml

@@ -0,0 +1,8 @@
+img_size: 512
+prompt_dim: 64
+downscale: 16
+dropout_rate: 0.2
+perturbation_type: gaussian_blur/gaussian_noise
+perturbation_intensity: 0.75
+sam_config_file: sam2.1_hiera_b+.yaml
+sam_checkpoint: sam2configs/sam2.1_hiera_base_plus.pt

checkpoints/model_epoch22_batch999_score1.1114.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:898e64f831f6e7e4e8486c0f27b139cbaa44119c449fad2570ab4f6b331ef421
+size 297395157

checkpoints/model_epoch22_batch999_score1.1114_params.yaml

@@ -0,0 +1,8 @@
+img_size: 512
+prompt_dim: 64
+downscale: 16
+dropout_rate: 0.15
+perturbation_type: gaussian_blur/gaussian_noise
+perturbation_intensity: 0.75
+sam_config_file: sam2.1_hiera_b+.yaml
+sam_checkpoint: sam2configs/sam2.1_hiera_base_plus.pt

demo/user_image/README.md

@@ -0,0 +1,14 @@
+# User Image Demo
+
+Place a single demo image here as:
+
+- `demo_input.png`
+
+Then run:
+
+```bash
+python -m detectivesam_inference.predict
+```
+
+When only a target image is available, the CLI reuses that image as its own reference source. That keeps the demo runnable, but it is not the canonical pairwise evaluation mode the model was trained for.
+

detectivesam_inference/__init__.py

@@ -0,0 +1,4 @@
+from detectivesam_inference.checkpoint import InferenceConfig
+from detectivesam_inference.runtime import DetectiveSAMRunner, PredictionResult
+
+__all__ = ["DetectiveSAMRunner", "InferenceConfig", "PredictionResult"]

detectivesam_inference/checkpoint.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+DEFAULT_CHECKPOINT = Path("checkpoints/model_epoch22_batch999_score1.1114.pth")
+CHECKPOINT_ALIASES = {
+    "detective_sam": DEFAULT_CHECKPOINT,
+    "detective_sam_sota": Path("checkpoints/detective_sam_sota.pth"),
+}
+
+
+@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
+
+    @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)
+    return 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",
+            )
+        ),
+    )
+
+
+def count_perturbation_streams(perturbation_type: str) -> int:
+    if perturbation_type == "none":
+        return 0
+    if "+" in perturbation_type:
+        return len(perturbation_type.split("+"))
+    if "/" in perturbation_type:
+        return len(perturbation_type.split("/"))
+    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.parent / "model_params.yaml",
+    ]
+    for candidate in candidate_paths:
+        if candidate.exists():
+            with candidate.open("r", encoding="utf-8") as handle:
+                loaded = 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)

detectivesam_inference/dataset.py

@@ -0,0 +1,249 @@
+from __future__ import annotations
+
+import hashlib
+from dataclasses import dataclass
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+import torchvision.transforms.functional as TF
+from PIL import Image
+from sam2.utils.transforms import SAM2Transforms
+from torch.utils.data import Dataset
+
+from detectivesam_inference.perturbations import (
+    add_gaussian_noise_deterministic,
+    apply_blur_to_image_tensor,
+    apply_jpeg_compression_to_tensor,
+)
+
+
+LEGACY_CONTRASTIVE_FLAG = False
+
+
+@dataclass
+class PreparedSample:
+    name: str
+    source_path: Path
+    target_path: Path
+    mask_path: Path | None
+    source_image: Image.Image
+    target_image: Image.Image
+    orig: torch.Tensor
+    streams: list[torch.Tensor]
+    mask: torch.Tensor | None
+
+
+def parse_perturbation_types(perturbation_type: str) -> list[str]:
+    if perturbation_type == "none":
+        return []
+    if "+" in perturbation_type:
+        return [item.strip() for item in perturbation_type.split("+")]
+    if "/" in perturbation_type:
+        return [item.strip() for item in perturbation_type.split("/")]
+    return [perturbation_type.strip()]
+
+
+def compute_perturbation_params(perturbation_intensity: float) -> dict[str, float | int]:
+    return {
+        "blur_sigma": perturbation_intensity * 2.0,
+        "jpeg_quality": max(10, int(95 - (perturbation_intensity * 56.67))),
+        "noise_std": perturbation_intensity * 0.2,
+    }
+
+
+def create_combined_mask(
+    mask_rgba: np.ndarray,
+) -> np.ndarray:
+    if mask_rgba.ndim == 2:
+        return (mask_rgba // 255).astype(np.uint8)
+
+    if mask_rgba.ndim == 3 and mask_rgba.shape[2] == 4:
+        alpha = mask_rgba[:, :, 3]
+        alpha_is_opaque = alpha.sum() == alpha.size * 255
+        if alpha_is_opaque:
+            foreground = (mask_rgba[:, :, 0] > 0).astype(np.uint8)
+            return cv2.resize(foreground, (512, 512), interpolation=cv2.INTER_NEAREST).astype(np.uint8)
+
+        _, binary = cv2.threshold(alpha, 0, 255, cv2.THRESH_BINARY)
+        return (1 - (binary // 255)).astype(np.uint8)
+
+    return (mask_rgba[:, :, 0] > 0).astype(np.uint8)
+
+
+def resize_triplet(
+    source_image: Image.Image,
+    target_image: Image.Image,
+    mask_image: Image.Image | None,
+    img_size: int,
+) -> tuple[Image.Image, Image.Image, Image.Image | None]:
+    source_resized = source_image.resize((img_size, img_size), Image.BILINEAR)
+    target_resized = target_image.resize((img_size, img_size), Image.BILINEAR)
+    if mask_image is None:
+        return source_resized, target_resized, None
+    return source_resized, target_resized, mask_image.resize((img_size, img_size), Image.NEAREST)
+
+
+def build_streams(
+    target_image: Image.Image,
+    perturbation_type: str,
+    perturbation_intensity: float,
+    seed: int,
+) -> list[torch.Tensor]:
+    perturbations = parse_perturbation_types(perturbation_type)
+    params = compute_perturbation_params(perturbation_intensity)
+    orig_tensor = TF.to_tensor(target_image)
+
+    streams: list[torch.Tensor] = []
+    for perturbation in perturbations:
+        if perturbation == "gaussian_blur":
+            streams.append(apply_blur_to_image_tensor(orig_tensor, sigma=float(params["blur_sigma"])))
+        elif perturbation == "jpeg_compression":
+            streams.append(apply_jpeg_compression_to_tensor(orig_tensor, quality=int(params["jpeg_quality"])))
+        elif perturbation == "gaussian_noise":
+            streams.append(
+                add_gaussian_noise_deterministic(
+                    orig_tensor,
+                    std=float(params["noise_std"]),
+                    seed=seed + len(streams),
+                )
+            )
+        else:
+            raise ValueError(f"Unsupported perturbation type: {perturbation}")
+    return streams
+
+
+def build_sample_seed(
+    source_path: Path,
+    target_path: Path,
+    mask_path: Path | None,
+    perturbation_type: str,
+    perturbation_intensity: float,
+) -> int:
+    sample_key = "|".join(
+        [
+            source_path.parent.parent.name,
+            source_path.stem,
+            target_path.parent.parent.name,
+            target_path.stem,
+            mask_path.stem if mask_path is not None else "no-mask",
+            perturbation_type,
+            f"{perturbation_intensity:.8f}",
+            str(LEGACY_CONTRASTIVE_FLAG),
+        ]
+    )
+    digest = hashlib.sha256(sample_key.encode("utf-8")).digest()
+    return int.from_bytes(digest[:8], byteorder="big", signed=False) % (2**31)
+
+
+def prepare_sample(
+    source_path: str | Path,
+    target_path: str | Path,
+    mask_path: str | Path | None,
+    img_size: int,
+    perturbation_type: str,
+    perturbation_intensity: float,
+) -> PreparedSample:
+    source_path = Path(source_path)
+    target_path = Path(target_path)
+    mask_path = Path(mask_path) if mask_path is not None else None
+
+    source_image = Image.open(source_path).convert("RGB")
+    target_image = Image.open(target_path).convert("RGB")
+    mask_image = Image.open(mask_path) if mask_path is not None else None
+    source_image, target_image, mask_image = resize_triplet(source_image, target_image, mask_image, img_size)
+
+    sample_seed = build_sample_seed(
+        source_path=source_path,
+        target_path=target_path,
+        mask_path=mask_path,
+        perturbation_type=perturbation_type,
+        perturbation_intensity=perturbation_intensity,
+    )
+    transforms = SAM2Transforms(resolution=img_size, mask_threshold=0.0)
+
+    orig_tensor = TF.to_tensor(target_image)
+    orig = transforms.transforms(orig_tensor).unsqueeze(0).squeeze(0)
+    streams_raw = build_streams(
+        target_image=target_image,
+        perturbation_type=perturbation_type,
+        perturbation_intensity=perturbation_intensity,
+        seed=sample_seed,
+    )
+    streams = [transforms.transforms(stream).unsqueeze(0).squeeze(0) for stream in streams_raw]
+
+    mask_tensor = None
+    if mask_image is not None:
+        binary_mask = create_combined_mask(
+            mask_rgba=np.array(mask_image),
+        )
+        if binary_mask.shape != (img_size, img_size):
+            binary_mask = cv2.resize(binary_mask, (img_size, img_size), interpolation=cv2.INTER_NEAREST)
+        mask_tensor = torch.tensor(binary_mask, dtype=torch.float32).unsqueeze(0)
+
+    return PreparedSample(
+        name=target_path.stem,
+        source_path=source_path,
+        target_path=target_path,
+        mask_path=mask_path,
+        source_image=source_image,
+        target_image=target_image,
+        orig=orig,
+        streams=streams,
+        mask=mask_tensor,
+    )
+
+
+class PairDataset(Dataset):
+    def __init__(
+        self,
+        root_dir: str | Path,
+        img_size: int,
+        perturbation_type: str,
+        perturbation_intensity: float,
+        max_samples: int | None = None,
+    ) -> None:
+        self.root_dir = Path(root_dir)
+        self.source_dir = self.root_dir / "source"
+        self.target_dir = self.root_dir / "target"
+        self.mask_dir = self.root_dir / "mask"
+        self.img_size = img_size
+        self.perturbation_type = perturbation_type
+        self.perturbation_intensity = perturbation_intensity
+
+        if not self.source_dir.exists() or not self.target_dir.exists():
+            raise FileNotFoundError(f"{root_dir} must contain source/ and target/ directories")
+
+        target_files = sorted(
+            path
+            for path in self.target_dir.iterdir()
+            if path.suffix.lower() in {".png", ".jpg", ".jpeg"}
+        )
+        if max_samples is not None:
+            target_files = target_files[:max_samples]
+        self.target_files = target_files
+
+    def __len__(self) -> int:
+        return len(self.target_files)
+
+    def __getitem__(self, index: int) -> PreparedSample:
+        target_path = self.target_files[index]
+        source_path = self.source_dir / target_path.name
+        if not source_path.exists():
+            png_fallback = self.source_dir / f"{target_path.stem}.png"
+            jpg_fallback = self.source_dir / f"{target_path.stem}.jpg"
+            source_path = png_fallback if png_fallback.exists() else jpg_fallback
+        if not source_path.exists():
+            raise FileNotFoundError(f"Could not find source image for {target_path.name}")
+
+        mask_candidate = self.mask_dir / f"{target_path.stem}.png"
+        mask_path = mask_candidate if mask_candidate.exists() else None
+        return prepare_sample(
+            source_path=source_path,
+            target_path=target_path,
+            mask_path=mask_path,
+            img_size=self.img_size,
+            perturbation_type=self.perturbation_type,
+            perturbation_intensity=self.perturbation_intensity,
+        )

detectivesam_inference/evaluate.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+from tqdm import tqdm
+
+from detectivesam_inference.dataset import PairDataset
+from detectivesam_inference.metrics import compute_f1, compute_iou, summarize_results
+from detectivesam_inference.runtime import DetectiveSAMRunner, get_repo_root
+from detectivesam_inference.visualization import save_prediction_outputs
+
+
+def parse_args() -> argparse.Namespace:
+    repo_root = get_repo_root()
+    parser = argparse.ArgumentParser(description="Evaluate DetectiveSAM on a dataset root with source/target/mask folders.")
+    parser.add_argument(
+        "--checkpoint",
+        default="detective_sam",
+        help="Checkpoint path or alias. Built-in aliases: detective_sam, detective_sam_sota.",
+    )
+    parser.add_argument("--dataset-root", default=str(repo_root / "demo" / "cocoglide"))
+    parser.add_argument("--output-dir", default=str(repo_root / "outputs" / "eval_demo"))
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--threshold", type=float, default=0.5)
+    parser.add_argument("--max-samples", type=int, default=None)
+    parser.add_argument("--num-visualizations", type=int, default=4)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    runner = DetectiveSAMRunner(checkpoint_path=args.checkpoint, device=args.device)
+    dataset = PairDataset(
+        root_dir=args.dataset_root,
+        img_size=runner.config.img_size,
+        perturbation_type=runner.config.perturbation_type,
+        perturbation_intensity=runner.config.perturbation_intensity,
+        max_samples=args.max_samples,
+    )
+
+    output_dir = Path(args.output_dir)
+    vis_dir = output_dir / "visualizations"
+    output_dir.mkdir(parents=True, exist_ok=True)
+    vis_dir.mkdir(parents=True, exist_ok=True)
+
+    per_sample_results: list[dict[str, float | str | None]] = []
+    for index, sample in enumerate(tqdm(dataset, desc="Evaluating")):
+        prediction = runner.predict_sample(sample, threshold=args.threshold)
+        gt_mask = sample.mask.squeeze().numpy().astype("uint8") if sample.mask is not None else None
+        per_sample_results.append(
+            {
+                "name": sample.name,
+                "iou": compute_iou(prediction.pred_mask, gt_mask) if gt_mask is not None else None,
+                "f1": compute_f1(prediction.pred_mask, gt_mask) if gt_mask is not None else None,
+            }
+        )
+
+        if index < args.num_visualizations:
+            save_prediction_outputs(
+                output_dir=vis_dir,
+                name=sample.name,
+                source_image=sample.source_image,
+                target_image=sample.target_image,
+                probability_map=prediction.probability,
+                pred_mask=prediction.pred_mask,
+                gt_mask=gt_mask,
+            )
+
+    payload = {
+        "checkpoint": str(runner.checkpoint_path.resolve()),
+        "dataset_root": str(Path(args.dataset_root).resolve()),
+        "threshold": args.threshold,
+        "summary": summarize_results(per_sample_results),
+        "samples": per_sample_results,
+    }
+    with (output_dir / "summary.json").open("w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2)
+
+    print(json.dumps(payload["summary"], indent=2))
+    print(f"Detailed results written to {output_dir / 'summary.json'}")
+
+
+if __name__ == "__main__":
+    main()

detectivesam_inference/metrics.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import numpy as np
+
+
+def compute_iou(pred_mask: np.ndarray, true_mask: np.ndarray) -> float:
+    pred = pred_mask.astype(bool)
+    true = true_mask.astype(bool)
+    intersection = np.logical_and(pred, true).sum()
+    union = np.logical_or(pred, true).sum()
+    if union == 0:
+        return 1.0
+    return float(intersection / union)
+
+
+def compute_f1(pred_mask: np.ndarray, true_mask: np.ndarray) -> float:
+    pred = pred_mask.astype(bool)
+    true = true_mask.astype(bool)
+    if not pred.any() and not true.any():
+        return 1.0
+    if not pred.any() or not true.any():
+        return 0.0
+
+    true_positive = np.logical_and(pred, true).sum()
+    false_positive = np.logical_and(pred, np.logical_not(true)).sum()
+    false_negative = np.logical_and(np.logical_not(pred), true).sum()
+
+    precision = true_positive / (true_positive + false_positive) if (true_positive + false_positive) else 0.0
+    recall = true_positive / (true_positive + false_negative) if (true_positive + false_negative) else 0.0
+    if precision + recall == 0:
+        return 0.0
+    return float(2 * precision * recall / (precision + recall))
+
+
+def summarize_results(results: list[dict[str, float | str | None]]) -> dict[str, int | float | None]:
+    if not results:
+        return {
+            "num_samples": 0,
+            "num_samples_with_gt": 0,
+            "mean_iou": None,
+            "mean_f1": None,
+        }
+
+    with_ground_truth = [item for item in results if item.get("iou") is not None]
+    return {
+        "num_samples": len(results),
+        "num_samples_with_gt": len(with_ground_truth),
+        "mean_iou": float(np.mean([item["iou"] for item in with_ground_truth])) if with_ground_truth else None,
+        "mean_f1": float(np.mean([item["f1"] for item in with_ground_truth])) if with_ground_truth else None,
+    }

detectivesam_inference/models/__init__.py

@@ -0,0 +1,3 @@
+from detectivesam_inference.models.forgerylocalizer import ForgeryLocalizer
+
+__all__ = ["ForgeryLocalizer"]

detectivesam_inference/models/adapters.py

@@ -0,0 +1,362 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+FeaturePyramid = list[torch.Tensor]
+StreamPyramid = list[list[torch.Tensor]]
+
+
+class SharedAdapter(nn.Module):
+    """Applies a residual adapter to each feature scale."""
+
+    def __init__(
+        self,
+        in_channels_list: list[int],
+        hidden_dim: int,
+        dropout_rate: float = 0.1,
+        max_streams: int = 2,
+    ) -> None:
+        super().__init__()
+        max_streams = max(max_streams, 1)
+
+        self.mlps_tune = nn.ModuleList(
+            nn.Conv2d(max_streams * channels, hidden_dim, kernel_size=1)
+            for channels in in_channels_list
+        )
+        self.mlps_bottleneck = nn.ModuleList(
+            nn.Sequential(nn.Conv2d(hidden_dim, hidden_dim, kernel_size=1), nn.GELU())
+            for _ in in_channels_list
+        )
+        self.mlp_up = nn.ModuleList(
+            nn.Conv2d(hidden_dim, channels, kernel_size=1)
+            for channels in in_channels_list
+        )
+        self.activation = nn.GELU()
+        self.dropout = nn.Dropout2d(p=dropout_rate)
+
+    def forward(
+        self,
+        stream_features: list[torch.Tensor],
+        unadapted: torch.Tensor,
+        scale_idx: int,
+    ) -> torch.Tensor:
+        fused_streams = torch.cat(stream_features, dim=1) if stream_features else unadapted
+        hidden = self.mlps_tune[scale_idx](fused_streams)
+        hidden = self.activation(hidden)
+        hidden = self.dropout(hidden)
+        hidden = self.mlps_bottleneck[scale_idx](hidden)
+        delta = self.mlp_up[scale_idx](hidden)
+        return unadapted + delta
+
+
+class RefineBlock(nn.Module):
+    """Refines the coarse mask with low-level features."""
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        low_channels: int,
+        out_channels: int = 1,
+        dropout_rate: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(hidden_dim + low_channels, hidden_dim, kernel_size=3, padding=1)
+        self.activation1 = nn.GELU()
+        self.dropout1 = nn.Dropout2d(p=dropout_rate)
+        self.conv2 = nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1)
+        self.activation2 = nn.GELU()
+        self.dropout2 = nn.Dropout2d(p=dropout_rate)
+        self.conv3 = nn.Conv2d(hidden_dim, out_channels, kernel_size=1)
+
+    def forward(
+        self,
+        attention_features: torch.Tensor,
+        low_features: torch.Tensor,
+        coarse_upsampled: torch.Tensor,
+    ) -> torch.Tensor:
+        refined = torch.cat([attention_features, low_features], dim=1)
+        refined = self.conv1(refined)
+        refined = self.activation1(refined)
+        refined = self.dropout1(refined)
+        refined = self.conv2(refined)
+        refined = self.activation2(refined)
+        refined = self.dropout2(refined)
+        delta = self.conv3(refined)
+        return coarse_upsampled + delta
+
+
+class CoarseProcessingBlock(nn.Module):
+    """Adds transformer-based coarse reasoning before refinement."""
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        attn_dim: int,
+        n_heads: int,
+        num_encoder_layers: int,
+        dropout_rate: float,
+        downscale: int,
+    ) -> None:
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.coarse_down = nn.Sequential(
+            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=downscale, stride=downscale, groups=hidden_dim),
+            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=1),
+            nn.GELU(),
+            nn.Dropout2d(p=dropout_rate),
+        )
+        self.pos_embed_conv = nn.Conv2d(2, hidden_dim, kernel_size=1)
+        self.pos_dropout = nn.Dropout2d(p=dropout_rate)
+        self.feat_proj = nn.Sequential(
+            nn.Linear(hidden_dim, attn_dim),
+            nn.Dropout(p=dropout_rate),
+        )
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=attn_dim,
+            nhead=n_heads,
+            dim_feedforward=attn_dim * 4,
+            dropout=dropout_rate,
+            activation="gelu",
+            batch_first=True,
+        )
+        self.transformer_encoder = nn.TransformerEncoder(
+            encoder_layer,
+            num_layers=num_encoder_layers,
+        )
+        self.transformer_out = nn.Sequential(
+            nn.Linear(attn_dim, hidden_dim),
+            nn.Dropout(p=dropout_rate),
+        )
+        self.residual_gate_conv = nn.Sequential(
+            nn.Conv2d(hidden_dim * 2, hidden_dim // 4, kernel_size=3, padding=1),
+            nn.GELU(),
+            nn.Dropout2d(p=dropout_rate),
+            nn.Conv2d(hidden_dim // 4, 1, kernel_size=1),
+        )
+        self.cached_pos_encodings: dict[tuple[int, int], torch.Tensor] = {}
+
+    def _generate_pos_encoding(self, height: int, width: int) -> torch.Tensor:
+        device = self.pos_embed_conv.weight.device
+        y_pos = torch.linspace(-1, 1, height, device=device).view(height, 1).expand(height, width)
+        x_pos = torch.linspace(-1, 1, width, device=device).view(1, width).expand(height, width)
+        pos_grid = torch.stack([y_pos, x_pos], dim=0).unsqueeze(0)
+        return self.pos_embed_conv(pos_grid)
+
+    def _get_positional_encoding(self, batch_size: int, height: int, width: int) -> torch.Tensor:
+        key = (height, width)
+        device = self.pos_embed_conv.weight.device
+        if key not in self.cached_pos_encodings:
+            self.cached_pos_encodings[key] = self._generate_pos_encoding(height, width).detach()
+
+        cached_encoding = self.cached_pos_encodings[key]
+        if cached_encoding.device != device:
+            cached_encoding = cached_encoding.to(device)
+            self.cached_pos_encodings[key] = cached_encoding
+        return cached_encoding.expand(batch_size, -1, -1, -1)
+
+    def forward(self, fused: torch.Tensor) -> torch.Tensor:
+        coarse_features = self.coarse_down(fused)
+        batch_size, _, height, width = coarse_features.shape
+
+        pos_embed = self._get_positional_encoding(batch_size, height, width)
+        pos_embed = self.pos_dropout(pos_embed)
+        coarse_with_position = coarse_features + pos_embed
+
+        feature_sequence = coarse_with_position.flatten(2).permute(0, 2, 1)
+        feature_sequence = self.feat_proj(feature_sequence)
+        transformer_output = self.transformer_encoder(feature_sequence)
+        hidden = self.transformer_out(transformer_output)
+        hidden = hidden.permute(0, 2, 1).view(batch_size, self.hidden_dim, height, width)
+
+        gate_input = torch.cat([hidden, coarse_features], dim=1)
+        residual_gate = torch.sigmoid(self.residual_gate_conv(gate_input))
+        return residual_gate * hidden + (1 - residual_gate) * coarse_features
+
+
+class FineProcessingBlock(nn.Module):
+    """Produces the coarse mask and uncertainty map."""
+
+    def __init__(self, hidden_dim: int, dropout_rate: float) -> None:
+        super().__init__()
+        self.feature_refinement = nn.Sequential(
+            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1),
+            nn.GELU(),
+            nn.Dropout2d(p=dropout_rate),
+            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1),
+            nn.GELU(),
+            nn.Dropout2d(p=dropout_rate),
+        )
+        self.coarse_head = nn.Conv2d(hidden_dim, 1, kernel_size=1)
+        self.uncertainty_head = nn.Conv2d(hidden_dim, 1, kernel_size=1)
+
+    def forward(
+        self,
+        hidden: torch.Tensor,
+        output_size: tuple[int, int],
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        hidden = self.feature_refinement(hidden)
+        coarse_logit = self.coarse_head(hidden)
+        uncertainty_logit = self.uncertainty_head(hidden)
+
+        coarse_mask = F.interpolate(
+            coarse_logit,
+            size=output_size,
+            mode="bilinear",
+            align_corners=False,
+        )
+        uncertainty_map = F.interpolate(
+            uncertainty_logit,
+            size=output_size,
+            mode="bilinear",
+            align_corners=False,
+        )
+        return hidden, coarse_mask, torch.sigmoid(uncertainty_map)
+
+
+class FeatureFusionBlockSpatial(nn.Module):
+    """Fuses original, adapted, and perturbed features with per-pixel attention."""
+
+    def __init__(
+        self,
+        in_channels_list: list[int],
+        hidden_dim: int = 128,
+        dropout_rate: float = 0.1,
+        max_streams: int = 2,
+        attn_reduction: int = 4,
+    ) -> None:
+        super().__init__()
+        self.num_streams = 2 + max_streams
+        self.att_conv = nn.ModuleList()
+        self.proj_conv = nn.ModuleList()
+
+        for channels in in_channels_list:
+            total_channels = channels * self.num_streams
+            mid_channels = max(total_channels // attn_reduction, 8)
+            self.att_conv.append(
+                nn.Sequential(
+                    nn.Conv2d(
+                        total_channels,
+                        mid_channels,
+                        kernel_size=3,
+                        padding=1,
+                        groups=self.num_streams,
+                        bias=False,
+                    ),
+                    nn.GELU(),
+                    nn.Conv2d(mid_channels, self.num_streams, kernel_size=1, bias=False),
+                )
+            )
+            self.proj_conv.append(
+                nn.Sequential(
+                    nn.Conv2d(channels, hidden_dim, kernel_size=1),
+                    nn.GELU(),
+                    nn.Dropout2d(p=dropout_rate),
+                )
+            )
+
+        fusion_channels = hidden_dim * len(in_channels_list)
+        self.fuse_project = nn.Sequential(
+            nn.Conv2d(fusion_channels, hidden_dim, kernel_size=1),
+            nn.GELU(),
+            nn.Dropout2d(p=dropout_rate),
+        )
+
+    def forward(
+        self,
+        adapted: FeaturePyramid,
+        unadapted: FeaturePyramid,
+        streams_unadapted: StreamPyramid,
+        output_size: tuple[int, int],
+    ) -> torch.Tensor:
+        fused_scales = []
+        for scale_idx, (att_head, projection) in enumerate(zip(self.att_conv, self.proj_conv)):
+            streams = [adapted[scale_idx], unadapted[scale_idx], *streams_unadapted[scale_idx]]
+            logits = att_head(torch.cat(streams, dim=1))
+            weights = F.softmax(logits, dim=1).unsqueeze(2)
+            fused = (torch.stack(streams, dim=1) * weights).sum(dim=1)
+            fused = projection(fused)
+            fused = F.interpolate(
+                fused,
+                size=output_size,
+                mode="bilinear",
+                align_corners=False,
+            )
+            fused_scales.append(fused)
+
+        return self.fuse_project(torch.cat(fused_scales, dim=1))
+
+
+class MaskAdapter(nn.Module):
+    """Builds the prompt mask passed into the SAM decoder."""
+
+    def __init__(
+        self,
+        hidden_dim: int = 256,
+        downscale: int = 16,
+        output_resolution: tuple[int, int] = (128, 128),
+        in_channels_list: list[int] | None = None,
+        attn_dim: int = 16,
+        n_heads: int = 4,
+        num_encoder_layers: int = 2,
+        dropout_rate: float = 0.1,
+        max_streams: int = 2,
+    ) -> None:
+        super().__init__()
+        channels = in_channels_list or [256, 32, 64]
+        self.downscale = downscale
+        self.output_resolution = output_resolution
+
+        self.feature_fusion = FeatureFusionBlockSpatial(
+            in_channels_list=channels,
+            hidden_dim=hidden_dim,
+            dropout_rate=dropout_rate,
+            max_streams=max_streams,
+        )
+        self.coarse_processor = CoarseProcessingBlock(
+            hidden_dim=hidden_dim,
+            attn_dim=attn_dim,
+            n_heads=n_heads,
+            num_encoder_layers=num_encoder_layers,
+            dropout_rate=dropout_rate,
+            downscale=downscale,
+        )
+        self.fine_processor = FineProcessingBlock(hidden_dim, dropout_rate)
+        self.spatial_gate = nn.Sequential(
+            nn.Conv2d(2, hidden_dim // 2, kernel_size=3, padding=1),
+            nn.GELU(),
+            nn.Dropout2d(p=dropout_rate),
+            nn.Conv2d(hidden_dim // 2, 1, kernel_size=1),
+            nn.Sigmoid(),
+        )
+        self.refine_head = RefineBlock(
+            hidden_dim=hidden_dim,
+            low_channels=32,
+            out_channels=1,
+            dropout_rate=dropout_rate,
+        )
+
+    def forward(
+        self,
+        adapted: FeaturePyramid,
+        streams_unadapted: StreamPyramid,
+        unadapted: FeaturePyramid,
+    ) -> torch.Tensor:
+        output_height, output_width = self.output_resolution
+        output_size = (output_height, output_width)
+        coarse_size = (output_height // self.downscale, output_width // self.downscale)
+
+        fused = self.feature_fusion(adapted, unadapted, streams_unadapted, output_size)
+        hidden = self.coarse_processor(fused)
+        if hidden.shape[-2:] != coarse_size:
+            hidden = F.adaptive_avg_pool2d(hidden, coarse_size)
+
+        hidden, coarse_mask, uncertainty_map = self.fine_processor(hidden, output_size)
+        attention_features = F.interpolate(hidden, size=output_size, mode="bilinear", align_corners=False)
+        low_features = F.interpolate(unadapted[1], size=output_size, mode="bilinear", align_corners=False)
+        refined_mask = self.refine_head(attention_features, low_features, coarse_mask)
+
+        spatial_gate = self.spatial_gate(torch.cat([coarse_mask, uncertainty_map], dim=1))
+        return spatial_gate * refined_mask + (1 - spatial_gate) * coarse_mask

detectivesam_inference/models/forgerylocalizer.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from sam2.build_sam import build_sam2
+from sam2.modeling.backbones.image_encoder import ImageEncoder
+from sam2.modeling.sam.mask_decoder import MaskDecoder
+from sam2.modeling.sam.prompt_encoder import PromptEncoder
+from sam2.modeling.sam2_base import SAM2Base
+from sam2.utils.transforms import SAM2Transforms
+
+from detectivesam_inference.models.adapters import MaskAdapter, SharedAdapter
+
+
+MODEL_CHANNELS = [256, 32, 64]
+MASK_ADAPTER_RESOLUTION = (128, 128)
+FeaturePyramid = list[torch.Tensor]
+StreamPyramid = list[list[torch.Tensor]]
+
+
+class ForgeryLocalizer(nn.Module):
+    """Inference-only DetectiveSAM model."""
+
+    def __init__(
+        self,
+        sam_config: str,
+        sam_checkpoint: str,
+        prompt_dim: int = 256,
+        output_resolution: tuple[int, int] = (512, 512),
+        downscale: int = 4,
+        dropout_rate: float = 0.1,
+        max_streams: int = 2,
+        device: str = "cpu",
+    ) -> None:
+        super().__init__()
+        self.output_resolution = output_resolution
+
+        sam = self._build_sam(
+            sam_config=sam_config,
+            sam_checkpoint=sam_checkpoint,
+            output_resolution=output_resolution,
+            device=device,
+        )
+        self.no_mem_embed = sam.no_mem_embed if hasattr(sam, "no_mem_embed") else None
+        self.encoder: ImageEncoder = sam.image_encoder
+        self.decoder: MaskDecoder = sam.sam_mask_decoder
+        self.sam_prompt_encoder: PromptEncoder = sam.sam_prompt_encoder
+        self.sam_prompt_encoder.image_embedding_size = MASK_ADAPTER_RESOLUTION
+
+        self._freeze_sam_components()
+        self.adapters = SharedAdapter(
+            hidden_dim=prompt_dim,
+            in_channels_list=MODEL_CHANNELS,
+            dropout_rate=dropout_rate,
+            max_streams=max_streams,
+        )
+        self.mask_adapter = MaskAdapter(
+            hidden_dim=prompt_dim,
+            output_resolution=MASK_ADAPTER_RESOLUTION,
+            downscale=downscale,
+            in_channels_list=MODEL_CHANNELS,
+            dropout_rate=dropout_rate,
+            max_streams=max_streams,
+        )
+        self.transforms = SAM2Transforms(resolution=sam.image_size, mask_threshold=0.0)
+
+    @staticmethod
+    def _build_sam(
+        *,
+        sam_config: str,
+        sam_checkpoint: str,
+        output_resolution: tuple[int, int],
+        device: str,
+    ) -> SAM2Base:
+        sam: SAM2Base = build_sam2(sam_config, sam_checkpoint, device=device)
+        sam.image_size = output_resolution[0]
+        return sam
+
+    def _freeze_sam_components(self) -> None:
+        for module in (self.encoder, self.decoder, self.sam_prompt_encoder):
+            for parameter in module.parameters():
+                parameter.requires_grad = False
+
+    def _project_sam_features(self, backbone_features: FeaturePyramid) -> tuple[torch.Tensor, FeaturePyramid]:
+        image_embeddings = backbone_features[-1]
+        if self.no_mem_embed is not None:
+            image_embeddings = image_embeddings + self.no_mem_embed.reshape(1, 256, 1, 1).detach()
+
+        high_res_features = [
+            self.decoder.conv_s0(backbone_features[0]),
+            self.decoder.conv_s1(backbone_features[1]),
+        ]
+        return image_embeddings, high_res_features
+
+    def _encode_original_and_streams(
+        self,
+        orig: torch.Tensor,
+        streams: list[torch.Tensor],
+    ) -> tuple[FeaturePyramid, StreamPyramid]:
+        orig_backbone_features = self.encoder(orig)["backbone_fpn"]
+        orig_image_embeddings, orig_high_res_features = self._project_sam_features(orig_backbone_features)
+
+        stream_image_embeddings: list[torch.Tensor] = []
+        stream_high_res_features_0: list[torch.Tensor] = []
+        stream_high_res_features_1: list[torch.Tensor] = []
+        for stream in streams:
+            stream_backbone_features = self.encoder(stream)["backbone_fpn"]
+            stream_image_embedding, stream_high_res_features = self._project_sam_features(stream_backbone_features)
+            stream_image_embeddings.append(stream_image_embedding)
+            stream_high_res_features_0.append(stream_high_res_features[0])
+            stream_high_res_features_1.append(stream_high_res_features[1])
+
+        unadapted = [orig_image_embeddings, orig_high_res_features[0], orig_high_res_features[1]]
+        streams_unadapted = [stream_image_embeddings, stream_high_res_features_0, stream_high_res_features_1]
+        return unadapted, streams_unadapted
+
+    def _apply_adapters(
+        self,
+        unadapted: FeaturePyramid,
+        streams_unadapted: StreamPyramid,
+    ) -> FeaturePyramid:
+        return [
+            self.adapters(streams_unadapted[scale_idx], unadapted[scale_idx], scale_idx)
+            for scale_idx in range(len(unadapted))
+        ]
+
+    def _prepare_decoder_inputs(
+        self,
+        adapted: FeaturePyramid,
+        mask_prompt: torch.Tensor,
+    ) -> tuple[torch.Tensor, FeaturePyramid, torch.Tensor]:
+        # Preserve the original interpolation steps so inference stays bit-exact.
+        image_embeddings = F.interpolate(
+            adapted[0],
+            size=adapted[0].shape[-2:],
+            mode="bilinear",
+            align_corners=False,
+        )
+        high_res_features = [
+            F.interpolate(
+                feature,
+                size=feature.shape[-2:],
+                mode="bilinear",
+                align_corners=False,
+            )
+            for feature in adapted[1:]
+        ]
+        mask_prompt = F.interpolate(
+            mask_prompt,
+            size=mask_prompt.shape[-2:],
+            mode="nearest",
+        )
+        return image_embeddings, high_res_features, mask_prompt
+
+    def forward(self, orig: torch.Tensor, streams: list[torch.Tensor]) -> torch.Tensor:
+        unadapted, streams_unadapted = self._encode_original_and_streams(orig, streams)
+        adapted = self._apply_adapters(unadapted, streams_unadapted)
+
+        mask_prompt = self.mask_adapter(adapted, streams_unadapted, unadapted)
+        image_embeddings, high_res_features, mask_prompt = self._prepare_decoder_inputs(adapted, mask_prompt)
+
+        self.sam_prompt_encoder.image_embedding_size = image_embeddings.shape[-2:]
+        sparse_prompt_embeddings, dense_prompt_embeddings = self.sam_prompt_encoder(
+            points=None,
+            boxes=None,
+            masks=mask_prompt,
+        )
+
+        mask_logits, _, _, _ = self.decoder(
+            image_embeddings=image_embeddings,
+            image_pe=self.sam_prompt_encoder.get_dense_pe(),
+            sparse_prompt_embeddings=sparse_prompt_embeddings,
+            dense_prompt_embeddings=dense_prompt_embeddings,
+            multimask_output=False,
+            repeat_image=False,
+            high_res_features=high_res_features,
+        )
+        return self.transforms.postprocess_masks(mask_logits, torch.Size(self.output_resolution))

detectivesam_inference/perturbations.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import io
+
+import torch
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from PIL import Image
+
+
+def create_spatial_gaussian_kernel(
+    kernel_size: int,
+    sigma: float,
+    channels: int,
+    device: torch.device,
+) -> torch.Tensor:
+    radius = kernel_size // 2
+    coords = torch.arange(-radius, radius + 1, device=device, dtype=torch.float32)
+    x_pos, y_pos = torch.meshgrid(coords, coords, indexing="ij")
+    gaussian = torch.exp(-(x_pos**2 + y_pos**2) / (2 * sigma**2))
+    gaussian /= gaussian.sum()
+    return gaussian.view(1, 1, kernel_size, kernel_size).repeat(channels, 1, 1, 1)
+
+
+def apply_spatial_gaussian_blur(
+    video_tensor: torch.Tensor,
+    sigma: float,
+    kernel_sizes: tuple[int, ...] = (5,),
+) -> torch.Tensor:
+    batch_size, channels, frames, height, width = video_tensor.shape
+    blurred = video_tensor
+    for kernel_size in kernel_sizes:
+        kernel = create_spatial_gaussian_kernel(kernel_size, sigma, channels, video_tensor.device)
+        padding = kernel_size // 2
+        reshaped = blurred.permute(0, 2, 1, 3, 4).reshape(batch_size * frames, channels, height, width)
+        reshaped = F.conv2d(reshaped, kernel, padding=padding, groups=channels)
+        blurred = reshaped.reshape(batch_size, frames, channels, height, width).permute(0, 2, 1, 3, 4)
+    return blurred
+
+
+def add_gaussian_noise_deterministic(
+    image: torch.Tensor,
+    *,
+    mean: float = 0.0,
+    std: float = 0.1,
+    seed: int,
+) -> torch.Tensor:
+    generator = torch.Generator(device=image.device)
+    generator.manual_seed(seed)
+    noise = torch.randn(image.shape, generator=generator, device=image.device, dtype=image.dtype) * std + mean
+    return torch.clamp(image + noise, 0.0, 1.0)
+
+
+def apply_blur_to_image_tensor(image_tensor: torch.Tensor, sigma: float) -> torch.Tensor:
+    if sigma <= 0:
+        return image_tensor.clone()
+    video_tensor = image_tensor.unsqueeze(0).unsqueeze(2)
+    blurred = apply_spatial_gaussian_blur(video_tensor, sigma=sigma, kernel_sizes=(5,))
+    return blurred.squeeze(0).squeeze(1)
+
+
+def apply_jpeg_compression(image: Image.Image, quality: int = 75) -> Image.Image:
+    buffer = io.BytesIO()
+    image.save(buffer, format="JPEG", quality=quality)
+    buffer.seek(0)
+    return Image.open(buffer).convert("RGB")
+
+
+def apply_jpeg_compression_to_tensor(image_tensor: torch.Tensor, quality: int) -> torch.Tensor:
+    image = TF.to_pil_image(image_tensor)
+    return TF.to_tensor(apply_jpeg_compression(image, quality=quality))

detectivesam_inference/predict.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+from detectivesam_inference.dataset import prepare_sample
+from detectivesam_inference.metrics import compute_f1, compute_iou
+from detectivesam_inference.runtime import DetectiveSAMRunner, get_repo_root
+from detectivesam_inference.visualization import save_prediction_outputs
+
+
+DEFAULT_DEMO_NAME = "banana_28809"
+
+
+def resolve_demo_defaults(repo_root: Path) -> tuple[Path | None, Path, Path | None, str]:
+    user_demo_target = repo_root / "demo" / "user_image" / "demo_input.png"
+    fallback_source = repo_root / "demo" / "cocoglide" / "source" / f"{DEFAULT_DEMO_NAME}.png"
+    fallback_target = repo_root / "demo" / "cocoglide" / "target" / f"{DEFAULT_DEMO_NAME}.png"
+    fallback_mask = repo_root / "demo" / "cocoglide" / "mask" / f"{DEFAULT_DEMO_NAME}.png"
+
+    if user_demo_target.exists():
+        return None, user_demo_target, None, "single_image"
+    return fallback_source, fallback_target, fallback_mask, "pair"
+
+
+def parse_args() -> argparse.Namespace:
+    repo_root = get_repo_root()
+    parser = argparse.ArgumentParser(description="Run DetectiveSAM on a single source/target pair.")
+    parser.add_argument(
+        "--checkpoint",
+        default="detective_sam",
+        help="Checkpoint path or alias. Built-in aliases: detective_sam, detective_sam_sota.",
+    )
+    parser.add_argument("--source", default=None, help="Optional source image. If omitted, target is reused as source.")
+    parser.add_argument(
+        "--target",
+        default=None,
+        help="Target image. If omitted, uses demo/user_image/demo_input.png when present, else falls back to the bundled CocoGlide pair.",
+    )
+    parser.add_argument("--mask", default=None, help="Optional ground-truth mask for metrics.")
+    parser.add_argument("--output-dir", default=str(repo_root / "outputs" / "predict_demo"))
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--threshold", type=float, default=0.5)
+    return parser.parse_args()
+
+
+def resolve_input_paths(args: argparse.Namespace, repo_root: Path) -> tuple[Path, Path, Path | None, str]:
+    demo_source, demo_target, demo_mask, demo_mode = resolve_demo_defaults(repo_root)
+    if args.target is not None:
+        target_path = Path(args.target)
+        source_path = Path(args.source) if args.source else target_path
+        mask_path = Path(args.mask) if args.mask else None
+        return source_path, target_path, mask_path, "custom"
+
+    target_path = demo_target
+    source_path = Path(args.source) if args.source else (demo_source or target_path)
+    mask_path = Path(args.mask) if args.mask else demo_mask
+    return source_path, target_path, mask_path, demo_mode
+
+
+def main() -> None:
+    args = parse_args()
+    repo_root = get_repo_root()
+    source_path, target_path, mask_path, demo_mode = resolve_input_paths(args, repo_root)
+    reference_mode = "pair" if source_path != target_path else "target_as_source"
+
+    runner = DetectiveSAMRunner(checkpoint_path=args.checkpoint, device=args.device)
+    sample = prepare_sample(
+        source_path=source_path,
+        target_path=target_path,
+        mask_path=mask_path,
+        img_size=runner.config.img_size,
+        perturbation_type=runner.config.perturbation_type,
+        perturbation_intensity=runner.config.perturbation_intensity,
+    )
+    prediction = runner.predict_sample(sample, threshold=args.threshold)
+
+    gt_mask = sample.mask.squeeze().numpy().astype("uint8") if sample.mask is not None else None
+    summary = {
+        "sample": sample.name,
+        "checkpoint": str(runner.checkpoint_path.resolve()),
+        "demo_mode": demo_mode,
+        "reference_mode": reference_mode,
+        "source": str(sample.source_path),
+        "target": str(sample.target_path),
+        "mask": str(sample.mask_path) if sample.mask_path is not None else None,
+        "threshold": args.threshold,
+        "metrics": {
+            "iou": compute_iou(prediction.pred_mask, gt_mask) if gt_mask is not None else None,
+            "f1": compute_f1(prediction.pred_mask, gt_mask) if gt_mask is not None else None,
+        },
+    }
+
+    output_dir = Path(args.output_dir)
+    save_prediction_outputs(
+        output_dir=output_dir,
+        name=sample.name,
+        source_image=sample.source_image,
+        target_image=sample.target_image,
+        probability_map=prediction.probability,
+        pred_mask=prediction.pred_mask,
+        gt_mask=gt_mask,
+    )
+    output_dir.mkdir(parents=True, exist_ok=True)
+    with (output_dir / f"{sample.name}_summary.json").open("w", encoding="utf-8") as handle:
+        json.dump(summary, handle, indent=2)
+
+    print(json.dumps(summary, indent=2))
+
+
+if __name__ == "__main__":
+    main()

detectivesam_inference/runtime.py

@@ -0,0 +1,102 @@
+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,
+            downscale=self.config.downscale,
+            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 = checkpoint["model"] if isinstance(checkpoint, dict) and "model" in checkpoint else checkpoint
+        model.load_state_dict(state_dict)
+        model.eval()
+        return model
+
+    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)

detectivesam_inference/visualization.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+from PIL import Image, ImageOps
+
+
+def ensure_uint8(mask: np.ndarray) -> np.ndarray:
+    if mask.dtype == np.bool_:
+        return mask.astype(np.uint8) * 255
+    if mask.max() <= 1.0:
+        return (mask * 255).astype(np.uint8)
+    return mask.astype(np.uint8)
+
+
+def mask_to_image(mask: np.ndarray) -> Image.Image:
+    return Image.fromarray(ensure_uint8(mask), mode="L")
+
+
+def probability_to_image(probability: np.ndarray) -> Image.Image:
+    clipped = np.clip(probability, 0.0, 1.0)
+    return Image.fromarray((clipped * 255).astype(np.uint8), mode="L")
+
+
+def overlay_mask(
+    image: Image.Image,
+    mask: np.ndarray,
+    color: tuple[int, int, int],
+    alpha: float = 0.45,
+) -> Image.Image:
+    base = np.array(image.convert("RGB"), dtype=np.float32)
+    overlay = base.copy()
+    overlay[mask.astype(bool)] = (1.0 - alpha) * overlay[mask.astype(bool)] + alpha * np.array(color, dtype=np.float32)
+    return Image.fromarray(np.clip(overlay, 0, 255).astype(np.uint8), mode="RGB")
+
+
+def concat_images(images: list[Image.Image]) -> Image.Image:
+    widths, heights = zip(*(image.size for image in images))
+    canvas = Image.new("RGB", (sum(widths), max(heights)), color=(255, 255, 255))
+    x_offset = 0
+    for image in images:
+        canvas.paste(image.convert("RGB"), (x_offset, 0))
+        x_offset += image.width
+    return canvas
+
+
+def save_prediction_outputs(
+    output_dir: str | Path,
+    name: str,
+    source_image: Image.Image,
+    target_image: Image.Image,
+    probability_map: np.ndarray,
+    pred_mask: np.ndarray,
+    gt_mask: np.ndarray | None = None,
+) -> None:
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    probability_image = probability_to_image(probability_map)
+    pred_mask_image = mask_to_image(pred_mask)
+    pred_overlay = overlay_mask(target_image, pred_mask, color=(255, 0, 0))
+
+    comparison_images = [
+        source_image.convert("RGB"),
+        target_image.convert("RGB"),
+        ImageOps.colorize(probability_image, black="black", white="white").convert("RGB"),
+        pred_overlay,
+    ]
+
+    probability_image.save(output_dir / f"{name}_probability.png")
+    pred_mask_image.save(output_dir / f"{name}_pred_mask.png")
+    pred_overlay.save(output_dir / f"{name}_pred_overlay.png")
+
+    if gt_mask is not None:
+        gt_mask_image = mask_to_image(gt_mask)
+        gt_overlay = overlay_mask(target_image, gt_mask, color=(0, 255, 0))
+        gt_mask_image.save(output_dir / f"{name}_gt_mask.png")
+        gt_overlay.save(output_dir / f"{name}_gt_overlay.png")
+        comparison_images.append(gt_overlay)
+
+    concat_images(comparison_images).save(output_dir / f"{name}_comparison.png")

pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+pythonpath = .

requirements.txt

@@ -0,0 +1,9 @@
+torch>=2.4
+torchvision>=0.19
+numpy>=1.24
+Pillow>=10.0
+opencv-python-headless>=4.10
+hydra-core>=1.3
+sam2
+tqdm>=4.66
+

sam2configs/sam2.1_hiera_b+.yaml

@@ -0,0 +1,116 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 1
+    trunk:
+      _target_: sam2.modeling.backbones.hieradet.Hiera
+      embed_dim: 112
+      num_heads: 2
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [896, 448, 224, 112]
+      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+      fpn_interp_model: nearest
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  no_obj_embed_spatial: true
+  # use high-resolution feature map in the SAM mask decoder
+  use_high_res_features_in_sam: true
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: true
+  proj_tpos_enc_in_obj_ptrs: true
+  use_signed_tpos_enc_to_obj_ptrs: true
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  compile_image_encoder: False

sam2configs/sam2.1_hiera_base_plus.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a2345aede8715ab1d5d31b4a509fb160c5a4af1970f199d9054ccfb746c004c5
+size 323606802

tests/test_regression.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import pytest
+
+from detectivesam_inference.checkpoint import resolve_checkpoint_path
+from detectivesam_inference.dataset import PairDataset, prepare_sample
+from detectivesam_inference.metrics import compute_f1, compute_iou, summarize_results
+from detectivesam_inference.runtime import DetectiveSAMRunner, get_repo_root
+
+
+def assert_exact(value: float | None, expected: float | None, *, abs_tol: float = 1e-12) -> None:
+    assert value is not None
+    assert expected is not None
+    assert math.isclose(value, expected, rel_tol=0.0, abs_tol=abs_tol)
+
+
+@pytest.fixture(scope="module")
+def repo_root() -> Path:
+    return get_repo_root()
+
+
+@pytest.fixture(scope="module")
+def baseline_runner() -> DetectiveSAMRunner:
+    return DetectiveSAMRunner(checkpoint_path="detective_sam", device="cpu")
+
+
+@pytest.fixture(scope="module")
+def sota_runner() -> DetectiveSAMRunner:
+    return DetectiveSAMRunner(checkpoint_path="detective_sam_sota", device="cpu")
+
+
+def predict_metrics(
+    runner: DetectiveSAMRunner,
+    *,
+    source_path: Path,
+    target_path: Path,
+    mask_path: Path,
+) -> tuple[float, float]:
+    sample = prepare_sample(
+        source_path=source_path,
+        target_path=target_path,
+        mask_path=mask_path,
+        img_size=runner.config.img_size,
+        perturbation_type=runner.config.perturbation_type,
+        perturbation_intensity=runner.config.perturbation_intensity,
+    )
+    prediction = runner.predict_sample(sample, threshold=0.5)
+    true_mask = sample.mask.squeeze().numpy().astype("uint8")
+    return compute_iou(prediction.pred_mask, true_mask), compute_f1(prediction.pred_mask, true_mask)
+
+
+def test_checkpoint_alias_resolution(repo_root: Path) -> None:
+    assert resolve_checkpoint_path("detective_sam", repo_root) == repo_root / "checkpoints" / "model_epoch22_batch999_score1.1114.pth"
+    assert resolve_checkpoint_path("detective_sam_sota", repo_root) == repo_root / "checkpoints" / "detective_sam_sota.pth"
+
+
+def test_baseline_banana_demo_metrics(repo_root: Path, baseline_runner: DetectiveSAMRunner) -> None:
+    demo_root = repo_root / "demo" / "cocoglide"
+    iou, f1 = predict_metrics(
+        baseline_runner,
+        source_path=demo_root / "source" / "banana_28809.png",
+        target_path=demo_root / "target" / "banana_28809.png",
+        mask_path=demo_root / "mask" / "banana_28809.png",
+    )
+    assert_exact(iou, 0.8566427949370513)
+    assert_exact(f1, 0.9227868680750683)
+
+
+def test_sota_flux_demo_metrics(repo_root: Path, sota_runner: DetectiveSAMRunner) -> None:
+    demo_root = repo_root / "demo" / "flux_test"
+    iou, f1 = predict_metrics(
+        sota_runner,
+        source_path=demo_root / "source" / "548.png",
+        target_path=demo_root / "target" / "548.png",
+        mask_path=demo_root / "mask" / "548.png",
+    )
+    assert_exact(iou, 0.8703024868799283)
+    assert_exact(f1, 0.9306542583192329)
+
+
+def test_sota_qwen_demo_metrics(repo_root: Path, sota_runner: DetectiveSAMRunner) -> None:
+    demo_root = repo_root / "demo" / "qwen_test"
+    iou, f1 = predict_metrics(
+        sota_runner,
+        source_path=demo_root / "source" / "166.png",
+        target_path=demo_root / "target" / "166.png",
+        mask_path=demo_root / "mask" / "166.png",
+    )
+    assert_exact(iou, 0.8297306693388413)
+    assert_exact(f1, 0.9069429542203147)
+
+
+def test_baseline_cocoglide_eval_summary(repo_root: Path, baseline_runner: DetectiveSAMRunner) -> None:
+    dataset = PairDataset(
+        root_dir=repo_root / "demo" / "cocoglide",
+        img_size=baseline_runner.config.img_size,
+        perturbation_type=baseline_runner.config.perturbation_type,
+        perturbation_intensity=baseline_runner.config.perturbation_intensity,
+    )
+
+    per_sample_results: list[dict[str, float | str | None]] = []
+    for sample in dataset:
+        prediction = baseline_runner.predict_sample(sample, threshold=0.5)
+        true_mask = sample.mask.squeeze().numpy().astype("uint8")
+        per_sample_results.append(
+            {
+                "name": sample.name,
+                "iou": compute_iou(prediction.pred_mask, true_mask),
+                "f1": compute_f1(prediction.pred_mask, true_mask),
+            }
+        )
+
+    summary = summarize_results(per_sample_results)
+    assert summary["num_samples"] == 5
+    assert summary["num_samples_with_gt"] == 5
+    assert_exact(summary["mean_iou"], 0.5092573070481035)
+    assert_exact(summary["mean_f1"], 0.6509390858765342)
+
+    expected_by_name = {
+        "airplane_139871": (0.41829717560376584, 0.5898582931686339),
+        "banana_28809": (0.8566427949370513, 0.9227868680750683),
+        "giraffe_296969": (0.22833093957714018, 0.3717743031951054),
+        "train_221213": (0.547253866814856, 0.7073872989458688),
+        "tv_453722": (0.49576175830770386, 0.662888665997994),
+    }
+    for result in per_sample_results:
+        expected_iou, expected_f1 = expected_by_name[result["name"]]
+        assert_exact(result["iou"], expected_iou)
+        assert_exact(result["f1"], expected_f1)