biomed-llama_multimodal

app.py

@@ -135,6 +135,78 @@ MODALITY_PROMPTS = {
    "OCT": ["edema"]
 }
 
+def extract_modality_and_prompts(llm_output):
+    """
+    Extract modality and relevant prompts from LLM output
+    Returns: (modality_type, list_of_prompts)
+    """
+    llm_output = llm_output.lower()
+    
+    # Dictionary mapping keywords to modalities
+    modality_indicators = {
+        'dermatoscop': 'Dermoscopy',
+        'skin lesion': 'Dermoscopy',
+        'oct': 'OCT',
+        'optical coherence': 'OCT',
+        'fundus': 'Fundus',
+        'retina': 'Fundus',
+        'endoscop': 'Endoscopy',
+        'colon': 'Endoscopy',
+        'patholog': 'Pathology',
+        'tissue': 'Pathology',
+        'histolog': 'Pathology',
+        'x-ray': 'X-Ray-Chest',
+        'xray': 'X-Ray-Chest',
+        'chest radiograph': 'X-Ray-Chest',
+        'mri': None,  # Will be refined below
+        'magnetic resonance': None,  # Will be refined below
+        'ct': None,  # Will be refined below
+        'computed tomography': None,  # Will be refined below
+        'ultrasound': 'Ultrasound-Cardiac',
+        'sonograph': 'Ultrasound-Cardiac'
+    }
+    
+    # First pass: Detect base modality
+    detected_modality = None
+    for keyword, modality in modality_indicators.items():
+        if keyword in llm_output:
+            detected_modality = modality
+            break
+    
+    # Second pass: Refine MRI and CT if detected
+    if detected_modality is None and ('mri' in llm_output or 'magnetic resonance' in llm_output):
+        if 'brain' in llm_output or 'flair' in llm_output:
+            detected_modality = 'MRI-FLAIR-Brain'
+        elif 'cardiac' in llm_output or 'heart' in llm_output:
+            detected_modality = 'MRI-Cardiac'
+        elif 'abdomen' in llm_output:
+            detected_modality = 'MRI-Abdomen'
+        elif 't1' in llm_output or 'contrast' in llm_output:
+            detected_modality = 'MRI-T1-Gd-Brain'
+        else:
+            detected_modality = 'MRI'
+    
+    if detected_modality is None and ('ct' in llm_output or 'computed tomography' in llm_output):
+        if 'chest' in llm_output or 'lung' in llm_output:
+            detected_modality = 'CT-Chest'
+        elif 'liver' in llm_output:
+            detected_modality = 'CT-Liver'
+        elif 'abdomen' in llm_output:
+            detected_modality = 'CT-Abdomen'
+        else:
+            detected_modality = 'CT'
+    
+    # If still no modality detected, return None
+    if not detected_modality:
+        return "", []
+    
+    # Get relevant prompts for the detected modality
+    if detected_modality in MODALITY_PROMPTS:
+        relevant_prompts = MODALITY_PROMPTS[detected_modality]
+    else:
+        relevant_prompts = []
+    
+    return detected_modality, relevant_prompts
 
 def extract_modality_from_llm(llm_output):
     """Extract modality from LLM output and map it to BIOMEDPARSE_MODES"""
@@ -278,7 +350,7 @@ def process_image(image_path, user_prompt, modality=None):
         #     f"Analyze this medical image considering the following context: {user_prompt}. "
         #     "Include modality, anatomical structures, and any abnormalities."
         # )
-        question = 'modality?'
+        question = 'What type of medical imaging modality is this? Please be specific.'
         msgs = [{'role': 'user', 'content': [pil_image, question]}]
         
         llm_response = ""
@@ -299,32 +371,29 @@ def process_image(image_path, user_prompt, modality=None):
         else:
             llm_response = "LLM not available. Please check LLM initialization logs."
         
-        detected_modality = extract_modality_from_llm(llm_response)
+        detected_modality, relevant_prompts = extract_modality_and_prompts(llm_response)
         if not detected_modality:
             detected_modality = "X-Ray-Chest"  # Fallback modality
-        
-        clinical_findings = extract_clinical_findings(llm_response, detected_modality)
-        if not clinical_findings:
-            clinical_findings = [detected_modality.split("-")[-1].lower()]
+            relevant_prompts = MODALITY_PROMPTS["X-Ray-Chest"]
         
         results = []
         analysis_results = []
         colors = [(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255)]
         
-        for idx, finding in enumerate(clinical_findings):
+        for idx, prompt in enumerate(relevant_prompts):
             try:
-                mask_list = interactive_infer_image(model, pil_image, [finding])
+                mask_list = interactive_infer_image(model, pil_image, [prompt])
                 if not mask_list or len(mask_list) == 0:
-                    analysis_results.append(f"No mask generated for '{finding}'")
+                    analysis_results.append(f"No mask generated for '{prompt}'")
                     continue
                 
                 pred_mask = mask_list[0]
-                if pred_mask is None or not pred_mask.any():
-                    analysis_results.append(f"Empty mask generated for '{finding}'")
-                    continue
+                # if pred_mask is None or not pred_mask.any():
+                #     analysis_results.append(f"Empty mask generated for '{finding}'")
+                #     continue
                 
-                p_value = check_mask_stats(image, pred_mask.astype(np.uint8) * 255, detected_modality, finding)
-                analysis_results.append(f"P-value for '{finding}' ({detected_modality}): {p_value:.4f}")
+                # p_value = check_mask_stats(image, pred_mask.astype(np.uint8) * 255, detected_modality, finding)
+                # analysis_results.append(f"P-value for '{finding}' ({detected_modality}): {p_value:.4f}")
                 
                 overlay_image = image.copy()
                 color = colors[idx % len(colors)]
@@ -333,26 +402,36 @@ def process_image(image_path, user_prompt, modality=None):
                     overlay_image[mask_indices] = color
                     results.append(overlay_image)
             except Exception as e:
-                print(f"Error processing finding {finding}: {str(e)}")
-                analysis_results.append(f"Failed to process '{finding}': {str(e)}")
+                print(f"Error processing finding {prompt}: {str(e)}")
+                analysis_results.append(f"Failed to process '{prompt}': {str(e)}")
         
         if not results:
             results = [image]  # Return original image if no overlays were created
             
-        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"
+        detailed_analysis = ""
+        # try:
+        analysis_prompt = f"Give the modality, organ, analysis, abnormalities (if any), treatment (if abnormalities are present) for this image. Focus more on the user question. which is: {user_prompt}"
+        msgs = [{'role': 'user', 'content': [pil_image, analysis_prompt]}]
         
-        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
+        # llm_response = ""
+        if llm_model and llm_tokenizer:
+            try:
+                for new_text in llm_model.chat(
+                    image=pil_image,
+                    msgs=msgs,
+                    tokenizer=llm_tokenizer,
+                    sampling=True,
+                    temperature=0.95,
+                    stream=True
+                ):
+                    detailed_analysis += new_text
+            except Exception as e:
+                print(f"LLM chat error: {str(e)}")
+                detailed_analysis = "LLM analysis failed. Proceeding with basic analysis."
+        else:
+            detailed_analysis = "LLM not available. Please check LLM initialization logs."
         
-        return results, combined_analysis, detected_modality
+        return results, detailed_analysis, detected_modality
 
     except Exception as e:
         error_msg = f"⚠️ An error occurred: {str(e)}"