Fix SAM v1 off-by-one in hidden_states indexing

[LevanBokeria]

Fixes #66.

The bug

extract_sam_2d_features / extract_sam_volume_features read hidden_states[idx] with the
raw SAM_CONFIGS layer_indices, but hidden_states[0] is the patch embedding. Every SAM v1
layer was therefore shifted one block earlier and the final block was never read:
(2, 5, 8, 11) read blocks 1, 4, 7, 10.

Two things made this worth fixing beyond tidiness:

• DINOv3 was unaffected, so the head-to-head was skewed. get_intermediate_layers(n=list(...))
  is genuinely block-indexed (for i, blk in enumerate(self.blocks): ... if i in blocks_to_take).
• SAM v1 classification already used the true final block (cache_sam_cls_features loops all
  12 blocks manually), so segmentation and classification meant different things by "SAM's last
  layer".

Changes

• SAM_CONFIGS layer_indices are now block indices, with the +1 applied at read time.
  sam-vit-base moves (2, 5, 8, 11) → (3, 6, 9, 11), matching DINOV3_CONFIGS at equal
  depth; large/huge values are unchanged but now resolve one block later, reaching blocks 23
  and 31. This also changes the layers_ cache tag, so a stale cache cannot be silently reused.
• New features.py::_select_hidden_state — the single place the convention lives. Raises a
  named IndexError (config index and resolved position) instead of silently reading the wrong
  entry or a bare IndexError.
• extract_sam_volume_features now takes a required keyword-only hidden_state_offset. It
  serves both families (run_segmentation.py routes --backbone sam and --backbone sam2
  through it), so a bare idx + 1 there would have broken SAM2. The offset is carried in
  load_backbone metadata (1 for SAM v1, 0 for SAM2) rather than guessed at the call site.
• hidden_state_offset recorded in stdout and in the committed results JSON, so old and new
  runs are distinguishable after the fact.
• tests/test_feature_layer_indexing.py — 7 synthetic tests (no weights or data). The key
  one asserts SAM v1 reads [hs[4], hs[7], hs[10], hs[12]] for (3, 6, 9, 11); it fails on the
  old code. Two more pin SAM2 to offset 0 through both extractors.
• Docs: an AGENTS.md gotcha on the two conventions, and a resolution note on N2 in the
  2026-09-01 refresher so the "decide this before extracting" instruction isn't re-litigated.

SAM2 is unchanged

Verified on sam2.1-hiera-base-plus: 25 hidden states / 24 blocks, Stage 3 (448ch, 64×64) spans
hs[6..21], and the configured (6, 11, 16, 21) land inside Stage 3 with hs[21] = its last
Stage-3 block. Applying +1 would read hs[22] (896ch, 32×32) and break the decoder.
extract_sam2_2d_features now routes through the helper with offset=0 — explicit, and
behaviourally identical.

Verification

• End-to-end on real facebook/sam-vit-base weights: the last extracted layer now equals
  hidden_states[12] (block 11) and no longer equals hidden_states[11].
• pytest tests/test_feature_layer_indexing.py tests/test_segmentation_metrics.py → 14 passed.
• pre-commit run --files <changed> → ruff, ruff-format and mypy pass on every changed file.

Notes for the reviewer

• Re-caching is required for SAM v1. The ACDC SAM v1 cache was already deleted and no other
  SAM cache exists, so nothing stale is left behind. Expect Dice to move slightly: on
  sam-vit-base the previously-skipped block sits at cosine 0.983 / ‖Δ‖/‖h‖ ≈ 0.24 from the one
  that was read, so this is a correctness fix, not an expected win for SAM.
• Two pre-existing problems were left alone, both reproducible on re-cache-all-features:
  tests/test_classification_multi_dataset.py fails at collection (MONAI's torch.jit.interface
DeprecationWarning vs filterwarnings = ["error"]), and the pinned ruff-format (0.5.5)
  reformats files the repo has committed under a newer ruff — pre-commit run -a rewrites ~130
  results JSONs and several scripts untouched by this PR.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB