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

#---------------------------------
#++++++++     Model     ++++++++++
#---------------------------------

def load_biomedclip_model():
    """Loads the BiomedCLIP model and prepares it with LeGrad."""
    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
    return model, preprocess, tokenizer, device

def classify_image_with_biomedclip(editor_value, prompts, model, preprocess, tokenizer, device):
    """Classifies the image with the given text prompts using BiomedCLIP."""
    if editor_value is None:
        return None, None

    image = editor_value["composite"]

    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)

    image_input = preprocess(image).unsqueeze(0).to(device)
    text_inputs = tokenizer(prompts).to(device)

    text_embeddings = model.encode_text(text_inputs, normalize=True)
    image_embeddings = model.encode_image(image_input, normalize=True)

    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()]
    )

    explanation_maps = explanation_maps.squeeze(0).detach().cpu().numpy()
    explanation_map = (explanation_maps * 255).astype(np.uint8)

    return probabilities, explanation_map

def prepare_output_image(image, explanation_map):
    """Prepares the output image by blending the original image with the explanation map."""
    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()

    explanation_image_resized = cv2.resize(
        explanation_image, (image.width, image.height)
    )

    explanation_image_resized = cv2.normalize(
        explanation_image_resized, None, 0, 255, cv2.NORM_MINMAX
    )

    explanation_colormap = cv2.applyColorMap(
        explanation_image_resized.astype(np.uint8), cv2.COLORMAP_JET
    )

    image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

    alpha = 0.5
    blended_image = cv2.addWeighted(image_cv, 1 - alpha, explanation_colormap, alpha, 0)

    blended_image_rgb = cv2.cvtColor(blended_image, cv2.COLOR_BGR2RGB)
    output_image = Image.fromarray(blended_image_rgb)
    return output_image

#---------------------------------
#++++++++     Gradio     ++++++++++
#---------------------------------

def update_output(editor_value, prompts_input, model, preprocess, tokenizer, device):
    """Main function to update the output based on image and prompts."""
    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, model, preprocess, tokenizer, device
    )

    if probabilities is None:
        return None, "Please upload and annotate an image."

    prob_text = "\n".join(
        [
            f"{prompt}: {prob*100:.2f}%"
            for prompt, prob in zip(prompts_list, probabilities)
        ]
    )

    image = editor_value["composite"]
    output_image = prepare_output_image(image, explanation_map)

    return output_image, prob_text

def clear_inputs():
    return None, ""

# Load model outside the Gradio blocks
model, preprocess, tokenizer, device = load_biomedclip_model()


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("---")

    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("---")

    with gr.Row():
        with gr.Column():
            image_editor = gr.ImageEditor(
                label="Upload and Annotate Image",
                type="pil",
                interactive=True,
                mirror_webcam=False,
                layers=False,
                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
            )

    inputs = [image_editor, prompts_input]
    outputs = [output_image, prob_text]
    submit_button.click(fn=update_output, inputs=inputs, outputs=outputs, 
                    _js=None,
                    api_name=None,
                    scroll_to_output=True,
                    show_progress=True,
                    queue=True,
                    batch=False,
                    preprocess=True,
                    postprocess=True,
                    cancels=None,
                    show_loading_status=True,
                    scroll_to_output_id=None,
                    model=model, preprocess=preprocess, tokenizer=tokenizer, device=device
                    )
if __name__ == "__main__":
    demo.launch(share=True)