| """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]}).") |
|
|
| |
| 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") |
|
|
