| """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 |
|
|
| |
| 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: |
| 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) |
| 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) |
| 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] |
| grid = (2 * src / (size - 1) - 1).flip(-1) |
| 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: |
| 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") |
|
|
