File size: 6,653 Bytes
e1b0735 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | """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)
|