Update app.py

app.py

@@ -1,109 +1,186 @@
+import gradio as gr
 import torch
 from PIL import Image
-import gradio as gr
-import spaces
-from transformers import AutoProcessor, AutoModel, CLIPVisionModel
-import torch.nn.functional as F
+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 tokenizer."""
-    biomedclip_model_name = 'microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224'
-    processor = AutoProcessor.from_pretrained(biomedclip_model_name)
-    config = AutoModel.from_pretrained(biomedclip_model_name).config
-    vision_model = CLIPVisionModel.from_pretrained(config.vision_config._name_or_path, torch_dtype=torch.float16).cuda().eval()
-    text_model = AutoModel.from_pretrained(config.text_config._name_or_path).cuda().eval()
-    return vision_model, text_model, processor
-
-def compute_similarity(image, text, vision_model, text_model, biomedclip_processor):
-    """Computes similarity scores using BiomedCLIP."""
-    with torch.no_grad():
-        inputs = biomedclip_processor(text=text, images=image, return_tensors="pt", padding=True).to(text_model.device)
-        text_embeds = text_model(**inputs).last_hidden_state[:,0,:] # Extract the [CLS] token
-        image_inputs = biomedclip_processor(images=image, return_tensors="pt").to(vision_model.device)
-        image_embeds = vision_model(**image_inputs).last_hidden_state[:,0,:] # Extract the image embedding
-    image_embeds = F.normalize(image_embeds, dim=-1)
-    text_embeds = F.normalize(text_embeds, dim=-1)
-    similarity = (text_embeds @ image_embeds.transpose(-1, -2)).squeeze()
-    return similarity
+    """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 gradio_reset(chat_state, img_list, similarity_output):
-    """Resets the chat state and image list."""
-    if chat_state is not None:
-        chat_state.messages = []
-    if img_list is not None:
-        img_list = []
-    return None, gr.update(value=None, interactive=True), gr.update(placeholder='Please upload your medical image first', interactive=False), gr.update(value="Upload & Start Analysis", interactive=True), chat_state, img_list, gr.update(value="", visible=False)
-
-def upload_img(gr_img, text_input, chat_state, similarity_output):
-    """Handles image upload."""
-    if gr_img is None:
-        return None, None, gr.update(interactive=True), chat_state, None, gr.update(visible=False)
-    img_list = [gr_img]
-    return gr.update(interactive=False), gr.update(interactive=True, placeholder='Type and press Enter'), gr.update(value="Start Analysis", interactive=False), chat_state, img_list, gr.update(visible=True)
-
-def gradio_ask(user_message, chatbot, chat_state):
-    """Handles user input."""
-    if not user_message:
-        return gr.update(interactive=True, placeholder='Input should not be empty!'), chatbot, chat_state
-    chatbot = chatbot + [[user_message, None]]
-    return '', chatbot, chat_state
-
-@spaces.GPU
-def gradio_answer(chatbot, chat_state, img_list, vision_model, text_model, biomedclip_processor, similarity_output):
-    """Computes and displays similarity scores."""
-    if not img_list:
-        return chatbot, chat_state, img_list, similarity_output
-
-    similarity_score = compute_similarity(img_list[0], chatbot[-1][0], vision_model, text_model, biomedclip_processor)
-    print(f'Similarity Score is: {similarity_score}')
-
-    similarity_text = f"Similarity Score: {similarity_score:.3f}"
-    chatbot[-1][1] = similarity_text
-    return chatbot, chat_state, img_list, gr.update(value=similarity_text, visible=True)
-
-
-title = """<h1 align="center">Medical Image Analysis Tool</h1>"""
-description = """<h3>Upload medical images, ask questions, and receive a similarity score.</h3>"""
-examples_list=[
-                    ["./case1.png", "Analyze the X-ray for any abnormalities."],
-                    ["./case2.jpg", "What type of disease may be present?"],
-                    ["./case1.png","What is the anatomical structure shown here?"]
-                ]
-
-# Load models and related resources outside of the Gradio block for loading on startup
-vision_model, text_model, biomedclip_processor = load_biomedclip_model()
+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."
 
-with gr.Blocks() as demo:
-    gr.Markdown(title)
-    gr.Markdown(description)
+    probabilities, explanation_map = classify_image_with_biomedclip(
+        editor_value, prompts_list, model, preprocess, tokenizer, device
+    )
 
-    with gr.Row():
-        with gr.Column(scale=0.5):
-            image = gr.Image(type="pil", label="Medical Image")
-            upload_button = gr.Button(value="Upload & Start Analysis", interactive=True, variant="primary")
-            clear = gr.Button("Restart")
+    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.Column():
-            chat_state = gr.State()
-            img_list = gr.State()
-            chatbot = gr.Chatbot(label='Medical Analysis')
-            text_input = gr.Textbox(label='Analysis Query', placeholder='Please upload your medical image first', interactive=False)
-            similarity_output = gr.Textbox(label="Similarity Score", visible=False, interactive=False)
-            gr.Examples(examples=examples_list, inputs=[image, text_input])
 
-    upload_button.click(upload_img, [image, text_input, chat_state, similarity_output], [image, text_input, upload_button, chat_state, img_list, similarity_output])
+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("---")
 
-    text_input.submit(gradio_ask, [text_input, chatbot, chat_state], [text_input, chatbot, chat_state]).then(
-        gradio_answer, [chatbot, chat_state, img_list, vision_model, text_model, biomedclip_processor, similarity_output], [chatbot, chat_state, img_list, similarity_output]
+    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."
     )
-    clear.click(gradio_reset, [chat_state, img_list, similarity_output], [chatbot, image, text_input, upload_button, chat_state, img_list, similarity_output], queue=False)
 
-demo.launch()
+    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)