bomedllamv2 integrated

app.py

@@ -269,7 +269,7 @@ def update_example_prompts(modality):
 def process_image(image_path, user_prompt, modality=None):
     try:
         if not image_path:
-            raise ValueError("Please upload an image")
+            return [], "Please upload an image", "No modality detected"
         
         image = read_rgb(image_path)
         pil_image = Image.fromarray(image)
@@ -282,47 +282,62 @@ def process_image(image_path, user_prompt, modality=None):
         
         llm_response = ""
         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
+            try:
+                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
+            except Exception as e:
+                print(f"LLM chat error: {str(e)}")
+                llm_response = "LLM analysis failed. Proceeding with basic analysis."
         else:
             llm_response = "LLM not available. Please check LLM initialization logs."
         
         detected_modality = extract_modality_from_llm(llm_response)
         if not detected_modality:
-            # Fallback if modality wasn't detected
-            detected_modality = "X-Ray-Chest"
+            detected_modality = "X-Ray-Chest"  # Fallback modality
         
         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]]
+            clinical_findings = [detected_modality.split("-")[-1].lower()]
         
         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):
-            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}")
-            
-            overlay_image = image.copy()
-            color = colors[idx % len(colors)]
-            overlay_image[pred_mask > 0.5] = color
-            results.append(overlay_image)
+            try:
+                mask_list = interactive_infer_image(model, pil_image, [finding])
+                if not mask_list or len(mask_list) == 0:
+                    analysis_results.append(f"No mask generated for '{finding}'")
+                    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
+                
+                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)]
+                mask_indices = pred_mask > 0.5
+                if mask_indices.any():
+                    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)}")
         
+        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)]
@@ -341,7 +356,7 @@ def process_image(image_path, user_prompt, modality=None):
     except Exception as e:
         error_msg = f"⚠️ An error occurred: {str(e)}"
         print(f"Error details: {str(e)}", flush=True)
-        return None, error_msg
+        return [image] if 'image' in locals() else [], error_msg, "Error detecting modality"
 
 with gr.Blocks() as demo:
     gr.HTML(MARKDOWN)