integrated biomedllama

app.py

@@ -285,33 +285,14 @@ def process_image(image_path, user_prompt, modality=None):
         image = read_rgb(image_path)
         pil_image = Image.fromarray(image)
         
-        # Stage 1: Get modality identification only
-        modality_question = "What is the medical imaging modality of this image? Please respond with only the modality type (CT, MRI, X-Ray, Pathology, Ultrasound, Endoscopy, Fundus, Dermoscopy, or OCT) and if applicable, the specific body part or sequence (e.g., CT-Chest, MRI-FLAIR-Brain)."
-        modality_msgs = [{'role': 'user', 'content': [pil_image, modality_question]}]
-        
-        modality_response = ""
-        for new_text in llm_model.chat(
-            image=pil_image,
-            msgs=modality_msgs,
-            tokenizer=llm_tokenizer,
-            sampling=True,
-            temperature=0.7,  # Lower temperature for more focused response
-            stream=True
-        ):
-            modality_response += new_text
-            
-        detected_modality = extract_modality_from_llm(modality_response)
-        if not detected_modality:
-            raise ValueError("Could not determine image modality")
-            
-        # Stage 2: Get detailed analysis with known modality
-        analysis_question = f"This is a {detected_modality} image. {user_prompt} Please analyze the anatomical structures and any abnormalities visible in this image."
-        analysis_msgs = [{'role': 'user', 'content': [pil_image, analysis_question]}]
+        # 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]}]
         
         llm_response = ""
         for new_text in llm_model.chat(
             image=pil_image,
-            msgs=analysis_msgs,
+            msgs=msgs,
             tokenizer=llm_tokenizer,
             sampling=True,
             temperature=0.95,
@@ -319,39 +300,50 @@ def process_image(image_path, user_prompt, modality=None):
         ):
             llm_response += new_text
             
-        # Extract findings using known modality
+        # 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)
         
-        # Rest of the processing remains the same
+        # 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)]
+        colors = [(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255)]  # Different colors for different findings
         
         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()
             color = colors[idx % len(colors)]
             overlay_image[pred_mask > 0.5] = color
             results.append(overlay_image)
         
-        enhanced_response = f"Modality: {detected_modality}\n\n"
-        enhanced_response += llm_response + "\n\nSegmentation Results:\n"
+        # 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"
         
-        combined_analysis = "\n".join(analysis_results)
+        # 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)}"
         print(f"Error details: {str(e)}", flush=True)
-        return None, error_msg, None
+        return None, error_msg
 
 # Define Gradio interface
 with gr.Blocks() as demo: