File size: 5,077 Bytes
5718512 |
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 |
import shutil
import traceback
from io import BytesIO
from urllib.parse import urlparse
import cv2
import numpy as np
import pydicom
import requests
import torch
import torch.nn.functional as F
from PIL import Image
from transformers import BitImageProcessor, BlipImageProcessor
@torch.no_grad()
def model_inference(image, text, model, image_processor, tokenizer):
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.sigmoid()[0]
return similarity_prob, similarity_map
@torch.no_grad()
def model_inference_multiple_text(image, text_list, model, image_processor, tokenizer):
# TODO: batch inference
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)
if isinstance(image_processor, BlipImageProcessor):
# XrayDINOv2
interpolated_scores = F.interpolate(
scores,
size=(height, width),
mode="bilinear",
align_corners=False,
)
interpolated_scores = interpolated_scores.squeeze(1)
elif isinstance(image_processor, BitImageProcessor):
shortest = min(height, width)
interpolated_scores = F.interpolate(
scores,
size=(shortest, shortest),
mode="bilinear",
align_corners=False,
)
cropped_left = (width - shortest) // 2
cropped_top = (height - shortest) // 2
original_size_map = torch.ones(height, width) * -999
original_size_map[
cropped_top : cropped_top + shortest, cropped_left : cropped_left + shortest
] = interpolated_scores.view(shortest, shortest)
interpolated_scores = original_size_map
interpolated_scores = interpolated_scores.unsqueeze(0)
return interpolated_scores
# copy from https://github.com/MIT-LCP/mimic-code/issues/1013
def dicom_to_pil_image(input_file_path, save_dir=None):
"""
Extract the image from a DICOM file and return it as a PIL.Image object.
Args:
input_file_path (str): Path to the input DICOM file.
Returns:
PIL.Image.Image: Processed image.
"""
try:
# Read the DICOM and extract the image.
dcm_file = pydicom.dcmread(input_file_path)
raw_image = dcm_file.pixel_array
assert len(raw_image.shape) == 2, "Expecting single channel (grayscale) image."
# Normalize pixels to be in [0, 255].
raw_image = raw_image - raw_image.min()
normalized_image = raw_image / raw_image.max()
rescaled_image = (normalized_image * 255).astype(np.uint8)
# Correct image inversion.
if dcm_file.PhotometricInterpretation == "MONOCHROME1":
rescaled_image = cv2.bitwise_not(rescaled_image)
# Perform histogram equalization.
final_image = cv2.equalizeHist(rescaled_image)
# Convert to PIL Image and return
image = Image.fromarray(final_image)
if save_dir is not None:
shutil.copy2(input_file_path, save_dir)
return image
except Exception:
print(traceback.format_exc())
def load_image(image):
"""
Load an image from a file path or a PIL.Image object.
Args:
image (str or PIL.Image.Image): Path to the image file or a PIL.Image object.
Returns:
PIL.Image.Image: Processed image.
"""
if isinstance(image, str):
if image.lower().endswith(".dcm"):
image = dicom_to_pil_image(image)
elif (
image.lower().endswith(".png")
or image.lower().endswith(".jpg")
or image.lower().endswith(".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)}")
return image
|