Upload clean DetectiveSAMv2 inference bundle

.gitattributes

@@ -1,35 +1,10 @@
-*.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/banana_28809.png filter=lfs diff=lfs merge=lfs -text
+demo/cocoglide/source/train_221213.png filter=lfs diff=lfs merge=lfs -text
+demo/cocoglide/target/banana_28809.png filter=lfs diff=lfs merge=lfs -text
+demo/cocoglide/target/train_221213.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,9 @@
+__pycache__/
+*.pyc
+.DS_Store
+.venv/
+venv/
+.pytest_cache/
+outputs/
+checkpoints/*.pth
+sam2configs/sam2.1_hiera_base_plus.pt

README.md

@@ -0,0 +1,177 @@
+# DetectiveSAMv2
+
+DetectiveSAMv2 is an inference-only image forgery localization bundle built around SAM2. The full runnable bundle with model weights is hosted on Hugging Face:
+
+- https://huggingface.co/Gertlek/DetectiveSAMv2
+
+## What is bundled
+
+- Inference checkpoint configs under `checkpoints/`
+- SAM2 config 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_v2`
+- `detective_sam`
+- `detective_sam_sota`
+
+The v2 bundle ships `detective_sam_v2`. The v1 aliases are retained for compatibility if those older checkpoint files are placed under `checkpoints/`.
+
+## Architecture
+
+This branch ports the v2 inference architecture from the training workspace while keeping the package layout, CLIs, and regression-test structure intact.
+
+Main changes:
+
+- Evidence features built from unadapted target features, raw perturbation-stream features, and adapted target deltas.
+- Spatial cross-attention feature adapters for flexible information sharing across target and perturbation streams.
+- Transformer evidence mask adapter for direct prompt-mask prediction.
+- V1 and V2 adapter code split into `adapters_v1.py` and `adapters_v2.py`; `adapters.py` remains a compatibility re-export.
+- Three perturbation streams by default: Gaussian blur, JPEG compression, and Gaussian noise.
+- JSON and YAML checkpoint sidecar support, including training checkpoint containers with `model` or `ema` keys.
+
+The `detective_sam_v2` alias expects a checkpoint at `checkpoints/detective_sam_v2.pth` with the sidecar config included in this repo.
+
+### Benchmark Scores
+
+Scores below are IoU / F1 in percent for the 10-epoch v2 checkpoint `dsamv2_best6info_10ep_20260623_060047/best_model.pth`. Paper baseline numbers are from the Detective SAM ICLR 2026 paper tables for the original Detective SAM model.
+
+| Dataset | Split | Samples | DetectiveSAMv2 | Paper Detective SAM | Delta |
+| --- | --- | ---: | ---: | ---: | ---: |
+| MagicBrush | trained on | 528 | 62.73 / 74.38 | 46.48 / 57.55 | +16.25 / +16.83 |
+| SIDA | trained on | 6,911 | 55.34 / 65.72 | 54.55 / 65.29 | +0.79 / +0.43 |
+| UltraEdit 10k | trained on | 10,000 | 67.03 / 76.71 | 27.74 / 35.54 | +39.29 / +41.17 |
+| CoCoGLIDE | not trained on | 512 | 47.17 / 58.83 | 44.74 / 51.50 | +2.43 / +7.33 |
+| AutoSplice | not trained on | 3,621 | 58.39 / 71.40 | 46.90 / 60.30 | +11.49 / +11.10 |
+| NanoBanana | not trained on | 84 | 33.81 / 46.79 | 19.34 / 20.77 | +14.47 / +26.02 |
+| FLUX-Bench | not trained on | 647 | 46.10 / 58.02 | 18.70 / 21.28 | +27.40 / +36.74 |
+| QWEN-Bench | not trained on | 194 | 46.32 / 57.94 | 20.41 / 22.29 | +25.91 / +35.65 |
+
+Grouped results from the same run:
+
+| Group | Samples | Weighted IoU / F1 | Simple IoU / F1 |
+| --- | ---: | ---: | ---: |
+| Trained-on | 17,439 | 62.26 / 72.28 | 61.70 / 72.27 |
+| Not-trained-on | 5,058 | 54.81 / 67.49 | 46.36 / 58.60 |
+| Overall | 22,497 | 60.59 / 71.20 | 52.11 / 63.72 |
+
+The not-trained-on group follows the requested reporting split and includes NanoBanana, CoCoGLIDE, AutoSplice, QWEN-Bench, and FLUX-Bench.
+
+## Setup
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+```
+
+Then download the large weights from Hugging Face:
+
+```bash
+pip install -U huggingface_hub
+hf download Gertlek/DetectiveSAMv2 \
+  checkpoints/detective_sam_v2.pth \
+  sam2configs/sam2.1_hiera_base_plus.pt \
+  --local-dir .
+```
+
+The expected checkpoint paths are also documented in `checkpoints/README.md` and `sam2configs/README.md`.
+
+## Hugging Face Usage
+
+For the simplest setup, clone the Hugging Face repo directly:
+
+```bash
+git lfs install
+git clone https://huggingface.co/Gertlek/DetectiveSAMv2
+cd DetectiveSAMv2
+python -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+python -m detectivesam_inference.predict \
+  --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 \
+  --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 \
+  --output-dir outputs/poster_baseline
+```
+
+### 3. Bundled SOTA examples
+
+Flux example:
+
+```bash
+python -m detectivesam_inference.predict \
+  --checkpoint detective_sam_v2 \
+  --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_v2 \
+  --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
+```
+
+### 4. Bundled CocoGlide subset sweep
+
+Use this to evaluate the bundled banana and train CocoGlide demo pairs.
+
+```bash
+python -m detectivesam_inference.evaluate \
+  --checkpoint detective_sam_v2 \
+  --dataset-root demo/cocoglide \
+  --output-dir outputs/poster_eval_cocoglide \
+  --num-visualizations 2
+```
+
+## 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 YAML or JSON sidecars in `checkpoints/`; you only need the alias or checkpoint path.
+- The Hugging Face repo bundles the v2 `.pth` checkpoint and SAM2 `.pt` weight file.
+- This repo does not include training code or training-only dependencies.

checkpoints/README.md

@@ -0,0 +1,13 @@
+# DetectiveSAMv2 Checkpoints
+
+The Hugging Face bundle ships the v2 `.pth` checkpoint file.
+
+The full runnable bundle is hosted on Hugging Face:
+
+- https://huggingface.co/Gertlek/DetectiveSAMv2
+
+To run inference from a lightweight checkout, download this file from Hugging Face into this directory:
+
+- `detective_sam_v2.pth`
+
+The small YAML sidecar in this directory is kept in Git so the `detective_sam_v2` alias resolves once the `.pth` file is present locally.

checkpoints/detective_sam_v2_params.yaml

@@ -0,0 +1,10 @@
+img_size: 512
+prompt_dim: 96
+downscale: 8
+dropout_rate: 0.1
+perturbation_type: gaussian_blur+jpeg_compression+gaussian_noise
+perturbation_intensity: 0.5
+sam_config_file: sam2.1_hiera_b+.yaml
+sam_checkpoint: sam2configs/sam2.1_hiera_base_plus.pt
+adapter_type: spatial_cross_attention
+mask_adapter_type: transformer

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,161 @@
+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": Path("checkpoints/model_epoch22_batch999_score1.1114.pth"),
+    "detective_sam_sota": Path("checkpoints/detective_sam_sota.pth"),
+    "detective_sam_v2": DEFAULT_CHECKPOINT,
+}
+
+
+@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 = "conv"
+    mask_adapter_type: str = "coarse"
+
+    @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",
+            )
+        ),
+        adapter_type=str(_resolve_param(params, "adapter_type", section="model_config", default="conv")),
+        mask_adapter_type=str(_resolve_param(params, "mask_adapter_type", section="model_config", default="coarse")),
+    )
+
+
+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)

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_v2",
+        help="Checkpoint path or alias. Built-in aliases: detective_sam, detective_sam_sota, detective_sam_v2.",
+    )
+    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,37 @@
+from detectivesam_inference.models.adapters_v1 import (
+    CoarseProcessingBlock,
+    FeatureFusionBlockSpatial,
+    FeaturePyramid,
+    FineProcessingBlock,
+    MaskAdapter,
+    RefineBlock,
+    SharedAdapter,
+    StreamPyramid,
+)
+from detectivesam_inference.models.adapters_v2 import (
+    CrossAttentionSharedAdapter,
+    FlexibleEvidenceBlock,
+    FlexibleEvidenceProjector,
+    SpatialCrossAttentionSharedAdapter,
+    StreamEvidenceBuilder,
+    Transformer2DBlock,
+    TransformerEvidenceMaskAdapter,
+)
+
+__all__ = [
+    "CoarseProcessingBlock",
+    "CrossAttentionSharedAdapter",
+    "FeatureFusionBlockSpatial",
+    "FeaturePyramid",
+    "FineProcessingBlock",
+    "FlexibleEvidenceBlock",
+    "FlexibleEvidenceProjector",
+    "MaskAdapter",
+    "RefineBlock",
+    "SharedAdapter",
+    "SpatialCrossAttentionSharedAdapter",
+    "StreamEvidenceBuilder",
+    "StreamPyramid",
+    "Transformer2DBlock",
+    "TransformerEvidenceMaskAdapter",
+]

detectivesam_inference/models/adapters_v1.py

@@ -0,0 +1,366 @@
+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/adapters_v2.py

@@ -0,0 +1,651 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional
+
+
+FeaturePyramid = list[torch.Tensor]
+StreamPyramid = list[list[torch.Tensor]]
+
+
+def _group_count(channels: int, max_groups: int = 8) -> int:
+    for groups in range(min(max_groups, channels), 0, -1):
+        if channels % groups == 0:
+            return groups
+    return 1
+
+
+class FlexibleEvidenceBlock(nn.Module):
+    def __init__(self, channels: int, dropout_rate: float = 0.1, expansion: int = 2) -> None:
+        super().__init__()
+        expanded = channels * expansion
+        self.norm = nn.GroupNorm(_group_count(channels), channels)
+        self.expand = nn.Conv2d(channels, expanded, kernel_size=1)
+        self.depthwise = nn.Conv2d(expanded, expanded, kernel_size=3, padding=1, groups=expanded)
+        self.act = nn.GELU()
+        self.channel_gate = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(expanded, max(channels // 4, 8), kernel_size=1),
+            nn.GELU(),
+            nn.Conv2d(max(channels // 4, 8), expanded, kernel_size=1),
+            nn.Sigmoid(),
+        )
+        self.project = nn.Conv2d(expanded, channels, kernel_size=1)
+        self.dropout = nn.Dropout2d(p=dropout_rate)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        hidden = self.norm(x)
+        hidden = self.expand(hidden)
+        hidden = self.depthwise(hidden)
+        hidden = self.act(hidden)
+        hidden = hidden * self.channel_gate(hidden)
+        hidden = self.project(hidden)
+        hidden = self.dropout(hidden)
+        return residual + hidden
+
+
+class FlexibleEvidenceProjector(nn.Module):
+    """Channel and local-spatial evidence projector used by v2 adapters."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_dim: int,
+        dropout_rate: float = 0.1,
+        depth: int = 2,
+    ) -> None:
+        super().__init__()
+        self.input_proj = nn.Sequential(
+            nn.Conv2d(in_channels, hidden_dim, kernel_size=1),
+            nn.GroupNorm(_group_count(hidden_dim), hidden_dim),
+            nn.GELU(),
+            nn.Dropout2d(p=dropout_rate),
+        )
+        self.blocks = nn.Sequential(
+            *[
+                FlexibleEvidenceBlock(hidden_dim, dropout_rate=dropout_rate)
+                for _ in range(depth)
+            ]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.blocks(self.input_proj(x))
+
+
+class StreamEvidenceBuilder(nn.Module):
+    """Build target evidence from raw target and perturbation-stream features."""
+
+    def __init__(
+        self,
+        in_channels_list: list[int],
+        hidden_dim: int,
+        dropout_rate: float = 0.1,
+        max_streams: int = 3,
+    ) -> None:
+        super().__init__()
+        self.max_streams = max(1, max_streams)
+        self.projections = nn.ModuleList(
+            [
+                FlexibleEvidenceProjector(
+                    in_channels=channels * (1 + 3 * self.max_streams),
+                    hidden_dim=hidden_dim,
+                    dropout_rate=dropout_rate,
+                    depth=3,
+                )
+                for channels in in_channels_list
+            ]
+        )
+
+    def forward(
+        self,
+        unadapted: FeaturePyramid,
+        streams_unadapted: StreamPyramid,
+    ) -> FeaturePyramid:
+        evidence = []
+        for scale_idx, (target, streams) in enumerate(zip(unadapted, streams_unadapted)):
+            fixed_streams = list(streams[: self.max_streams])
+            while len(fixed_streams) < self.max_streams:
+                fixed_streams.append(torch.zeros_like(target))
+
+            abs_diffs = [(target - stream).abs() for stream in fixed_streams]
+            signed_diffs = [target - stream for stream in fixed_streams]
+            raw = torch.cat([target, *fixed_streams, *abs_diffs, *signed_diffs], dim=1)
+            evidence.append(self.projections[scale_idx](raw))
+        return evidence
+
+
+class CrossAttentionSharedAdapter(nn.Module):
+    """Per-scale cross-attention adapter over perturbation streams."""
+
+    def __init__(
+        self,
+        in_channels_list: list[int],
+        hidden_dim: int,
+        dropout_rate: float = 0.1,
+        max_streams: int = 2,
+        num_heads: int = 4,
+    ) -> None:
+        super().__init__()
+        self.num_scales = len(in_channels_list)
+        self.max_streams = max(1, max_streams)
+        self.num_heads = max(1, num_heads)
+        if hidden_dim % self.num_heads != 0:
+            raise ValueError(f"hidden_dim ({hidden_dim}) must be divisible by num_heads ({self.num_heads})")
+
+        self.head_dim = hidden_dim // self.num_heads
+        self.scale = self.head_dim**-0.5
+        self.q_proj = nn.ModuleList(nn.Conv2d(channels, hidden_dim, kernel_size=1) for channels in in_channels_list)
+        self.k_proj = nn.ModuleList(nn.Conv2d(channels, hidden_dim, kernel_size=1) for channels in in_channels_list)
+        self.v_proj = nn.ModuleList(nn.Conv2d(channels, hidden_dim, kernel_size=1) for channels in in_channels_list)
+        self.out_proj = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.Conv2d(hidden_dim, hidden_dim, kernel_size=1),
+                    nn.GELU(),
+                    nn.Dropout2d(p=dropout_rate),
+                    nn.Conv2d(hidden_dim, channels, kernel_size=1),
+                )
+                for channels in in_channels_list
+            ]
+        )
+        self.evidence_q_proj = nn.ModuleList(
+            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=1)
+            for _ in in_channels_list
+        )
+        self.evidence_context_proj = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.Conv2d(hidden_dim, hidden_dim, kernel_size=1),
+                    nn.GELU(),
+                    nn.Dropout2d(p=dropout_rate),
+                )
+                for _ in in_channels_list
+            ]
+        )
+        self.residual_gate = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.Conv2d(channels + hidden_dim, channels, kernel_size=1),
+                    nn.Sigmoid(),
+                )
+                for channels in in_channels_list
+            ]
+        )
+        self.dropout = nn.Dropout2d(p=dropout_rate)
+
+    def forward(
+        self,
+        stream_features: list[torch.Tensor],
+        unadapted: torch.Tensor,
+        scale_idx: int,
+        return_delta: bool = False,
+        evidence: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if not stream_features:
+            delta = torch.zeros_like(unadapted)
+            adapted = unadapted
+            if return_delta:
+                return adapted, delta
+            return adapted
+
+        batch_size, _, height, width = unadapted.shape
+        stream_count = len(stream_features)
+        query = self.q_proj[scale_idx](unadapted)
+        if evidence is not None:
+            if evidence.shape[-2:] != query.shape[-2:]:
+                evidence = F.interpolate(evidence, size=query.shape[-2:], mode="bilinear", align_corners=False)
+            query = query + self.evidence_q_proj[scale_idx](evidence)
+        query = query.view(batch_size, self.num_heads, self.head_dim, 1, height, width)
+
+        keys = [self.k_proj[scale_idx](feature) for feature in stream_features]
+        values = [self.v_proj[scale_idx](feature) for feature in stream_features]
+        key = torch.stack(keys, dim=1).view(batch_size, stream_count, self.num_heads, self.head_dim, height, width)
+        value = torch.stack(values, dim=1).view(batch_size, stream_count, self.num_heads, self.head_dim, height, width)
+        key = key.permute(0, 2, 3, 1, 4, 5)
+        value = value.permute(0, 2, 3, 1, 4, 5)
+
+        attention_logits = (query * key).sum(dim=2) * self.scale
+        attention = F.softmax(attention_logits, dim=2)
+        context = (attention.unsqueeze(2) * value).sum(dim=3)
+        context = context.reshape(batch_size, self.num_heads * self.head_dim, height, width)
+        if evidence is not None:
+            context = context + self.evidence_context_proj[scale_idx](evidence)
+        context = self.dropout(context)
+
+        raw_delta = self.out_proj[scale_idx](context)
+        gate = self.residual_gate[scale_idx](torch.cat([unadapted, context], dim=1))
+        delta = gate * raw_delta
+        adapted = unadapted + delta
+        if return_delta:
+            return adapted, delta
+        return adapted
+
+
+class SpatialCrossAttentionSharedAdapter(CrossAttentionSharedAdapter):
+    """Cross-attention adapter with global low-res and windowed mid-res spatial mixing."""
+
+    def __init__(
+        self,
+        in_channels_list: list[int],
+        hidden_dim: int,
+        dropout_rate: float = 0.1,
+        max_streams: int = 2,
+        num_heads: int = 4,
+        global_spatial_scales: tuple[int, ...] = (0,),
+        windowed_spatial_scales: tuple[int, ...] = (2,),
+        window_size: int = 8,
+    ) -> None:
+        super().__init__(
+            in_channels_list=in_channels_list,
+            hidden_dim=hidden_dim,
+            dropout_rate=dropout_rate,
+            max_streams=max_streams,
+            num_heads=num_heads,
+        )
+        self.global_spatial_scales = set(global_spatial_scales)
+        self.windowed_spatial_scales = set(windowed_spatial_scales)
+        self.window_size = window_size
+
+    def _apply_residual(
+        self,
+        context: torch.Tensor,
+        unadapted: torch.Tensor,
+        scale_idx: int,
+        return_delta: bool,
+        evidence: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if evidence is not None:
+            if evidence.shape[-2:] != context.shape[-2:]:
+                evidence = F.interpolate(evidence, size=context.shape[-2:], mode="bilinear", align_corners=False)
+            context = context + self.evidence_context_proj[scale_idx](evidence)
+        context = self.dropout(context)
+        raw_delta = self.out_proj[scale_idx](context)
+        gate = self.residual_gate[scale_idx](torch.cat([unadapted, context], dim=1))
+        delta = gate * raw_delta
+        adapted = unadapted + delta
+        if return_delta:
+            return adapted, delta
+        return adapted
+
+    def _global_spatial_context(
+        self,
+        stream_features: list[torch.Tensor],
+        unadapted: torch.Tensor,
+        scale_idx: int,
+        evidence: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        batch_size, _, height, width = unadapted.shape
+        stream_count = len(stream_features)
+        num_tokens = height * width
+        query = self.q_proj[scale_idx](unadapted)
+        if evidence is not None:
+            if evidence.shape[-2:] != query.shape[-2:]:
+                evidence = F.interpolate(evidence, size=query.shape[-2:], mode="bilinear", align_corners=False)
+            query = query + self.evidence_q_proj[scale_idx](evidence)
+        query = query.view(batch_size, self.num_heads, self.head_dim, num_tokens).transpose(2, 3)
+
+        keys = [self.k_proj[scale_idx](feature) for feature in stream_features]
+        values = [self.v_proj[scale_idx](feature) for feature in stream_features]
+        key = torch.stack(keys, dim=1).view(batch_size, stream_count, self.num_heads, self.head_dim, num_tokens)
+        value = torch.stack(values, dim=1).view(batch_size, stream_count, self.num_heads, self.head_dim, num_tokens)
+        key = key.permute(0, 2, 1, 4, 3).reshape(batch_size, self.num_heads, stream_count * num_tokens, self.head_dim)
+        value = value.permute(0, 2, 1, 4, 3).reshape(batch_size, self.num_heads, stream_count * num_tokens, self.head_dim)
+
+        attention = torch.matmul(query, key.transpose(-1, -2)) * self.scale
+        attention = F.softmax(attention, dim=-1)
+        context = torch.matmul(attention, value)
+        return context.transpose(2, 3).reshape(batch_size, self.num_heads * self.head_dim, height, width)
+
+    def _window_partition_q(self, x: torch.Tensor, window_size: int) -> torch.Tensor:
+        batch_size, heads, head_dim, height, width = x.shape
+        return (
+            x.permute(0, 3, 4, 1, 2)
+            .reshape(batch_size, height // window_size, window_size, width // window_size, window_size, heads, head_dim)
+            .permute(0, 1, 3, 5, 2, 4, 6)
+            .reshape(batch_size * (height // window_size) * (width // window_size), heads, window_size * window_size, head_dim)
+        )
+
+    def _window_partition_kv(self, x: torch.Tensor, window_size: int) -> torch.Tensor:
+        batch_size, streams, heads, head_dim, height, width = x.shape
+        return (
+            x.permute(0, 4, 5, 1, 2, 3)
+            .reshape(
+                batch_size,
+                height // window_size,
+                window_size,
+                width // window_size,
+                window_size,
+                streams,
+                heads,
+                head_dim,
+            )
+            .permute(0, 1, 3, 6, 2, 4, 5, 7)
+            .reshape(
+                batch_size * (height // window_size) * (width // window_size),
+                heads,
+                window_size * window_size * streams,
+                head_dim,
+            )
+        )
+
+    def _window_unpartition(
+        self,
+        x: torch.Tensor,
+        batch_size: int,
+        height: int,
+        width: int,
+        window_size: int,
+    ) -> torch.Tensor:
+        heads = x.shape[1]
+        head_dim = x.shape[-1]
+        windows_h = height // window_size
+        windows_w = width // window_size
+        return (
+            x.reshape(batch_size, windows_h, windows_w, heads, window_size, window_size, head_dim)
+            .permute(0, 3, 6, 1, 4, 2, 5)
+            .reshape(batch_size, heads, head_dim, height, width)
+        )
+
+    def _windowed_spatial_context(
+        self,
+        stream_features: list[torch.Tensor],
+        unadapted: torch.Tensor,
+        scale_idx: int,
+        evidence: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        batch_size, _, height, width = unadapted.shape
+        stream_count = len(stream_features)
+        window_size = self.window_size
+        pad_h = (window_size - height % window_size) % window_size
+        pad_w = (window_size - width % window_size) % window_size
+        padded_h = height + pad_h
+        padded_w = width + pad_w
+
+        query = self.q_proj[scale_idx](unadapted)
+        if evidence is not None:
+            if evidence.shape[-2:] != query.shape[-2:]:
+                evidence = F.interpolate(evidence, size=query.shape[-2:], mode="bilinear", align_corners=False)
+            query = query + self.evidence_q_proj[scale_idx](evidence)
+        query = query.view(batch_size, self.num_heads, self.head_dim, height, width)
+        keys = [
+            self.k_proj[scale_idx](feature).view(batch_size, self.num_heads, self.head_dim, height, width)
+            for feature in stream_features
+        ]
+        values = [
+            self.v_proj[scale_idx](feature).view(batch_size, self.num_heads, self.head_dim, height, width)
+            for feature in stream_features
+        ]
+        if pad_h or pad_w:
+            query = F.pad(query, (0, pad_w, 0, pad_h))
+            keys = [F.pad(key, (0, pad_w, 0, pad_h)) for key in keys]
+            values = [F.pad(value, (0, pad_w, 0, pad_h)) for value in values]
+
+        key = torch.stack(keys, dim=1)
+        value = torch.stack(values, dim=1)
+        query_windows = self._window_partition_q(query, window_size)
+        key_windows = self._window_partition_kv(key, window_size)
+        value_windows = self._window_partition_kv(value, window_size)
+
+        attention = torch.matmul(query_windows, key_windows.transpose(-1, -2)) * self.scale
+        attention = F.softmax(attention, dim=-1)
+        context_windows = torch.matmul(attention, value_windows)
+        context = self._window_unpartition(context_windows, batch_size, padded_h, padded_w, window_size)
+        return context[:, :, :, :height, :width].reshape(batch_size, self.num_heads * self.head_dim, height, width)
+
+    def forward(
+        self,
+        stream_features: list[torch.Tensor],
+        unadapted: torch.Tensor,
+        scale_idx: int,
+        return_delta: bool = False,
+        evidence: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if not stream_features:
+            delta = torch.zeros_like(unadapted)
+            adapted = unadapted
+            if return_delta:
+                return adapted, delta
+            return adapted
+        if scale_idx in self.global_spatial_scales:
+            context = self._global_spatial_context(stream_features, unadapted, scale_idx, evidence)
+            return self._apply_residual(context, unadapted, scale_idx, return_delta, evidence)
+        if scale_idx in self.windowed_spatial_scales:
+            context = self._windowed_spatial_context(stream_features, unadapted, scale_idx, evidence)
+            return self._apply_residual(context, unadapted, scale_idx, return_delta, evidence)
+        return super().forward(stream_features, unadapted, scale_idx, return_delta, evidence)
+
+
+
+
+class Transformer2DBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_heads: int = 4,
+        dropout_rate: float = 0.1,
+        window_size: Optional[int] = None,
+        mlp_ratio: int = 2,
+    ) -> None:
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.window_size = window_size
+        self.norm1 = nn.LayerNorm(hidden_dim)
+        self.attn = nn.MultiheadAttention(
+            hidden_dim,
+            num_heads=num_heads,
+            dropout=dropout_rate,
+            batch_first=True,
+        )
+        self.norm2 = nn.GroupNorm(_group_count(hidden_dim), hidden_dim)
+        expanded = hidden_dim * mlp_ratio
+        self.mlp = nn.Sequential(
+            nn.Conv2d(hidden_dim, expanded, kernel_size=1),
+            nn.GELU(),
+            nn.Conv2d(expanded, expanded, kernel_size=3, padding=1, groups=expanded),
+            nn.GELU(),
+            nn.Dropout2d(p=dropout_rate),
+            nn.Conv2d(expanded, hidden_dim, kernel_size=1),
+            nn.Dropout2d(p=dropout_rate),
+        )
+
+    def _global_attention(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, channels, height, width = x.shape
+        sequence = x.flatten(2).transpose(1, 2)
+        sequence_norm = self.norm1(sequence)
+        attention_out, _ = self.attn(sequence_norm, sequence_norm, sequence_norm, need_weights=False)
+        return attention_out.transpose(1, 2).reshape(batch_size, channels, height, width)
+
+    def _window_partition(self, x: torch.Tensor, window_size: int) -> torch.Tensor:
+        batch_size, channels, height, width = x.shape
+        return (
+            x.reshape(batch_size, channels, height // window_size, window_size, width // window_size, window_size)
+            .permute(0, 2, 4, 3, 5, 1)
+            .reshape(batch_size * (height // window_size) * (width // window_size), window_size * window_size, channels)
+        )
+
+    def _window_unpartition(
+        self,
+        windows: torch.Tensor,
+        batch_size: int,
+        height: int,
+        width: int,
+        window_size: int,
+    ) -> torch.Tensor:
+        channels = windows.shape[-1]
+        return (
+            windows.reshape(
+                batch_size,
+                height // window_size,
+                width // window_size,
+                window_size,
+                window_size,
+                channels,
+            )
+            .permute(0, 5, 1, 3, 2, 4)
+            .reshape(batch_size, channels, height, width)
+        )
+
+    def _window_attention(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, _, height, width = x.shape
+        window_size = self.window_size
+        pad_h = (window_size - height % window_size) % window_size
+        pad_w = (window_size - width % window_size) % window_size
+        x_padded = F.pad(x, (0, pad_w, 0, pad_h)) if pad_h or pad_w else x
+        _, _, padded_h, padded_w = x_padded.shape
+        windows = self._window_partition(x_padded, window_size)
+        windows_norm = self.norm1(windows)
+        attention_windows, _ = self.attn(windows_norm, windows_norm, windows_norm, need_weights=False)
+        out = self._window_unpartition(attention_windows, batch_size, padded_h, padded_w, window_size)
+        return out[:, :, :height, :width]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.window_size is None:
+            x = x + self._global_attention(x)
+        else:
+            x = x + self._window_attention(x)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class TransformerEvidenceMaskAdapter(nn.Module):
+    """V2 mask adapter with evidence projection and efficient 2D transformer sharing."""
+
+    def __init__(
+        self,
+        hidden_dim: int = 128,
+        out_channels: int = 1,
+        output_resolution: tuple[int, int] = (128, 128),
+        in_channels_list: list[int] | None = None,
+        dropout_rate: float = 0.1,
+        num_heads: int = 4,
+        window_size: int = 8,
+    ) -> None:
+        super().__init__()
+        channels = in_channels_list or [256, 32, 64]
+        self.output_resolution = output_resolution
+        self.hidden_dim = hidden_dim
+        self.mask_evidence_proj = nn.ModuleList(
+            [
+                FlexibleEvidenceProjector(
+                    in_channels=hidden_dim + (2 * in_channels),
+                    hidden_dim=hidden_dim,
+                    dropout_rate=dropout_rate,
+                    depth=2,
+                )
+                for in_channels in channels
+            ]
+        )
+        self.low_blocks = nn.Sequential(
+            *[
+                Transformer2DBlock(hidden_dim, num_heads=num_heads, dropout_rate=dropout_rate, window_size=None)
+                for _ in range(2)
+            ]
+        )
+        self.mid_fuse = FlexibleEvidenceProjector(
+            in_channels=hidden_dim * 2,
+            hidden_dim=hidden_dim,
+            dropout_rate=dropout_rate,
+            depth=1,
+        )
+        self.mid_blocks = nn.Sequential(
+            *[
+                Transformer2DBlock(hidden_dim, num_heads=num_heads, dropout_rate=dropout_rate, window_size=window_size)
+                for _ in range(2)
+            ]
+        )
+        self.high_fuse = FlexibleEvidenceProjector(
+            in_channels=hidden_dim * 2,
+            hidden_dim=hidden_dim,
+            dropout_rate=dropout_rate,
+            depth=1,
+        )
+        self.high_blocks = nn.Sequential(
+            *[
+                Transformer2DBlock(hidden_dim, num_heads=num_heads, dropout_rate=dropout_rate, window_size=window_size)
+                for _ in range(2)
+            ]
+        )
+        self.mask_head = nn.Sequential(
+            nn.GroupNorm(_group_count(hidden_dim), hidden_dim),
+            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1),
+            nn.GELU(),
+            nn.Dropout2d(p=dropout_rate),
+            nn.Conv2d(hidden_dim, out_channels, kernel_size=1),
+        )
+        self.aux_head = nn.Sequential(
+            nn.GroupNorm(_group_count(hidden_dim), hidden_dim),
+            nn.Conv2d(hidden_dim, out_channels, kernel_size=1),
+        )
+
+    def _scale_evidence(
+        self,
+        adapted: FeaturePyramid,
+        unadapted: FeaturePyramid,
+        evidence: FeaturePyramid | None,
+    ) -> FeaturePyramid:
+        projected = []
+        for scale_idx, projection in enumerate(self.mask_evidence_proj):
+            if evidence is None:
+                evidence_i = torch.zeros(
+                    adapted[scale_idx].shape[0],
+                    self.hidden_dim,
+                    adapted[scale_idx].shape[-2],
+                    adapted[scale_idx].shape[-1],
+                    dtype=adapted[scale_idx].dtype,
+                    device=adapted[scale_idx].device,
+                )
+            else:
+                evidence_i = evidence[scale_idx]
+                if evidence_i.shape[-2:] != adapted[scale_idx].shape[-2:]:
+                    evidence_i = F.interpolate(
+                        evidence_i,
+                        size=adapted[scale_idx].shape[-2:],
+                        mode="bilinear",
+                        align_corners=False,
+                    )
+            adapter_delta = adapted[scale_idx] - unadapted[scale_idx]
+            projected.append(
+                projection(torch.cat([evidence_i, adapted[scale_idx], adapter_delta.abs()], dim=1))
+            )
+        return projected
+
+    def forward(
+        self,
+        adapted: FeaturePyramid,
+        streams_unadapted: StreamPyramid,
+        unadapted: FeaturePyramid,
+        evidence: FeaturePyramid | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+        del streams_unadapted
+        high_target = self.output_resolution
+        scale_features = self._scale_evidence(adapted, unadapted, evidence)
+
+        low = self.low_blocks(scale_features[0])
+        mid = scale_features[2] + F.interpolate(
+            low,
+            size=scale_features[2].shape[-2:],
+            mode="bilinear",
+            align_corners=False,
+        )
+        mid = self.mid_fuse(torch.cat([mid, scale_features[2]], dim=1))
+        mid = self.mid_blocks(mid)
+
+        high = scale_features[1] + F.interpolate(
+            mid,
+            size=scale_features[1].shape[-2:],
+            mode="bilinear",
+            align_corners=False,
+        )
+        high = self.high_fuse(torch.cat([high, scale_features[1]], dim=1))
+        high = self.high_blocks(high)
+
+        if high.shape[-2:] != high_target:
+            high = F.interpolate(high, size=high_target, mode="bilinear", align_corners=False)
+        mask_logits = self.mask_head(high)
+        aux_logits = self.aux_head(mid)
+        return mask_logits, aux_logits, None
+
+

detectivesam_inference/models/forgerylocalizer.py

@@ -0,0 +1,305 @@
+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_v1 import MaskAdapter, SharedAdapter
+from detectivesam_inference.models.adapters_v2 import (
+    CrossAttentionSharedAdapter,
+    SpatialCrossAttentionSharedAdapter,
+    StreamEvidenceBuilder,
+    TransformerEvidenceMaskAdapter,
+)
+
+
+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,
+        adapter_type: str = "conv",
+        mask_adapter_type: str = "coarse",
+        device: str = "cpu",
+    ) -> None:
+        super().__init__()
+        self.output_resolution = output_resolution
+        self.mask_adapter_type = mask_adapter_type
+
+        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 = self._build_feature_adapter(
+            adapter_type=adapter_type,
+            prompt_dim=prompt_dim,
+            dropout_rate=dropout_rate,
+            max_streams=max_streams,
+        )
+        self.evidence_builder = (
+            StreamEvidenceBuilder(
+                in_channels_list=MODEL_CHANNELS,
+                hidden_dim=prompt_dim,
+                dropout_rate=dropout_rate,
+                max_streams=max_streams,
+            )
+            if self._uses_evidence(adapter_type=adapter_type, mask_adapter_type=mask_adapter_type)
+            else None
+        )
+        self.mask_adapter = self._build_mask_adapter(
+            mask_adapter_type=mask_adapter_type,
+            prompt_dim=prompt_dim,
+            downscale=downscale,
+            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
+
+    @staticmethod
+    def _uses_evidence(*, adapter_type: str, mask_adapter_type: str) -> bool:
+        return adapter_type != "conv" or mask_adapter_type == "transformer"
+
+    @staticmethod
+    def _build_feature_adapter(
+        *,
+        adapter_type: str,
+        prompt_dim: int,
+        dropout_rate: float,
+        max_streams: int,
+    ) -> nn.Module:
+        adapter_classes = {
+            "conv": SharedAdapter,
+            "cross_attention": CrossAttentionSharedAdapter,
+            "spatial_cross_attention": SpatialCrossAttentionSharedAdapter,
+        }
+        if adapter_type not in adapter_classes:
+            raise ValueError(f"Unknown adapter_type: {adapter_type}")
+        return adapter_classes[adapter_type](
+            hidden_dim=prompt_dim,
+            in_channels_list=MODEL_CHANNELS,
+            dropout_rate=dropout_rate,
+            max_streams=max_streams,
+        )
+
+    @staticmethod
+    def _build_mask_adapter(
+        *,
+        mask_adapter_type: str,
+        prompt_dim: int,
+        downscale: int,
+        dropout_rate: float,
+        max_streams: int,
+    ) -> nn.Module:
+        if mask_adapter_type == "coarse":
+            return 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,
+            )
+        if mask_adapter_type == "transformer":
+            del downscale, max_streams
+            return TransformerEvidenceMaskAdapter(
+                hidden_dim=prompt_dim,
+                output_resolution=MASK_ADAPTER_RESOLUTION,
+                in_channels_list=MODEL_CHANNELS,
+                dropout_rate=dropout_rate,
+            )
+        raise ValueError(f"Unknown mask_adapter_type: {mask_adapter_type}")
+
+    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,
+        evidence: FeaturePyramid | None,
+    ) -> FeaturePyramid:
+        if evidence is None:
+            return [
+                self.adapters(streams_unadapted[scale_idx], unadapted[scale_idx], scale_idx)
+                for scale_idx in range(len(unadapted))
+            ]
+        return [
+            self.adapters(
+                streams_unadapted[scale_idx],
+                unadapted[scale_idx],
+                scale_idx,
+                evidence=evidence[scale_idx],
+            )
+            for scale_idx in range(len(unadapted))
+        ]
+
+    def _build_mask_prompt(
+        self,
+        adapted: FeaturePyramid,
+        streams_unadapted: StreamPyramid,
+        unadapted: FeaturePyramid,
+        evidence: FeaturePyramid | None,
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor | None]]:
+        if self.mask_adapter_type == "transformer":
+            mask_prompt, coarse_mask, detection_logit = self.mask_adapter(
+                adapted,
+                streams_unadapted,
+                unadapted,
+                evidence=evidence,
+            )
+            return mask_prompt, {
+                "coarse_mask": coarse_mask,
+                "detection_logit": detection_logit,
+            }
+
+        mask_prompt = self.mask_adapter(adapted, streams_unadapted, unadapted)
+        return mask_prompt, {"coarse_mask": None, "detection_logit": None}
+
+    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],
+        output_extras: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor | None]]:
+        unadapted, streams_unadapted = self._encode_original_and_streams(orig, streams)
+        evidence = (
+            self.evidence_builder(unadapted, streams_unadapted)
+            if self.evidence_builder is not None
+            else None
+        )
+        adapted = self._apply_adapters(unadapted, streams_unadapted, evidence)
+
+        mask_prompt, extras = self._build_mask_prompt(adapted, streams_unadapted, unadapted, evidence)
+        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, iou_pred, sam_tokens_out, object_score_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,
+        )
+        processed_mask_logits = self.transforms.postprocess_masks(mask_logits, torch.Size(self.output_resolution))
+        if output_extras:
+            extras.update(
+                {
+                    "iou_pred": iou_pred,
+                    "sam_tokens_out": sam_tokens_out,
+                    "object_score_logits": object_score_logits,
+                    "mask_prompt": mask_prompt,
+                }
+            )
+            return processed_mask_logits, extras
+        return processed_mask_logits

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_v2",
+        help="Checkpoint path or alias. Built-in aliases: detective_sam, detective_sam_sota, detective_sam_v2.",
+    )
+    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,138 @@
+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,
+            adapter_type=self.config.adapter_type,
+            mask_adapter_type=self.config.mask_adapter_type,
+            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)

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,8 @@
+torch>=2.5.1,<2.8
+torchvision>=0.20.1,<0.23
+numpy>=1.24
+Pillow>=10.0
+opencv-python-headless>=4.10
+hydra-core>=1.3
+sam2
+tqdm>=4.66

sam2configs/README.md

@@ -0,0 +1,13 @@
+# SAM2 Weights
+
+The Hugging Face bundle ships the SAM2 Hiera Base Plus checkpoint.
+
+The full runnable bundle is hosted on Hugging Face:
+
+- https://huggingface.co/Gertlek/DetectiveSAMv2
+
+To run inference from a lightweight checkout, download the checkpoint file from Hugging Face and place it in this directory with this exact name:
+
+- `sam2.1_hiera_base_plus.pt`
+
+The bundled inference configs already point to that path.

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

tests/test_regression.py

@@ -0,0 +1,189 @@
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import pytest
+
+from detectivesam_inference.checkpoint import load_inference_config, 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.models.adapters_v1 import MaskAdapter
+from detectivesam_inference.models.adapters_v2 import TransformerEvidenceMaskAdapter
+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:
+    checkpoint_path = resolve_checkpoint_path("detective_sam", get_repo_root())
+    if not checkpoint_path.exists():
+        pytest.skip(f"Missing optional checkpoint: {checkpoint_path}")
+    return DetectiveSAMRunner(checkpoint_path="detective_sam", device="cpu")
+
+
+@pytest.fixture(scope="module")
+def sota_runner() -> DetectiveSAMRunner:
+    checkpoint_path = resolve_checkpoint_path("detective_sam_sota", get_repo_root())
+    if not checkpoint_path.exists():
+        pytest.skip(f"Missing optional checkpoint: {checkpoint_path}")
+    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"
+    assert resolve_checkpoint_path("detective_sam_v2", repo_root) == repo_root / "checkpoints" / "detective_sam_v2.pth"
+
+
+def test_v2_checkpoint_sidecar(repo_root: Path) -> None:
+    config = load_inference_config(repo_root / "checkpoints" / "detective_sam_v2.pth")
+    assert config.prompt_dim == 96
+    assert config.downscale == 8
+    assert config.max_streams == 3
+    assert config.perturbation_type == "gaussian_blur+jpeg_compression+gaussian_noise"
+    assert config.adapter_type == "spatial_cross_attention"
+    assert config.mask_adapter_type == "transformer"
+
+
+def test_json_checkpoint_sidecar(tmp_path: Path) -> None:
+    checkpoint_path = tmp_path / "best_model.pth"
+    checkpoint_path.touch()
+    (tmp_path / "model_params.json").write_text(
+        """
+{
+    "model_config": {
+    "prompt_dim": 96,
+    "downscale": 8,
+    "dropout_rate": 0.1,
+    "adapter_type": "spatial_cross_attention",
+    "mask_adapter_type": "transformer"
+  },
+  "training_config": {"img_size": 512},
+  "data_config": {
+    "perturbation_type": "gaussian_blur+jpeg_compression+gaussian_noise",
+    "perturbation_intensity": 0.5
+  },
+  "sam_config": {
+    "sam_config_file": "sam2.1_hiera_b+.yaml",
+    "sam_checkpoint": "sam2configs/sam2.1_hiera_base_plus.pt"
+  }
+}
+""",
+        encoding="utf-8",
+    )
+    config = load_inference_config(checkpoint_path)
+    assert config.prompt_dim == 96
+    assert config.max_streams == 3
+    assert config.adapter_type == "spatial_cross_attention"
+    assert config.mask_adapter_type == "transformer"
+
+
+def test_adapter_modules_are_split() -> None:
+    assert MaskAdapter.__module__ == "detectivesam_inference.models.adapters_v1"
+    assert TransformerEvidenceMaskAdapter.__module__ == "detectivesam_inference.models.adapters_v2"
+
+
+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)