Update app.py

app.py

@@ -66,7 +66,7 @@ def initialize_app():
     return True
 
 
-def process_image_task(image_path: str, task_text: str, agent: CellposeAgent) -> str:
+def process_image_task(image_path: str, task_text: str, agent: CellposeAgent) -> tuple[str, str | None]:
     """
     Process a user task with the CellposeAgent.
     
@@ -76,20 +76,55 @@ def process_image_task(image_path: str, task_text: str, agent: CellposeAgent) ->
         agent: Initialized CellposeAgent instance
     
     Returns:
-        str: Agent's response
+        tuple: (agent's text response, path to segmented image or None)
     """
     if not image_path:
-        return "⚠️ Please upload an image first."
+        return "⚠️ Please upload an image first.", None
     
-    if not task_text:
-        task_text = f"What parameters would work best for my image {image_path}?"
+    # Always include the full image path in the task
+    if not task_text or task_text.strip() == "":
+        task = f"What parameters would work best for my image {image_path}?"
+    else:
+        if image_path not in task_text:
+            task = f"For image {image_path}: {task_text}"
+        else:
+            task = task_text
     
     try:
-        result = agent.run(task_text)
+        result = agent.run(task)
         get_client().flush()
-        return result
+        
+        # Extract output image path from the agent's response if segmentation was run
+        output_image_path = None
+        try:
+            # Check if the response contains a segmentation output path
+            if "output_path" in result or "_cellpose_sam_overlay.png" in result:
+                # Parse the response to find the output path
+                import json
+                import re
+                
+                # Try to find JSON in the response
+                json_match = re.search(r'\{[^{}]*"output_path"[^{}]*\}', result)
+                if json_match:
+                    result_data = json.loads(json_match.group())
+                    output_image_path = result_data.get("output_path")
+                else:
+                    # Try to find the path in the text
+                    path_match = re.search(r'([^\s]+_cellpose_sam_overlay\.png)', result)
+                    if path_match:
+                        output_image_path = path_match.group(1)
+                
+                # Verify the file exists
+                if output_image_path and not Path(output_image_path).exists():
+                    output_image_path = None
+                    
+        except Exception as e:
+            print(f"Could not extract output image path: {e}")
+        
+        return result, output_image_path
+        
     except Exception as e:
-        return f"❌ Error processing task: {str(e)}"
+        return f"❌ Error processing task: {str(e)}", None
 
 
 def create_gradio_interface():
@@ -124,7 +159,7 @@ def create_gradio_interface():
                 # Task input
                 task_input = gr.Textbox(
                     label="Task / Question",
-                    placeholder="e.g., 'What parameters would work best for this image?' or leave empty for default",
+                    placeholder="e.g., 'What parameters would work best?' or 'Run segmentation' (image path added automatically)",
                     lines=3
                 )
                 
@@ -145,19 +180,26 @@ def create_gradio_interface():
                 )
             
             with gr.Column(scale=1):
-                # Output
+                # Text output
                 output = gr.Textbox(
                     label="Agent Response",
-                    lines=20,
-                    max_lines=30,
+                    lines=12,
+                    max_lines=20,
                     show_copy_button=True
                 )
+                
+                # Image output for segmentation results
+                output_image = gr.Image(
+                    label="Segmentation Result",
+                    type="filepath",
+                    height=400
+                )
         
         # Event handler
         submit_btn.click(
             fn=lambda img, task: process_image_task(img, task, agent),
             inputs=[image_input, task_input],
-            outputs=output
+            outputs=[output, output_image]
         )
         
         gr.Markdown(
@@ -178,6 +220,7 @@ def create_gradio_interface():
             3. Visually analyzes your image to validate recommendations
             4. Runs segmentation and checks quality
             5. Refines parameters if needed (up to 2 iterations)
+            6. **Displays the segmented overlay image with colored cell masks**
             """
         )