integrated biomedllama

app.py

@@ -16,31 +16,15 @@ from tqdm import tqdm
 import sys
 from pathlib import Path
 from huggingface_hub import login
-# from dotenv import load_dotenv
 from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
 
-# For Hugging Face Spaces, secrets are automatically loaded as environment variables
+
 token = os.getenv("HF_TOKEN")
 if token:
     login(token=token)
-# Clear Hugging Face cache
-# cache_dirs = [
-#     "/home/user/.cache/huggingface/",
-#     "/home/user/.cache/torch/",
-#     "/home/user/.cache/pip/"
-# ]
-
-# for cache_dir in cache_dirs:
-#     if os.path.exists(cache_dir):
-#         print(f"Clearing cache: {cache_dir}")
-#         shutil.rmtree(cache_dir, ignore_errors=True)
-# Add the current directory to Python path
 current_dir = Path(__file__).parent
 sys.path.append(str(current_dir))
-# sys.path.append("./BiomedParse/")
-# BIOMEDPARSE_PATH = Path(__file__).parent / "BiomedParse"
-# sys.path.append(str(BIOMEDPARSE_PATH))
-# sys.path.append(str(BIOMEDPARSE_PATH / "BiomedParse"))  # Add the inner BiomedParse directory
+
 from modeling.BaseModel import BaseModel
 from modeling import build_model
 from utilities.arguments import load_opt_from_config_files
@@ -51,7 +35,7 @@ from inference_utils.processing_utils import read_rgb
 
 import spaces
 
-# breakpoint()
+
 MARKDOWN = """
 <div align="center" style="padding: 20px 0;">
     <h1 style="font-size: 3em; margin: 0;">
@@ -154,6 +138,68 @@ MODALITY_PROMPTS = {
    "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',
+            'chest': 'CT-Chest',
+            'liver': 'CT-Liver'
+        },
+        'mri': {
+            'abdomen': 'MRI-Abdomen',
+            'cardiac': 'MRI-Cardiac',
+            'heart': 'MRI-Cardiac',
+            '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
+        },
+        'x-ray': {'chest': 'X-Ray-Chest'},
+        'ultrasound': {'cardiac': 'Ultrasound-Cardiac', 'heart': 'Ultrasound-Cardiac'},
+        'endoscopy': {'': 'Endoscopy'},
+        'fundus': {'': 'Fundus'},
+        'dermoscopy': {'': 'Dermoscopy'},
+        'oct': {'': 'OCT'},
+        'pathology': {'': 'Pathology'}
+    }
+    
+    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
+            return next(iter(subtypes.values()))
+    
+    return None
+
+def extract_clinical_findings(llm_output, modality):
+    """Extract relevant clinical findings that match available anatomical sites in BIOMEDPARSE_MODES"""
+    available_sites = BIOMEDPARSE_MODES.get(modality, [])
+    findings = []
+    
+    # Convert sites to lowercase for case-insensitive matching
+    sites_lower = {site.lower(): site for site in available_sites}
+    
+    # Look for each available site in the LLM output
+    for site_lower, original_site in sites_lower.items():
+        if site_lower in llm_output.lower():
+            findings.append(original_site)
+    
+    # Add additional findings from MODALITY_PROMPTS if available
+    if modality in MODALITY_PROMPTS:
+        for prompt in MODALITY_PROMPTS[modality]:
+            if prompt.lower() in llm_output.lower() and prompt not in findings:
+                findings.append(prompt)
+    
+    return findings
+
 def on_mode_dropdown_change(selected_mode):
     if selected_mode in IMAGE_INFERENCE_MODES:
         # Show modality dropdown and hide other inputs initially
@@ -231,72 +277,68 @@ def update_example_prompts(modality):
 @spaces.GPU
 @torch.inference_mode()
 @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
-def process_image(image_path, text_prompts, modality):
+def process_image(image_path, user_prompt, modality=None):
     try:
-        # Input validation
         if not image_path:
             raise ValueError("Please upload an image")
-        if not text_prompts or text_prompts.strip() == "":
-            raise ValueError("Please enter prompts for analysis")
-        if not modality:
-            raise ValueError("Please select a modality")
-
-        # Original BiomedParse processing
+        
         image = read_rgb(image_path)
-        text_prompts = [prompt.strip() for prompt in text_prompts.split(',')]
-        pred_masks = interactive_infer_image(model, Image.fromarray(image), text_prompts)
+        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."
+        msgs = [{'role': 'user', 'content': [pil_image, question]}]
         
-        # Prepare outputs
+        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
+        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
+        clinical_findings = extract_clinical_findings(llm_response, detected_modality)
+        
+        # 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
         
-        # Process with BiomedParse
-        for i, prompt in enumerate(text_prompts):
-            p_value = check_mask_stats(image, pred_masks[i] * 255, modality, prompt)
-            analysis_results.append(f"P-value for '{prompt}' ({modality}): {p_value:.4f}")
+        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)
+            analysis_results.append(f"P-value for '{finding}' ({detected_modality}): {p_value:.4f}")
             
+            # Create colored overlay
             overlay_image = image.copy()
-            overlay_image[pred_masks[i] > 0.5] = [255, 0, 0]
+            color = colors[idx % len(colors)]
+            overlay_image[pred_mask > 0.5] = color
             results.append(overlay_image)
         
-        # Process with LLM only if available
-        if llm_model is not None and llm_tokenizer is not None:
-            print("LLM model and tokenizer are available")
-            try:
-                pil_image = Image.fromarray(image)
-                question = 'Give the modality, organ, analysis, abnormalities (if any), treatment (if abnormalities are present)?'
-                msgs = [{'role': 'user', 'content': [pil_image, question]}]
-                
-                print("Starting LLM inference...")
-                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
-                print(f"LLM generated response: {llm_response}")
-                
-                # Make the combined analysis more visible
-                combined_analysis = "\n\n" + "="*50 + "\n"
-                combined_analysis += "BiomedParse Analysis:\n"
-                combined_analysis += "\n".join(analysis_results)
-                combined_analysis += "\n\n" + "="*50 + "\n"
-                combined_analysis += "LLM Analysis:\n"
-                combined_analysis += llm_response
-                combined_analysis += "\n" + "="*50
-                
-            except Exception as e:
-                print(f"LLM analysis failed with error: {str(e)}")
-                combined_analysis = "\n".join(analysis_results)
-        else:
-            print("LLM model or tokenizer is not available")
-            combined_analysis = "\n".join(analysis_results)
+        # 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)]
+            enhanced_response += f"- {finding} (shown in {color_name})\n"
         
-        return results, combined_analysis
+        combined_analysis = "\n\n" + "="*50 + "\n"
+        combined_analysis += "BiomedParse Analysis:\n"
+        combined_analysis += "\n".join(analysis_results)
+        combined_analysis += "\n\n" + "="*50 + "\n"
+        combined_analysis += "Enhanced LLM Analysis:\n"
+        combined_analysis += enhanced_response
+        combined_analysis += "\n" + "="*50
+        
+        return results, combined_analysis, detected_modality
 
     except Exception as e:
         error_msg = f"⚠️ An error occurred: {str(e)}"
@@ -309,33 +351,45 @@ with gr.Blocks() as demo:
     with gr.Row():        
         with gr.Column():
             image_input = gr.Image(type="filepath", label="Input Image")
-            prompts_input = gr.Textbox(
-                lines=2, 
-                placeholder="Enter prompts separated by commas...", 
-                label="Prompts"
+            prompt_input = gr.Textbox(
+                lines=4, 
+                placeholder="Ask any question about the medical image...", 
+                label="Question/Prompt"
             )
-            modality_dropdown = gr.Dropdown(
-                choices=list(BIOMEDPARSE_MODES.keys()),
-                value=list(BIOMEDPARSE_MODES.keys())[0],
-                label="Modality"
+            detected_modality = gr.Textbox(
+                label="Detected Modality", 
+                interactive=False,
+                visible=True
             )
-            submit_btn = gr.Button("Submit")
+            submit_btn = gr.Button("Analyze")
+            
         with gr.Column():
-            output_gallery = gr.Gallery(label="Findings")
-            pvalue_output = gr.Textbox(
-                label="Results", 
+            output_gallery = gr.Gallery(
+                label="Segmentation Results", 
+                show_label=True,
+                columns=[2], 
+                height="auto"
+            )
+            analysis_output = gr.Textbox(
+                label="Analysis", 
                 interactive=False,
-                show_label=True
+                show_label=True,
+                lines=10
             )
-            with gr.Accordion("Example Prompts by Modality", open=False):
-                for modality, prompts in MODALITY_PROMPTS.items():
-                    prompt_str = ", ".join(prompts)
-                    gr.Markdown(f"**{modality}**: {prompt_str}")
-    # Add error handling for the submit button
+            
+    # Examples section
+    gr.Examples(
+        examples=IMAGE_PROCESSING_EXAMPLES,
+        inputs=[image_input, prompt_input],
+        outputs=[output_gallery, analysis_output, detected_modality],
+        cache_examples=True,
+    )
+    
+    # Connect the submit button to the process_image function
     submit_btn.click(
         fn=process_image,
-        inputs=[image_input, prompts_input, modality_dropdown],
-        outputs=[output_gallery, pvalue_output],
+        inputs=[image_input, prompt_input],
+        outputs=[output_gallery, analysis_output, detected_modality],
         api_name="process"
     )