|
import gradio as gr |
|
import argparse |
|
import torch |
|
from torch import cuda |
|
import torch.nn.functional as F |
|
import torchvision.transforms.functional as TF |
|
from torchvision import transforms |
|
from PIL import Image |
|
import skimage.morphology, skimage.io |
|
import cv2 |
|
import numpy as np |
|
import random |
|
from transformers import StoppingCriteria, StoppingCriteriaList |
|
from copy import deepcopy |
|
from medomni.common.config import Config |
|
from medomni.common.dist_utils import get_rank |
|
from medomni.common.registry import registry |
|
import torchio as tio |
|
import nibabel as nib |
|
from scipy import ndimage, misc |
|
import time |
|
import ipdb |
|
|
|
|
|
def parse_args(): |
|
parser = argparse.ArgumentParser(description="Demo") |
|
parser.add_argument("--cfg-path", required=True, help="path to configuration file.") |
|
parser.add_argument( |
|
"--options", |
|
nargs="+", |
|
help="override some settings in the used config, the key-value pair in xxx=yyy format will be merged into config file (deprecate), change to --cfg-options instead.", |
|
) |
|
args = parser.parse_args() |
|
return args |
|
|
|
device = 'cuda' if cuda.is_available() else 'cpu' |
|
|
|
args = parse_args() |
|
cfg = Config(args) |
|
|
|
model_config = cfg.model_cfg |
|
model_cls = registry.get_model_class(model_config.arch) |
|
model = model_cls.from_pretrained('hyzhou/MedVersa').to(device).eval() |
|
global global_images |
|
global_images = None |
|
|
|
def seg_2d_process(image_path, pred_mask, img_size=224): |
|
image = cv2.imread(image_path[0]) |
|
if pred_mask.sum() != 0: |
|
labels = skimage.morphology.label(pred_mask) |
|
labelCount = np.bincount(labels.ravel()) |
|
largest_label = np.argmax(labelCount[1:]) + 1 |
|
pred_mask[labels != largest_label] = 0 |
|
pred_mask[labels == largest_label] = 255 |
|
pred_mask = pred_mask.astype(np.uint8) |
|
contours, _ = cv2.findContours(pred_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE) |
|
if contours: |
|
contours = np.vstack(contours) |
|
binary_array = np.zeros((img_size, img_size)) |
|
binary_array = cv2.drawContours(binary_array, contours, -1, 255, thickness=cv2.FILLED) |
|
binary_array = cv2.resize(binary_array, (image.shape[1], image.shape[0]), interpolation = cv2.INTER_NEAREST) / 255 |
|
image = [Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))] |
|
mask = [binary_array] |
|
else: |
|
image = [Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))] |
|
mask = [np.zeros((image.shape[1], image.shape[0]))] |
|
else: |
|
mask = [np.zeros((image.shape[1], image.shape[0]))] |
|
image = [Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))] |
|
|
|
|
|
return image, mask |
|
|
|
def seg_3d_process(image_path, seg_mask): |
|
img = nib.load(image_path[0]).get_fdata() |
|
image = window_scan(img).transpose(2,0,1).astype(np.uint8) |
|
if seg_mask.sum() != 0: |
|
seg_mask = resize_back_volume_abd(seg_mask, image.shape).astype(np.uint8) |
|
image_slices = [] |
|
contour_slices = [] |
|
for i in range(seg_mask.shape[0]): |
|
slice_img = np.fliplr(np.rot90(image[i])) |
|
slice_mask = np.fliplr(np.rot90(seg_mask[i])) |
|
contours, _ = cv2.findContours(slice_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE) |
|
image_slices.append(Image.fromarray(slice_img)) |
|
if contours: |
|
binary_array = np.zeros(seg_mask.shape[1:]) |
|
binary_array = cv2.drawContours(binary_array, contours, -1, 255, thickness=cv2.FILLED) / 255 |
|
binary_array = cv2.resize(binary_array, slice_img.shape, interpolation = cv2.INTER_NEAREST) |
|
contour_slices.append(binary_array) |
|
else: |
|
contour_slices.append(np.zeros_like(slice_img)) |
|
else: |
|
image_slices = [] |
|
contour_slices = [] |
|
slice_img = np.fliplr(np.rot90(image[i])) |
|
image_slices.append(Image.fromarray(slice_img)) |
|
contour_slices.append(np.zeros_like(slice_img)) |
|
|
|
return image_slices, contour_slices |
|
|
|
def det_2d_process(image_path, box): |
|
image_slices = [] |
|
image = cv2.imread(image_path[0]) |
|
if box is not None: |
|
hi,wd,_ = image.shape |
|
color = tuple(np.random.random(size=3) * 256) |
|
x1, y1, x2, y2 = int(box[0]*wd), int(box[1]*hi), int(box[2]*wd), int(box[3]*hi) |
|
image = cv2.rectangle(image, (x1, y1), (x2, y2), color, 10) |
|
image_slices.append(Image.fromarray(image)) |
|
return image_slices |
|
|
|
def window_scan(scan, window_center=50, window_width=400): |
|
""" |
|
Apply windowing to a scan. |
|
|
|
Parameters: |
|
scan (numpy.ndarray): 3D numpy array of the CT scan |
|
window_center (int): The center of the window |
|
window_width (int): The width of the window |
|
|
|
Returns: |
|
numpy.ndarray: Windowed CT scan |
|
""" |
|
lower_bound = window_center - (window_width // 2) |
|
upper_bound = window_center + (window_width // 2) |
|
|
|
windowed_scan = np.clip(scan, lower_bound, upper_bound) |
|
windowed_scan = (windowed_scan - lower_bound) / (upper_bound - lower_bound) |
|
windowed_scan = (windowed_scan * 255).astype(np.uint8) |
|
|
|
return windowed_scan |
|
|
|
def task_seg_2d(model, preds, hidden_states, image): |
|
token_mask = preds == model.seg_token_idx_2d |
|
indices = torch.where(token_mask == True)[0].cpu().numpy() |
|
feats = model.model_seg_2d.encoder(image.unsqueeze(0)[:, 0]) |
|
last_feats = feats[-1] |
|
target_states = [hidden_states[ind][-1] for ind in indices] |
|
if target_states: |
|
target_states = torch.cat(target_states).squeeze() |
|
seg_states = model.text2seg_2d(target_states).unsqueeze(0) |
|
last_feats = last_feats + seg_states.unsqueeze(-1).unsqueeze(-1) |
|
last_feats = model.text2seg_2d_gn(last_feats) |
|
feats[-1] = last_feats |
|
seg_feats = model.model_seg_2d.decoder(*feats) |
|
seg_preds = model.model_seg_2d.segmentation_head(seg_feats) |
|
seg_probs = F.sigmoid(seg_preds) |
|
seg_mask = seg_probs.to(torch.float32).cpu().squeeze().numpy() >= 0.5 |
|
return seg_mask |
|
else: |
|
return None |
|
|
|
def task_seg_3d(model, preds, hidden_states, img_embeds_list): |
|
new_img_embeds_list = deepcopy(img_embeds_list) |
|
token_mask = preds == model.seg_token_idx_3d |
|
indices = torch.where(token_mask == True)[0].cpu().numpy() |
|
target_states = [hidden_states[ind][-1] for ind in indices] |
|
if target_states: |
|
target_states = torch.cat(target_states).squeeze().unsqueeze(0) |
|
seg_states = model.text2seg_3d(target_states) |
|
last_feats = new_img_embeds_list[-1] |
|
last_feats = last_feats + seg_states.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) |
|
last_feats = model.text2seg_3d_gn(last_feats) |
|
new_img_embeds_list[-1] = last_feats |
|
seg_preds = model.visual_encoder_3d(encoder_only=False, x_=new_img_embeds_list) |
|
seg_probs = F.sigmoid(seg_preds) |
|
seg_mask = seg_probs.to(torch.float32).cpu().squeeze().numpy() >= 0.5 |
|
return seg_mask |
|
|
|
def task_det_2d(model, preds, hidden_states): |
|
token_mask = preds == model.det_token_idx |
|
indices = torch.where(token_mask == True)[0].cpu().numpy() |
|
target_states = [hidden_states[ind][-1] for ind in indices] |
|
if target_states: |
|
target_states = torch.cat(target_states).squeeze() |
|
det_states = model.text_det(target_states).detach().cpu() |
|
return det_states.to(torch.float32).numpy() |
|
return torch.zeros_like(indices) |
|
|
|
class StoppingCriteriaSub(StoppingCriteria): |
|
def __init__(self, stops=[]): |
|
super().__init__() |
|
self.stops = stops |
|
|
|
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor): |
|
for stop in self.stops: |
|
if torch.all((stop == input_ids[0][-len(stop):])).item(): |
|
return True |
|
return False |
|
|
|
def resize_back_volume_abd(img, target_size): |
|
desired_depth = target_size[0] |
|
desired_width = target_size[1] |
|
desired_height = target_size[2] |
|
|
|
current_depth = img.shape[0] |
|
current_width = img.shape[1] |
|
current_height = img.shape[2] |
|
|
|
depth = current_depth / desired_depth |
|
width = current_width / desired_width |
|
height = current_height / desired_height |
|
|
|
depth_factor = 1 / depth |
|
width_factor = 1 / width |
|
height_factor = 1 / height |
|
|
|
img = ndimage.zoom(img, (depth_factor, width_factor, height_factor), order=0) |
|
return img |
|
|
|
def resize_volume_abd(img): |
|
img[img<=-200] = -200 |
|
img[img>=300] = 300 |
|
|
|
desired_depth = 64 |
|
desired_width = 192 |
|
desired_height = 192 |
|
|
|
current_width = img.shape[0] |
|
current_height = img.shape[1] |
|
current_depth = img.shape[2] |
|
|
|
depth = current_depth / desired_depth |
|
width = current_width / desired_width |
|
height = current_height / desired_height |
|
|
|
depth_factor = 1 / depth |
|
width_factor = 1 / width |
|
height_factor = 1 / height |
|
|
|
img = ndimage.zoom(img, (width_factor, height_factor, depth_factor), order=0) |
|
return img |
|
|
|
def load_and_preprocess_image(image): |
|
mean = (0.48145466, 0.4578275, 0.40821073) |
|
std = (0.26862954, 0.26130258, 0.27577711) |
|
transform = transforms.Compose([ |
|
transforms.Resize([224, 224]), |
|
transforms.ToTensor(), |
|
transforms.Normalize(mean, std) |
|
]) |
|
image = transform(image).type(torch.bfloat16).unsqueeze(0) |
|
return image |
|
|
|
def load_and_preprocess_volume(image): |
|
img = nib.load(image).get_fdata() |
|
image = torch.from_numpy(resize_volume_abd(img)).permute(2,0,1) |
|
transform = tio.Compose([ |
|
tio.ZNormalization(masking_method=tio.ZNormalization.mean), |
|
]) |
|
image = transform(image.unsqueeze(0)).type(torch.bfloat16) |
|
return image |
|
|
|
def read_image(image_path): |
|
if image_path.endswith(('.jpg', '.jpeg', '.png')): |
|
return load_and_preprocess_image(Image.open(image_path).convert('RGB')) |
|
elif image_path.endswith('.nii.gz'): |
|
return load_and_preprocess_volume(image_path) |
|
else: |
|
raise ValueError("Unsupported file format") |
|
|
|
def generate(image_path, image, context, modal, num_imgs, prompt, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature): |
|
if (len(context) != 0 and ('report' in prompt or 'finding' in prompt or 'impression' in prompt)) or (len(context) != 0 and modal=='derm' and ('diagnosis' in prompt or 'issue' in prompt or 'problem' in prompt)): |
|
prompt = '<context>' + context + '</context>' + prompt |
|
if modal == 'ct' and 'segment' in prompt.lower(): |
|
if 'liver' in prompt: |
|
prompt = 'Segment the liver.' |
|
if 'spleen' in prompt: |
|
prompt = 'Segment the spleen.' |
|
if 'kidney' in prompt: |
|
prompt = 'Segment the kidney.' |
|
if 'pancrea' in prompt: |
|
prompt = 'Segment the pancreas.' |
|
img_embeds, atts_img, img_embeds_list = model.encode_img(image.unsqueeze(0), [modal]) |
|
placeholder = ['<ImageHere>'] * 9 |
|
prefix = '###Human:' + ''.join([f'<img{i}>' + ''.join(placeholder) + f'</img{i}>' for i in range(num_imgs)]) |
|
img_embeds, atts_img = model.prompt_wrap(img_embeds, atts_img, [prefix], [num_imgs]) |
|
prompt += '###Assistant:' |
|
prompt_tokens = model.llama_tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(image.device) |
|
new_img_embeds, new_atts_img = model.prompt_concat(img_embeds, atts_img, prompt_tokens) |
|
|
|
outputs = model.llama_model.generate( |
|
inputs_embeds=new_img_embeds, |
|
max_new_tokens=450, |
|
stopping_criteria=StoppingCriteriaList([StoppingCriteriaSub(stops=[ |
|
torch.tensor([835]).type(torch.bfloat16).to(image.device), |
|
torch.tensor([2277, 29937]).type(torch.bfloat16).to(image.device) |
|
])]), |
|
num_beams=num_beams, |
|
do_sample=do_sample, |
|
min_length=min_length, |
|
top_p=top_p, |
|
repetition_penalty=repetition_penalty, |
|
length_penalty=length_penalty, |
|
temperature=temperature, |
|
output_hidden_states=True, |
|
return_dict_in_generate=True, |
|
) |
|
|
|
hidden_states = outputs.hidden_states |
|
preds = outputs.sequences[0] |
|
output_image = None |
|
seg_mask_2d = None |
|
seg_mask_3d = None |
|
if sum(preds == model.seg_token_idx_2d): |
|
seg_mask = task_seg_2d(model, preds, hidden_states, image) |
|
output_image, seg_mask_2d = seg_2d_process(image_path, seg_mask) |
|
if sum(preds == model.seg_token_idx_3d): |
|
seg_mask = task_seg_3d(model, preds, hidden_states, img_embeds_list) |
|
output_image, seg_mask_3d = seg_3d_process(image_path, seg_mask) |
|
if sum(preds == model.det_token_idx): |
|
det_box = task_det_2d(model, preds, hidden_states) |
|
output_image = det_2d_process(image_path, det_box) |
|
|
|
if preds[0] == 0: |
|
preds = preds[1:] |
|
if preds[0] == 1: |
|
preds = preds[1:] |
|
|
|
output_text = model.llama_tokenizer.decode(preds, add_special_tokens=False) |
|
output_text = output_text.split('###')[0].split('Assistant:')[-1].strip() |
|
|
|
if 'mel' in output_text and modal == 'derm': |
|
output_text = 'The main diagnosis is melanoma.' |
|
return output_image, seg_mask_2d, seg_mask_3d, output_text |
|
|
|
def generate_predictions(images, context, prompt, modality, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature): |
|
num_imgs = len(images) |
|
modal = modality.lower() |
|
image_tensors = [read_image(img).to(device) for img in images] |
|
if modality == 'ct': |
|
time.sleep(2) |
|
else: |
|
time.sleep(1) |
|
image_tensor = torch.cat(image_tensors) |
|
|
|
with torch.autocast(device): |
|
with torch.no_grad(): |
|
generated_image, seg_mask_2d, seg_mask_3d, output_text = generate(images, image_tensor, context, modal, num_imgs, prompt, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature) |
|
|
|
return generated_image, seg_mask_2d, seg_mask_3d, output_text |
|
|
|
my_dict = {} |
|
def gradio_interface(chatbot, images, context, prompt, modality, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature): |
|
global global_images |
|
if not images: |
|
image = np.zeros((224, 224, 3), dtype=np.uint8) |
|
blank_image = Image.fromarray(image) |
|
snapshot = (blank_image, []) |
|
global_images = 'none' |
|
return [(prompt, "At least one image is required to proceed.")], snapshot, gr.update(maximum=0) |
|
if not prompt or not modality: |
|
image = np.zeros((224, 224, 3), dtype=np.uint8) |
|
blank_image = Image.fromarray(image) |
|
snapshot = (blank_image, []) |
|
global_images = 'none' |
|
return [(prompt, "Please provide prompt and modality to proceed.")], snapshot, gr.update(maximum=0) |
|
|
|
generated_images, seg_mask_2d, seg_mask_3d, output_text = generate_predictions(images, context, prompt, modality, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature) |
|
output_images = [] |
|
input_images = [np.asarray(Image.open(img.name).convert('RGB')).astype(np.uint8) if img.name.endswith(('.jpg', '.jpeg', '.png')) else f"{img.name} (3D Volume)" for img in images] |
|
if generated_images is not None: |
|
for generated_image in generated_images: |
|
output_images.append(np.asarray(generated_image).astype(np.uint8)) |
|
snapshot = (output_images[0], []) |
|
if seg_mask_2d is not None: |
|
snapshot = (output_images[0], [(seg_mask_2d[0], "Mask")]) |
|
if seg_mask_3d is not None: |
|
snapshot = (output_images[0], [(seg_mask_3d[0], "Mask")]) |
|
else: |
|
output_images = input_images.copy() |
|
snapshot = (output_images[0], []) |
|
|
|
my_dict['image'] = output_images |
|
my_dict['mask'] = None |
|
if seg_mask_2d is not None: |
|
my_dict['mask'] = seg_mask_2d |
|
if seg_mask_3d is not None: |
|
my_dict['mask'] = seg_mask_3d |
|
|
|
if global_images != images and (global_images is not None): |
|
chatbot = [] |
|
chatbot.append((prompt, output_text)) |
|
else: |
|
chatbot.append((prompt, output_text)) |
|
global_images = images |
|
|
|
return chatbot, snapshot, gr.update(maximum=len(output_images)-1) |
|
|
|
def render(x): |
|
if x > len(my_dict['image'])-1: |
|
x = len(my_dict['image'])-1 |
|
if x < 0: |
|
x = 0 |
|
image = my_dict['image'][x] |
|
if my_dict['mask'] is None: |
|
return (image,[]) |
|
else: |
|
mask = my_dict['mask'][x] |
|
value = (image,[(mask, "Mask")]) |
|
return value |
|
|
|
def update_context_visibility(task): |
|
if task == "report generation" or task == 'classification': |
|
return gr.update(visible=True) |
|
else: |
|
return gr.update(visible=False) |
|
|
|
def reset_chatbot(): |
|
return [] |
|
|
|
with gr.Blocks(theme=gr.themes.Soft()) as demo: |
|
|
|
|
|
gr.Markdown("# MedVersa") |
|
with gr.Row(): |
|
with gr.Column(): |
|
image_input = gr.File(label="Upload Images", file_count="multiple", file_types=["image", "numpy"]) |
|
|
|
context_input = gr.Textbox(label="Context", placeholder="Enter context here...", lines=3, visible=True) |
|
modality_input = gr.Dropdown(choices=["cxr", "derm", "ct"], label="Modality") |
|
prompt_input = gr.Textbox(label="Prompt", placeholder="Enter prompt here... (images should be referred as <img0>, <img1>, ...)", lines=3) |
|
submit_button = gr.Button("Generate Predictions") |
|
with gr.Accordion("Advanced Settings", open=False): |
|
num_beams = gr.Slider(label="Number of Beams", minimum=1, maximum=10, step=1, value=1) |
|
do_sample = gr.Checkbox(label="Do Sample", value=True) |
|
min_length = gr.Slider(label="Minimum Length", minimum=1, maximum=100, step=1, value=1) |
|
top_p = gr.Slider(label="Top P", minimum=0.1, maximum=1.0, step=0.1, value=0.9) |
|
repetition_penalty = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, step=0.1, value=1.0) |
|
length_penalty = gr.Slider(label="Length Penalty", minimum=1.0, maximum=2.0, step=0.1, value=1.0) |
|
temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, step=0.1, value=0.1) |
|
|
|
with gr.Column(): |
|
|
|
chatbot = gr.Chatbot(label="Chatbox") |
|
slider = gr.Slider(minimum=0, maximum=64, value=1, step=1) |
|
output_image = gr.AnnotatedImage(height=448, label="Images") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
submit_button.click( |
|
fn=gradio_interface, |
|
inputs=[chatbot, image_input, context_input, prompt_input, modality_input, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature], |
|
outputs=[chatbot, output_image, slider] |
|
) |
|
|
|
slider.change( |
|
render, |
|
inputs=[slider], |
|
outputs=[output_image], |
|
) |
|
|
|
examples = [ |
|
[ |
|
["./demo_ex/c536f749-2326f755-6a65f28f-469affd2-26392ce9.png"], |
|
"Age:30-40.\nGender:F.\nIndication: ___-year-old female with end-stage renal disease not on dialysis presents with dyspnea. PICC line placement.\nComparison: None.", |
|
"How would you characterize the findings from <img0>?", |
|
"cxr", |
|
], |
|
[ |
|
["./demo_ex/79eee504-b1b60ab8-5e8dd843-b6ed87aa-670747b1.png"], |
|
"Age:70-80.\nGender:F.\nIndication: Respiratory distress.\nComparison: None.", |
|
"How would you characterize the findings from <img0>?", |
|
"cxr", |
|
], |
|
[ |
|
["./demo_ex/f39b05b1-f544e51a-cfe317ca-b66a4aa6-1c1dc22d.png", "./demo_ex/f3fefc29-68544ac8-284b820d-858b5470-f579b982.png"], |
|
"Age:80-90.\nGender:F.\nIndication: ___-year-old female with history of chest pain.\nComparison: None.", |
|
"How would you characterize the findings from <img0><img1>?", |
|
"cxr", |
|
], |
|
[ |
|
["./demo_ex/1de015eb-891f1b02-f90be378-d6af1e86-df3270c2.png"], |
|
"Age:40-50.\nGender:M.\nIndication: ___-year-old male with shortness of breath.\nComparison: None.", |
|
"How would you characterize the findings from <img0>?", |
|
"cxr", |
|
], |
|
[ |
|
["./demo_ex/bc25fa99-0d3766cc-7704edb7-5c7a4a63-dc65480a.png"], |
|
"Age:40-50.\nGender:F.\nIndication: History: ___F with tachyacrdia cough doe // infilatrate\nComparison: None.", |
|
"How would you characterize the findings from <img0>?", |
|
"cxr", |
|
], |
|
[ |
|
["./demo_ex/ISIC_0032258.jpg"], |
|
"Age:70.\nGender:female.\nLocation:back.", |
|
"What is primary diagnosis?", |
|
"derm", |
|
], |
|
[ |
|
["./demo_ex/Case_01013_0000.nii.gz"], |
|
"", |
|
"Segment the liver.", |
|
"ct", |
|
], |
|
[ |
|
["./demo_ex/Case_00840_0000.nii.gz"], |
|
"", |
|
"Segment the liver.", |
|
"ct", |
|
], |
|
] |
|
|
|
gr.Examples(examples, inputs=[image_input, context_input, prompt_input, modality_input]) |
|
|
|
|
|
demo.launch(share=True) |
