bomedllamv2 integrated

app.py

@@ -9,16 +9,13 @@ import cv2
 import gradio as gr
 import numpy as np
 import spaces
-# import supervision as sv
 import torch
 from PIL import Image
 from tqdm import tqdm
-import sys
 from pathlib import Path
 from huggingface_hub import login
 from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
 
-
 token = os.getenv("HF_TOKEN")
 if token:
     login(token=token)
@@ -135,14 +132,14 @@ MODALITY_PROMPTS = {
    "Endoscopy": ["neoplastic polyp", "polyp", "non-neoplastic polyp"],
    "Fundus": ["optic cup", "optic disc"],
    "Dermoscopy": ["lesion", "melanoma"],
-   "OCT": ["edema"] }
+   "OCT": ["edema"]
+}
 
 
 def extract_modality_from_llm(llm_output):
     """Extract modality from LLM output and map it to BIOMEDPARSE_MODES"""
     llm_output = llm_output.lower()
     
-    # Direct modality mapping
     modality_keywords = {
         'ct': {
             'abdomen': 'CT-Abdomen',
@@ -156,7 +153,7 @@ def extract_modality_from_llm(llm_output):
             'flair': 'MRI-FLAIR-Brain',
             't1': 'MRI-T1-Gd-Brain',
             'contrast': 'MRI-T1-Gd-Brain',
-            'brain': 'MRI-FLAIR-Brain'  # default to FLAIR if just "brain" is mentioned
+            'brain': 'MRI-FLAIR-Brain'
         },
         'x-ray': {'chest': 'X-Ray-Chest'},
         'ultrasound': {'cardiac': 'Ultrasound-Cardiac', 'heart': 'Ultrasound-Cardiac'},
@@ -169,12 +166,9 @@ def extract_modality_from_llm(llm_output):
     
     for modality, subtypes in modality_keywords.items():
         if modality in llm_output:
-            # For modalities with subtypes, try to find the specific subtype
-            if subtypes:
-                for keyword, specific_modality in subtypes.items():
-                    if not keyword or keyword in llm_output:
-                        return specific_modality
-            # For modalities without subtypes, return the direct mapping
+            for keyword, specific_modality in subtypes.items():
+                if not keyword or keyword in llm_output:
+                    return specific_modality
             return next(iter(subtypes.values()))
     
     return None
@@ -202,7 +196,6 @@ def extract_clinical_findings(llm_output, modality):
 
 def on_mode_dropdown_change(selected_mode):
     if selected_mode in IMAGE_INFERENCE_MODES:
-        # Show modality dropdown and hide other inputs initially
         return [
             gr.Dropdown(visible=True, choices=list(BIOMEDPARSE_MODES.keys()), label="Modality"),
             gr.Dropdown(visible=True, label="Anatomical Site"),
@@ -210,7 +203,6 @@ def on_mode_dropdown_change(selected_mode):
             gr.Textbox(visible=False)
         ]
     else:
-        # Original behavior for other modes
         return [
             gr.Dropdown(visible=False),
             gr.Dropdown(visible=False),
@@ -223,7 +215,6 @@ def on_modality_change(modality):
         return gr.Dropdown(choices=BIOMEDPARSE_MODES[modality], visible=True)
     return gr.Dropdown(visible=False)
 
-
 def initialize_model():
     opt = load_opt_from_config_files(["configs/biomedparse_inference.yaml"])
     pretrained_pth = 'hf_hub:microsoft/BiomedParse'
@@ -238,7 +229,6 @@ def initialize_model():
 def initialize_llm():
     try:
         print("Starting LLM initialization...")
-        # Add quantization config
         quantization_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_compute_dtype=torch.float16
@@ -273,7 +263,6 @@ def update_example_prompts(modality):
         return f"Example prompts for {modality}:\n" + ", ".join(examples)
     return ""
 
-# Utility functions
 @spaces.GPU
 @torch.inference_mode()
 @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
@@ -285,46 +274,55 @@ def process_image(image_path, user_prompt, modality=None):
         image = read_rgb(image_path)
         pil_image = Image.fromarray(image)
         
-        # Step 1: Get LLM analysis
-        question = f"Analyze this medical image considering the following context: {user_prompt}. Include modality, anatomical structures, and any abnormalities."
+        question = (
+            f"Analyze this medical image considering the following context: {user_prompt}. "
+            "Include modality, anatomical structures, and any abnormalities."
+        )
         msgs = [{'role': 'user', 'content': [pil_image, question]}]
         
         llm_response = ""
-        for new_text in llm_model.chat(
-            image=pil_image,
-            msgs=msgs,
-            tokenizer=llm_tokenizer,
-            sampling=True,
-            temperature=0.95,
-            stream=True
-        ):
-            llm_response += new_text
-            
-        # Step 2: Extract modality from LLM output
+        if llm_model and llm_tokenizer:
+            for new_text in llm_model.chat(
+                image=pil_image,
+                msgs=msgs,
+                tokenizer=llm_tokenizer,
+                sampling=True,
+                temperature=0.95,
+                stream=True
+            ):
+                llm_response += new_text
+        else:
+            llm_response = "LLM not available. Please check LLM initialization logs."
+        
         detected_modality = extract_modality_from_llm(llm_response)
         if not detected_modality:
-            raise ValueError("Could not determine image modality from LLM output")
-            
-        # Step 3: Extract relevant clinical findings
+            # Fallback if modality wasn't detected
+            detected_modality = "X-Ray-Chest"
+        
         clinical_findings = extract_clinical_findings(llm_response, detected_modality)
+        if not clinical_findings:
+            # Fallback if no findings are detected
+            clinical_findings = [detected_modality.split("-")[-1]]
         
-        # Step 4: Generate masks for each finding
         results = []
         analysis_results = []
-        colors = [(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255)]  # Different colors for different findings
+        colors = [(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255)]
         
         for idx, finding in enumerate(clinical_findings):
-            pred_mask = interactive_infer_image(model, pil_image, [finding])[0]
-            p_value = check_mask_stats(image, pred_mask * 255, detected_modality, finding)
+            mask_list = interactive_infer_image(model, pil_image, [finding])
+            if not mask_list:
+                analysis_results.append(f"No mask found for '{finding}'.")
+                continue
+            
+            pred_mask = mask_list[0]
+            p_value = check_mask_stats(image, pred_mask*255, detected_modality, finding)
             analysis_results.append(f"P-value for '{finding}' ({detected_modality}): {p_value:.4f}")
             
-            # Create colored overlay
             overlay_image = image.copy()
             color = colors[idx % len(colors)]
             overlay_image[pred_mask > 0.5] = color
             results.append(overlay_image)
         
-        # Update LLM response with color references
         enhanced_response = llm_response + "\n\nSegmentation Results:\n"
         for idx, finding in enumerate(clinical_findings):
             color_name = ["red", "green", "blue", "yellow", "magenta"][idx % len(colors)]
@@ -345,7 +343,6 @@ def process_image(image_path, user_prompt, modality=None):
         print(f"Error details: {str(e)}", flush=True)
         return None, error_msg
 
-# Define Gradio interface
 with gr.Blocks() as demo:
     gr.HTML(MARKDOWN)    
     with gr.Row():        
@@ -376,17 +373,7 @@ with gr.Blocks() as demo:
                 show_label=True,
                 lines=10
             )
-            
-    # Examples section - Fixed version
-    # gr.Examples(
-    #     examples=IMAGE_PROCESSING_EXAMPLES,
-    #     inputs=[image_input, prompt_input],
-    #     outputs=[output_gallery, analysis_output, detected_modality],
-    #     fn=process_image,
-    #     cache_examples=True,
-    # )
     
-    # Connect the submit button to the process_image function
     submit_btn.click(
         fn=process_image,
         inputs=[image_input, prompt_input],