import gradio as gr import torch from PIL import Image import open_clip import numpy as np from LeGrad.legrad import LeWrapper, LePreprocess import cv2 import numpy as np from PIL import Image # Load BiomedCLIP model device = "cuda" if torch.cuda.is_available() else "cpu" model_name = "hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224" model, preprocess = open_clip.create_model_from_pretrained( model_name=model_name, device=device ) tokenizer = open_clip.get_tokenizer(model_name=model_name) model = LeWrapper(model) # Equip the model with LeGrad preprocess = LePreprocess( preprocess=preprocess, image_size=448 ) # Optional higher-res preprocessing def classify_image_with_biomedclip(editor_value, prompts): # editor_value is a dict with keys: 'background', 'layers', 'composite' # The 'composite' key contains the final annotated image if editor_value is None: return None, None # Get the composite image (background + annotations) image = editor_value["composite"] # Ensure image is in PIL format if not isinstance(image, Image.Image): image = Image.fromarray(image) # Preprocess and encode the image image_input = preprocess(image).unsqueeze(0).to(device) text_inputs = tokenizer(prompts).to(device) # Encode text and image text_embeddings = model.encode_text(text_inputs, normalize=True) image_embeddings = model.encode_image(image_input, normalize=True) # Generate probabilities (optional - not required for LeGrad explanations but included for completeness) similarity = ( model.logit_scale.exp() * image_embeddings @ text_embeddings.T ).softmax(dim=-1) probabilities = similarity[0].detach().cpu().numpy() explanation_maps = model.compute_legrad_clip( image=image_input, text_embedding=text_embeddings[probabilities.argmax()] ) # Convert explanation maps to heatmap explanation_maps = explanation_maps.squeeze(0).detach().cpu().numpy() explanation_map = (explanation_maps * 255).astype(np.uint8) # Rescale to [0, 255] return probabilities, explanation_map def update_output(editor_value, prompts_input): prompts_list = [p.strip() for p in prompts_input.split(",") if p.strip()] if not prompts_list: return None, "Please enter at least one prompt." probabilities, explanation_map = classify_image_with_biomedclip( editor_value, prompts_list ) if probabilities is None: return None, "Please upload and annotate an image." # Create probability display prob_text = "\n".join( [ f"{prompt}: {prob*100:.2f}%" for prompt, prob in zip(prompts_list, probabilities) ] ) # Prepare the explanation map overlay image = editor_value["composite"] if not isinstance(image, Image.Image): image = Image.fromarray(image) explanation_image = explanation_map[0] if isinstance(explanation_image, torch.Tensor): explanation_image = explanation_image.cpu().numpy() # Resize the explanation map to match the size of the original image explanation_image_resized = cv2.resize( explanation_image, (image.width, image.height) ) # Normalize the explanation map for proper colormap application explanation_image_resized = cv2.normalize( explanation_image_resized, None, 0, 255, cv2.NORM_MINMAX ) # Apply the colormap (e.g., COLORMAP_JET) explanation_colormap = cv2.applyColorMap( explanation_image_resized.astype(np.uint8), cv2.COLORMAP_JET ) # Convert the original image to a format that OpenCV understands (RGB to BGR) image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) # Blend the original image and the colormap alpha = 0.5 # Transparency factor blended_image = cv2.addWeighted(image_cv, 1 - alpha, explanation_colormap, alpha, 0) # Convert back to RGB for displaying with PIL or matplotlib blended_image_rgb = cv2.cvtColor(blended_image, cv2.COLOR_BGR2RGB) output_image = Image.fromarray(blended_image_rgb) return output_image, prob_text def clear_inputs(): return None, "" with gr.Blocks() as demo: gr.Markdown( "# ✨ Visual Prompt Engineering for Medical Vision Language Models in Radiology ✨", elem_id="main-header", ) gr.Markdown( "This tool applies **visual prompt engineering to improve the classification of medical images using the BiomedCLIP**[3], the current state of the art in zero-shot biomedical image classification. By uploading biomedical images (e.g., chest X-rays), you can manually annotate areas of interest directly on the image. These annotations serve as visual prompts, which guide the model's attention on the region of interest. This technique improves the model's ability to focus on subtle yet important details.\n\n" "After annotating and inputting text prompts (e.g., 'A chest X-ray with a benign/malignant lung nodule indicated by a red circle'), the tool returns classification results. These results are accompanied by **explainability maps** generated by **LeGrad** [3], which show where the model focused its attention, conditioned on the highest scoring text prompt. This helps to better interpret the model's decision-making process.\n\n" "In our paper **[Visual Prompt Engineering for Medical Vision Language Models in Radiology](https://arxiv.org/pdf/2408.15802)**, we show, that visual prompts such as arrows, circles, and contours improve the zero-shot classification of biomedical vision language models in radiology." ) gr.Markdown("---") # Horizontal rule for separation gr.Markdown( "## 📝 **How It Works**:\n" "1. **Upload** a biomedical image.\n" "2. **Annotate** the image using the built-in editor to highlight regions of interest.\n" "3. **Enter text prompts** separated by comma (e.g., 'A chest X-ray with a (benign/malignant) lung nodule indicated by a red circle').\n" "4. **Submit** to get class probabilities and an explainability map conditioned on the highest scoring text prompt." ) gr.Markdown("---") # Horizontal rule for separation with gr.Row(): with gr.Column(): image_editor = gr.ImageEditor( label="Upload and Annotate Image", type="pil", interactive=True, mirror_webcam=False, layers=False, # placeholder="Upload an image", scale=2, ) prompts_input = gr.Textbox( placeholder="Enter prompts, comma-separated", label="Text Prompts" ) submit_button = gr.Button("Submit", variant="primary") with gr.Column(): output_image = gr.Image( type="pil", label="Output Image with Explanation Map", ) prob_text = gr.Textbox( label="Class Probabilities", interactive=False, lines=10 ) # Manually trigger the computation with the submit button inputs = [image_editor, prompts_input] outputs = [output_image, prob_text] submit_button.click(fn=update_output, inputs=inputs, outputs=outputs) gr.Markdown("---") # Horizontal rule for separation gr.Markdown("### 📝 **References**:\n") gr.Markdown( "[1] Denner, S., Bujotzek, M., Bounias, D., Zimmerer, D., Stock, R., Jäger, P.F. and Maier-Hein, K., 2024. **Visual Prompt Engineering for Medical Vision Language Models in Radiology**. arXiv preprint arXiv:2408.15802." ) gr.Markdown( "[2] Zhang, S., Xu, Y., Usuyama, N., Bagga, J., Tinn, R., Preston, S., Rao, R., Wei, M., Valluri, N., Wong, C. and Lungren, M.P., 2023. **Large-scale domain-specific pretraining for biomedical vision-language processing**. arXiv preprint arXiv:2303.00915, 2(3), p.6.\n" ) gr.Markdown( "[3] Bousselham, W., Boggust, A., Chaybouti, S., Strobelt, H. and Kuehne, H., 2024. **LeGrad: An Explainability Method for Vision Transformers via Feature Formation Sensitivity**. arXiv preprint arXiv:2404.03214." ) if __name__ == "__main__": demo.launch(share=True)