data/ct_bank.py

"""CT token-bank builder (Gate 0).

Builds the held-out CT patch-token bank Z used by Phase 1 subspace constructions.
Tokens come ONLY from the frozen MedDINOv3 backbone over held-out CT slices. No labels
are read here; the builder operates purely on pixels + the frozen backbone.

Gate 0 criterion: token bank size >= 2e6 tokens [FIXED]. With 196 patch tokens per
224x224 slice, that is ~10,205 slices.

If `image_root` is not provided (pixel data unavailable — see loaders.py), the builder
returns a result with `data_gap=True` and `n_tokens=0`, which the Gate 0 runner records
honestly rather than substituting non-comparable data (IMPLEMENTATION_SPEC §7).
"""
from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np
import torch
from PIL import Image

from backbone.meddino import CT_MEAN, CT_STD, MedDINOv3Backbone

from .loaders import (
    iter_manifest,
    iter_slices_from_tree,
    load_scan_splits,
    resolve_image,
)


@dataclass
class BankResult:
    n_tokens: int
    n_slices: int
    out_path: str | None
    data_gap: bool
    gap_reason: str | None = None
    dim: int = 0
    meta: dict = field(default_factory=dict)


def _load_slice(path: str, image_size: int) -> torch.Tensor:
    img = Image.open(path).convert("RGB").resize((image_size, image_size), Image.BILINEAR)
    arr = np.asarray(img, dtype=np.float32) / 255.0
    arr = (arr - np.asarray(CT_MEAN, np.float32)) / np.asarray(CT_STD, np.float32)
    return torch.from_numpy(arr).permute(2, 0, 1)  # (3,H,W)


def build_token_bank_from_tree(
    backbone: MedDINOv3Backbone,
    image_root: str,
    splits_json_path: str,
    out_path: str,
    target_tokens: int = 2_000_000,
    held_out_split: str = "train",
    image_size: int = 224,
    batch_size: int = 32,
) -> BankResult:
    """Build the held-out CT token bank from a local raw/lidc tree + scan-level splits.

    Tokens come ONLY from the frozen backbone over slices whose scan is in
    `held_out_split`, keeping the bank disjoint from eval scans (no labels are read).
    """
    scan_splits = load_scan_splits(splits_json_path)
    slices = list(iter_slices_from_tree(image_root, scan_splits, held_out_split))
    if not slices:
        return BankResult(
            n_tokens=0, n_slices=0, out_path=None, data_gap=True,
            gap_reason=(
                f"No '{held_out_split}' CT slices found under image_root={image_root!r} "
                f"(scans matching split in {splits_json_path})."
            ),
        )

    chunks: list[torch.Tensor] = []
    total = 0
    n_slices = 0
    scans_used: set[str] = set()
    batch: list[torch.Tensor] = []
    batch_scans: list[str] = []

    def flush():
        nonlocal total, n_slices
        if not batch:
            return
        imgs = torch.stack(batch, 0)
        Z = backbone.extract_patch_tokens(imgs)
        chunks.append(Z.reshape(-1, Z.shape[-1]))
        total += chunks[-1].shape[0]
        n_slices += imgs.shape[0]
        scans_used.update(batch_scans)
        batch.clear()
        batch_scans.clear()

    for scan_id, png in slices:
        try:
            batch.append(_load_slice(png, image_size))
            batch_scans.append(scan_id)
        except Exception:
            continue
        if len(batch) >= batch_size:
            flush()
        if total >= target_tokens:
            break
    flush()

    bank = torch.cat(chunks, 0) if chunks else torch.empty(0)
    torch.save({"tokens": bank, "n_slices": n_slices, "split": held_out_split}, out_path)
    return BankResult(
        n_tokens=int(bank.shape[0]),
        n_slices=n_slices,
        out_path=out_path,
        data_gap=False,
        dim=int(bank.shape[-1]) if bank.numel() else 0,
        meta={
            "available_slices": len(slices),
            "scans_used": len(scans_used),
            "held_out_split": held_out_split,
        },
    )


def build_token_bank(
    backbone: MedDINOv3Backbone,
    manifest_local_path: str,
    image_root: str | None,
    out_path: str,
    target_tokens: int = 2_000_000,
    held_out_split: str = "train",
    image_size: int = 224,
    batch_size: int = 16,
) -> BankResult:
    records = list(iter_manifest(manifest_local_path, split=held_out_split))
    resolved = [
        (r, resolve_image(image_root, r.image_path)) for r in records
    ]
    available = [(r, p) for r, p in resolved if p is not None]

    if not available:
        return BankResult(
            n_tokens=0,
            n_slices=0,
            out_path=None,
            data_gap=True,
            gap_reason=(
                f"No CT slice pixels accessible. Manifest lists {len(records)} "
                f"held-out '{held_out_split}' slices, but image_root="
                f"{image_root!r} resolved 0 of them. The interim PNG bucket "
                f"hf://buckets/Chucks90/eryon-datasets is not readable with the "
                f"provided token. Provide a local LIDC slice mirror to build the bank."
            ),
            meta={"manifest_slices": len(records)},
        )

    chunks: list[torch.Tensor] = []
    total = 0
    n_slices = 0
    batch: list[torch.Tensor] = []

    def flush():
        nonlocal total, n_slices
        if not batch:
            return
        imgs = torch.stack(batch, 0)
        Z = backbone.extract_patch_tokens(imgs)  # (B,n,d)
        chunks.append(Z.reshape(-1, Z.shape[-1]))
        total += chunks[-1].shape[0]
        n_slices += imgs.shape[0]
        batch.clear()

    for r, p in available:
        batch.append(_load_slice(p, image_size))
        if len(batch) >= batch_size:
            flush()
        if total >= target_tokens:
            break
    flush()

    bank = torch.cat(chunks, 0) if chunks else torch.empty(0)
    torch.save({"tokens": bank, "n_slices": n_slices}, out_path)
    return BankResult(
        n_tokens=int(bank.shape[0]),
        n_slices=n_slices,
        out_path=out_path,
        data_gap=False,
        dim=int(bank.shape[-1]) if bank.numel() else 0,
        meta={"manifest_slices": len(records), "resolved_slices": len(available)},
    )