File size: 5,025 Bytes
32e07d6 | 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 | """Inference helpers for the GLINT-CXR vision-language model.
These utilities mirror the zero-shot grounding / phrase-similarity pipeline used
to evaluate GLINT on chest X-rays. ``model.compute_logits`` already removes the
CLS token from the returned similarity map (the model was trained with
``use_cls_token=False``), so the helpers here only reshape and interpolate.
"""
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
@torch.no_grad()
def model_inference(image, text, model, image_processor, tokenizer):
"""Return (similarity_prob, similarity_map) for one image and one phrase.
similarity_prob : scalar sigmoid probability that the phrase applies to the image.
similarity_map : per-pixel similarity map in [0, 1] (H x W), upsampled to the
original image resolution.
"""
image = load_image(image)
(width, height) = image.size
image_size = (height, width)
image_processor_outputs = image_processor(image)
processed_image = torch.FloatTensor(np.array(image_processor_outputs["pixel_values"])).to(model.device)
tokenized_text = tokenizer(
text,
padding=True,
truncation=True,
return_tensors="pt",
).to(model.device)
output = model.compute_logits(processed_image, [tokenized_text])
logits = output["logits"]
similarity_prob = logits.sigmoid()
similarity_scores = output["similarity_scores"]
similarity_scores = similarity_scores.view(-1)
similarity_scores = interpolate_similarity_scores(similarity_scores, image_size, image_processor)
similarity_map = similarity_scores[0]
return similarity_prob, similarity_map
def save_similarity_map(image, similarity_map, save_path, threshold=0.0, beta=0.5):
"""Save a JET heat-map of ``similarity_map`` overlaid on ``image`` (path or PIL.Image) as a PNG."""
import cv2
img_rgb = np.array(load_image(image))
sim = similarity_map.detach().float().cpu().numpy() if torch.is_tensor(similarity_map) else similarity_map
sim = cv2.resize(sim.astype(np.float32), (img_rgb.shape[1], img_rgb.shape[0]), interpolation=cv2.INTER_LINEAR)
sim = (np.clip(sim, threshold, 1.0) - threshold) / (1.0 - threshold + 1e-8)
sim_u8 = (sim * 255).astype(np.uint8)
heatmap = cv2.cvtColor(cv2.applyColorMap(sim_u8, cv2.COLORMAP_JET), cv2.COLOR_BGR2RGB)
Image.fromarray(cv2.addWeighted(img_rgb, 1 - beta, heatmap, beta, 0)).save(save_path)
@torch.no_grad()
def model_inference_multiple_text(image, text_list, model, image_processor, tokenizer):
probs, similarity_maps = [], []
for text in text_list:
prob, similarity_map = model_inference(image, text, model, image_processor, tokenizer)
probs.append(prob)
similarity_maps.append(similarity_map)
return torch.stack(probs), torch.stack(similarity_maps)
def interpolate_similarity_scores(similarity_scores, origin_size, image_processor):
(height, width) = origin_size
patch_size = int(similarity_scores.shape[-1] ** 0.5)
scores = similarity_scores.view(1, 1, patch_size, patch_size)
interpolated_scores = F.interpolate(
scores,
size=(height, width),
mode="bilinear",
align_corners=False,
)
interpolated_scores = interpolated_scores.squeeze(1)
return interpolated_scores
def dicom_to_pil_image(input_file_path):
"""Extract a grayscale image from a DICOM file and return it as a PIL.Image."""
import cv2
import pydicom
dcm_file = pydicom.dcmread(input_file_path)
raw_image = dcm_file.pixel_array
if raw_image.ndim != 2 or raw_image.size == 0:
raise ValueError(f"Expecting a non-empty single-channel DICOM image, got shape {raw_image.shape}")
raw_image = raw_image - raw_image.min()
rescaled_image = (raw_image / raw_image.max() * 255).astype(np.uint8)
if dcm_file.PhotometricInterpretation == "MONOCHROME1":
rescaled_image = cv2.bitwise_not(rescaled_image)
return Image.fromarray(cv2.equalizeHist(rescaled_image))
def load_image(image):
"""Load an image from a path (.png/.jpg/.jpeg/.dcm) or a PIL.Image object."""
if isinstance(image, str):
if image.lower().endswith(".dcm"):
image = dicom_to_pil_image(image)
elif image.lower().endswith((".png", ".jpg", ".jpeg")):
image = Image.open(image)
else:
raise ValueError(f"Invalid image type: {image}")
elif not isinstance(image, Image.Image):
raise ValueError(f"Invalid image type: {type(image)}")
if image.size[0] <= 0 or image.size[1] <= 0:
raise ValueError(f"Image has invalid size (W: {image.size[0]}, H: {image.size[1]}).")
# per-image min-max stretch to 8 bit, as in the zero-shot evaluation (also handles 16-bit inputs)
array = np.asarray(image, dtype=np.float32)
array = (array - array.min()) / max(float(array.max() - array.min()), 1e-8) * 255.0
return Image.fromarray(np.rint(array).astype(np.uint8)).convert("RGB")
|