inference_utils.py

"""Shared plumbing for the VISTA-PATH inference entrypoints.

Both ``inference.py`` (tissue-prompted / prompt-free whole-slide inference) and
``inference_bbx.py`` (box-prompted ROI inference) build the same model and use
the same sliding-window machinery; only the prompt source and the output
post-processing differ.
"""

import random

import numpy as np
import torch

from models.backbones import CustomSegmentationModel

# Text encoder (PLIP = CLIP fine-tuned on pathology) and the Mask2Former
# checkpoint the segmentation trunk was initialized from. Both must match the
# values used at training time or the released weights will not load.
BASE_MODEL_NAME = "vinid/plip"
MASK2FORMER_NAME = "facebook/mask2former-swin-small-ade-semantic"
SAM_PRETRAINED = "facebook/sam-vit-base"

# Prompt template used for every class name at training time.
TEXT_TEMPLATE = "an image of {}"


def set_seed(seed=42):
    random.seed(seed)                      # Python random module
    np.random.seed(seed)                   # NumPy
    torch.manual_seed(seed)                # PyTorch CPU
    torch.cuda.manual_seed(seed)           # PyTorch GPU
    torch.cuda.manual_seed_all(seed)       # All GPUs (if using multi-GPU)

    torch.backends.cudnn.deterministic = True   # Ensure deterministic behavior
    torch.backends.cudnn.benchmark = False


class SegWrapper(torch.nn.Module):
    """Thin wrapper kept for checkpoint compatibility.

    The released weights were saved from the HuggingFace Trainer, so every key
    is prefixed with ``model.``. Keeping the wrapper lets the checkpoint load
    strictly, without key surgery.
    """

    def __init__(self, model):
        super().__init__()
        self.model = model
        self._keys_to_ignore_on_save = None

    def forward(self, pixel_values_m2f, input_ids, attention_mask, labels=None, box=None):
        logits, _ = self.model(pixel_values_m2f, input_ids, attention_mask, box)
        loss = None
        if labels is not None:
            loss = torch.nn.functional.cross_entropy(logits, labels)
        return {"logits": logits.detach(), "loss": loss}


def add_model_args(parser):
    """Architecture / checkpoint arguments shared by both entrypoints."""
    parser.add_argument("--checkpoint_file", type=str, required=True,
                        help="Path to the released VISTA-PATH weights "
                             "(./checkpoints/pytorch_model.bin).")
    parser.add_argument("--num_queries", type=int, default=20,
                        help="Mask2Former queries. Must match training (20 for the "
                             "released checkpoint).")
    parser.add_argument("--m2f_image_size", type=int, default=512,
                        help="Resolution fed to the model. Must match the value used "
                             "at training time (512 for the released checkpoint).")
    parser.add_argument("--tune_mode", type=str, default="freeze")
    # Kept for backwards compatibility with the published command lines; the
    # current architecture derives its widths from the pretrained trunks.
    parser.add_argument("--d_model", type=int, default=512, help="Unused, kept for CLI compatibility.")
    parser.add_argument("--nhead", type=int, default=8)
    parser.add_argument("--num_layers", type=int, default=4, help="Unused, kept for CLI compatibility.")
    return parser


def load_model(args, device, bbx_random=None):
    """Instantiate VISTA-PATH and load the released checkpoint.

    ``bbx_random`` overrides ``args.bbx_random``; it is the probability of
    *dropping* the box prompt. Use 0 to always honour the box and 1 to run
    fully prompt-free.
    """
    core_model = CustomSegmentationModel(
        BASE_MODEL_NAME,
        args.d_model,
        args.nhead,
        args.num_layers,
        args.bbx_random if bbx_random is None else bbx_random,
        args.tune_mode,
        mask2former_name=MASK2FORMER_NAME,
        num_queries=args.num_queries,
        image_size=args.m2f_image_size,
        sam_pretrained=SAM_PRETRAINED,
    )
    model = SegWrapper(core_model)

    checkpoint_path = args.checkpoint_file
    if checkpoint_path.endswith(".safetensors"):
        from safetensors.torch import load_file
        checkpoint = load_file(checkpoint_path, device="cpu")
    else:
        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    # Accept both a bare state_dict and a training checkpoint dict.
    if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
        checkpoint = checkpoint["model_state_dict"]
    # Weights saved from a plain CustomSegmentationModel (no Trainer wrapper)
    # have no "model." prefix; add it so both layouts load strictly.
    if not any(k.startswith("model.") for k in checkpoint):
        checkpoint = {f"model.{k}": v for k, v in checkpoint.items()}

    model.load_state_dict(checkpoint)
    model = model.to(device)
    model.eval()
    return model


def encode_class_texts(processor, idx_to_class, class_indices):
    """Tokenize the prompt for every class once; workers reuse the cache."""
    text_cache = {}
    for c in class_indices:
        text = TEXT_TEMPLATE.format(idx_to_class[c])
        enc = processor(text=text, return_tensors="pt",
                        padding="max_length", truncation=True, max_length=77)
        text_cache[c] = (enc["input_ids"].squeeze(0), enc["attention_mask"].squeeze(0))
    return text_cache


def get_patch_positions(H, W, patch_size=1024, overlap=128):
    """Compute sliding-window patch top-left positions without loading image data."""
    overlap = max(0, min(overlap, patch_size - 1))
    stride_y = patch_size - overlap if H >= patch_size else H
    stride_x = patch_size - overlap if W >= patch_size else W

    def build_coords(limit, psize, stride):
        if limit <= psize:
            return [0]
        coords = list(range(0, limit - psize + 1, stride))
        if coords[-1] + psize < limit:
            coords.append(limit - psize)
        return coords

    y_coords = build_coords(H, min(H, patch_size), stride_y)
    x_coords = build_coords(W, min(W, patch_size), stride_x)
    return [(x, y) for y in y_coords for x in x_coords]


def generate_gaussian_weight_mask(height, width, sigma_scale=0.125):
    """Gaussian falloff toward patch edges, used to blend overlapping windows."""
    y = np.linspace(-1, 1, height)
    x = np.linspace(-1, 1, width)
    xv, yv = np.meshgrid(x, y)
    sigma_x = sigma_scale * 2
    sigma_y = sigma_scale * 2
    gauss = np.exp(-((xv**2) / (2 * sigma_x**2) + (yv**2) / (2 * sigma_y**2)))
    gauss /= gauss.max()
    return gauss.astype(np.float32)