File size: 10,269 Bytes
c5a246d db46ad7 c5a246d db46ad7 c5a246d db46ad7 c5a246d db46ad7 c5a246d db46ad7 c5a246d db46ad7 c5a246d db46ad7 c5a246d db46ad7 c5a246d db46ad7 c5a246d db46ad7 c5a246d | 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 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 | """Inference helpers for the GLINT-CT vision-language model.
CT volumes are preprocessed with the same MONAI pipeline used during GLINT
training/zero-shot evaluation (orientation IPL, resample to 1.5/0.75/0.75 mm,
HU window [-1000, 1000] -> [-1, 1], 240x384x384), then regrouped into
``[num_frames=80, 3, 384, 384]`` and ImageNet-normalised.
"""
import numpy as np
import torch
# GLINT-CT pretraining preprocessing (glint/configs/chest_ct_monai.yaml)
CT_MONAI_CONFIG = {
"orientation": "IPL",
"target_spacing": (1.5, 0.75, 0.75),
"roi_size": (240, 480, 480),
"final_size": (240, 384, 384),
"hu_range": (-1000.0, 1000.0),
"normalize_range": (-1.0, 1.0),
}
def build_ct_transform(orientd=True, raw_spacing=None):
from monai.transforms import (
CenterSpatialCropd,
Compose,
EnsureChannelFirstd,
EnsureTyped,
Lambdad,
LoadImaged,
Orientationd,
Resized,
ScaleIntensityRanged,
Spacingd,
SpatialPadd,
)
cfg = CT_MONAI_CONFIG
b_min, b_max = cfg["normalize_range"]
def _patch_affine(img):
if raw_spacing is not None and hasattr(img, "affine"):
sx, sy, sz = (float(x) for x in raw_spacing)
a = torch.eye(4, dtype=img.affine.dtype, device=img.affine.device)
a[0, 0], a[1, 1], a[2, 2] = sx, sy, sz
img.affine = a
return img
steps = [
LoadImaged(keys=["image"], image_only=False),
EnsureChannelFirstd(keys=["image"]),
*([Orientationd(keys=["image"], axcodes=cfg["orientation"])] if orientd else []),
*([Lambdad(keys=["image"], func=_patch_affine)] if raw_spacing is not None else []),
ScaleIntensityRanged(
keys=["image"],
a_min=cfg["hu_range"][0],
a_max=cfg["hu_range"][1],
b_min=b_min,
b_max=b_max,
clip=True,
),
Spacingd(keys=["image"], pixdim=cfg["target_spacing"], mode="bilinear"),
CenterSpatialCropd(keys=["image"], roi_size=cfg["roi_size"]),
SpatialPadd(keys=["image"], spatial_size=cfg["roi_size"], mode="constant", constant_values=b_min),
Resized(keys=["image"], spatial_size=cfg["final_size"], mode="trilinear"),
EnsureTyped(keys=["image"], dtype=torch.float32),
]
return Compose(steps)
def normalize_pixel_values_videos(pixel_values_videos):
"""(D, H, W) in [-1, 1] -> (num_frames, 3, H, W) ImageNet-normalised."""
s = pixel_values_videos.shape
num_channels = 3
pixel_values_videos = pixel_values_videos.reshape(-1, num_channels, s[-2], s[-1])
pixel_values_videos = (pixel_values_videos + 1) / 2
image_mean = torch.tensor([0.485, 0.456, 0.406], device=pixel_values_videos.device).view(1, 3, 1, 1)
image_std = torch.tensor([0.229, 0.224, 0.225], device=pixel_values_videos.device).view(1, 3, 1, 1)
pixel_values_videos = (pixel_values_videos - image_mean) / image_std
return pixel_values_videos
def _input_geometry(image, orientd, raw_spacing):
"""Input voxel grid (shape, affine) and its 4x4 voxel mapping onto the preprocessed grid."""
import nibabel as nib
affine = np.asarray(image.meta["original_affine"], dtype=np.float64)
shape = tuple(int(s) for s in image.meta["spatial_shape"])
ref_affine = affine
if raw_spacing is not None: # the pipeline replaced the affine with diag(raw_spacing) after reorientation
to_oriented = np.eye(4)
if orientd:
ornt = nib.orientations.ornt_transform(
nib.orientations.io_orientation(affine),
nib.orientations.axcodes2ornt(CT_MONAI_CONFIG["orientation"]),
)
to_oriented = np.linalg.inv(nib.orientations.inv_ornt_aff(ornt, shape))
ref_affine = np.diag([*map(float, raw_spacing), 1.0]) @ to_oriented
to_model = np.linalg.inv(np.asarray(image.affine, dtype=np.float64)) @ ref_affine
return {"shape": shape, "affine": affine, "to_model": to_model}
def load_ct_volume(volume, orientd=True, raw_spacing=None, return_geometry=False):
"""Load and preprocess a CT volume.
volume : path to a NIfTI (.nii/.nii.gz) file or a DICOM series directory, or a (240, 384, 384)
tensor/ndarray in HU that is already on the preprocessed grid (IPL, 1.5/0.75/0.75 mm
resampled, cropped/padded, resized); only the HU window is applied to arrays. Returns
pixel_values of shape (1, num_frames, 3, H, W), plus the input grid if ``return_geometry``
(``shape``, ``affine`` (None for array input) and ``to_model``, its voxel mapping onto the
preprocessed grid), which ``similarity_scores_to_volume`` maps the scores back onto.
"""
if isinstance(volume, str):
transform = build_ct_transform(orientd=orientd, raw_spacing=raw_spacing)
out = transform({"image": volume})
image = out["image"].squeeze(0) # (1, D, H, W) -> (D, H, W)
geometry = _input_geometry(out["image"], orientd, raw_spacing)
else:
image = torch.as_tensor(np.asarray(volume), dtype=torch.float32)
if tuple(image.shape) != tuple(CT_MONAI_CONFIG["final_size"]):
raise ValueError(
f"Array input must be on the preprocessed grid {CT_MONAI_CONFIG['final_size']}, got {tuple(image.shape)}"
)
(a_min, a_max), (b_min, b_max) = CT_MONAI_CONFIG["hu_range"], CT_MONAI_CONFIG["normalize_range"]
image = (image.clamp(a_min, a_max) - a_min) / (a_max - a_min) * (b_max - b_min) + b_min
geometry = {"shape": tuple(image.shape), "affine": None, "to_model": np.eye(4)}
pixel_values_videos = normalize_pixel_values_videos(image).unsqueeze(0) # (1, num_frames, 3, H, W)
return (pixel_values_videos, geometry) if return_geometry else pixel_values_videos
def _resample_to_input(volume, to_model, shape, chunk=32):
"""Trilinearly sample a preprocessed-grid volume at every input voxel; 0 outside the preprocessing crop."""
if np.allclose(to_model, np.eye(4)) and tuple(shape) == tuple(volume.shape):
return volume
device = volume.device
M = torch.as_tensor(to_model, dtype=torch.float32, device=device)
size = torch.tensor(volume.shape, dtype=torch.float32, device=device)
ni, nj, nk = shape
jk = torch.stack(torch.meshgrid(torch.arange(nj), torch.arange(nk), indexing="ij"), -1).reshape(-1, 2).to(device)
out = torch.empty(tuple(shape), dtype=torch.float32, device=device)
for i0 in range(0, ni, chunk):
ii = torch.arange(i0, min(i0 + chunk, ni), device=device)
ijk = torch.cat([ii.repeat_interleave(len(jk))[:, None], jk.repeat(len(ii), 1)], 1).float()
src = ijk @ M[:3, :3].T + M[:3, 3] # input voxel -> preprocessed voxel
grid = (2 * src / (size - 1) - 1).flip(-1) # (d, h, w) -> grid_sample's (x, y, z)
out[i0 : i0 + len(ii)] = torch.nn.functional.grid_sample(
volume[None, None], grid.view(1, -1, 1, 1, 3), mode="bilinear", padding_mode="zeros", align_corners=True
).view(len(ii), nj, nk)
return out
def similarity_scores_to_volume(similarity_scores, model, geometry=None):
"""Per-token scores ``[1, 1, num_tokens]`` -> similarity map on the input CT grid.
The token grid is upsampled (trilinear) to the 240x384x384 preprocessed grid and, with ``geometry`` from
``load_ct_volume``, resampled onto the input voxels so the map has the input's shape and affine (0 outside
the 360 mm preprocessing crop). Without ``geometry`` the map stays on the preprocessed grid.
"""
vc = model.config.vision_config
patch_size = getattr(vc, "patch_size", None) or 16
frame_length = vc.num_channels * (getattr(vc, "tubelet_size", None) or 2)
grid = (vc.input_slices // frame_length, vc.img_size // patch_size, vc.img_size // patch_size)
similarity_map = similarity_scores.float().reshape(grid)
similarity_map = torch.nn.functional.interpolate(
similarity_map[None, None], size=CT_MONAI_CONFIG["final_size"], mode="trilinear", align_corners=False
)[0, 0]
if geometry is None:
return similarity_map
return _resample_to_input(similarity_map, geometry["to_model"], geometry["shape"])
def save_similarity_map(similarity_map, affine, save_path):
"""Save the similarity map as NIfTI; with ``geometry["affine"]`` it overlays the input CT voxel for voxel."""
import nibabel as nib
sim = similarity_map.detach().float().cpu().numpy() if torch.is_tensor(similarity_map) else similarity_map
nib.save(nib.Nifti1Image(sim.astype(np.float32), affine), save_path)
def _as_pixel_values(volume, raw_spacing):
if torch.is_tensor(volume) and volume.dim() == 5: # already preprocessed by load_ct_volume
return volume
return load_ct_volume(volume, raw_spacing=raw_spacing)
@torch.no_grad()
def model_inference(volume, text, model, tokenizer, raw_spacing=None):
"""Return (similarity_prob, similarity_scores) for one CT volume and one phrase.
similarity_prob : scalar sigmoid probability that the phrase applies.
similarity_scores : per-token similarity in [0, 1] over the V-JEPA video tokens
(shape [1, 1, num_tokens]); ``similarity_scores_to_volume``
maps it onto the input CT grid.
volume may also be the ``pixel_values`` returned by ``load_ct_volume``.
"""
pixel_values = _as_pixel_values(volume, raw_spacing).to(model.device, dtype=model.dtype)
tokenized_text = tokenizer(
text,
padding=True,
truncation=True,
return_tensors="pt",
).to(model.device)
output = model.compute_logits(pixel_values, [tokenized_text])
similarity_prob = output["logits"].sigmoid()
return similarity_prob, output.get("similarity_scores")
@torch.no_grad()
def model_inference_multiple_text(volume, text_list, model, tokenizer, raw_spacing=None):
pixel_values = _as_pixel_values(volume, raw_spacing).to(model.device, dtype=model.dtype)
enc = tokenizer(text_list, padding=True, truncation=True, return_tensors="pt").to(model.device)
output = model.compute_logits(pixel_values, [enc])
return output["logits"].sigmoid(), output.get("similarity_scores")
|