app.py

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)