Spaces:
Runtime error
Runtime error
| """ | |
| Gradio web interface for CellposeAgent | |
| """ | |
| import os | |
| import json | |
| import gradio as gr | |
| from pathlib import Path | |
| from langfuse import get_client | |
| from openinference.instrumentation.smolagents import SmolagentsInstrumentor | |
| from smolagents.agents import ActionStep | |
| from config import settings | |
| from agents.agent import CellposeAgent | |
| from stores import neo4j_store | |
| from utils.prechecks import check_hf_persistent_storage | |
| def setup_observability(): | |
| """Initializes Langfuse and Smolagents instrumentation.""" | |
| try: | |
| if os.getenv("LANGFUSE_PUBLIC_KEY") and os.getenv("LANGFUSE_SECRET_KEY"): | |
| get_client() | |
| SmolagentsInstrumentor().instrument() | |
| print("✓ Observability and instrumentation initialized.") | |
| else: | |
| print("⚠️ Langfuse keys not found. Observability disabled.") | |
| except Exception as e: | |
| print(f"⚠️ Could not initialize observability: {e}") | |
| def initialize_app(): | |
| """Initialize the application and verify prerequisites.""" | |
| print("\n--- Initializing Cellpose Agent Application ---") | |
| # Setup observability | |
| setup_observability() | |
| # Configure LlamaIndex | |
| settings.configure_llama_index() | |
| # check for cellpose-db | |
| check_hf_persistent_storage( | |
| repo_id = "hmgill/Cellpose-DB", | |
| target = "cellpose_db", | |
| file_or_folder="folder" | |
| ) | |
| # check for cellpose sam | |
| check_hf_persistent_storage( | |
| repo_id = "hmgill/Cellpose-SAM-Checkpoint", | |
| target = "sam_vit_h_4b8939.pth", | |
| file_or_folder="file" | |
| ) | |
| # Verify knowledge graph is ready | |
| try: | |
| node_count, _ = neo4j_store.check_graph_status() | |
| if node_count == 0: | |
| print("\n❌ WARNING: The knowledge graph is empty.") | |
| print("Please run the setup script to build the knowledge graph:") | |
| print("\n python setup_kg.py\n") | |
| return False | |
| print(f"✓ Knowledge graph is ready with {node_count} nodes.") | |
| except Exception as e: | |
| print(f"❌ ERROR: Could not connect to Neo4j: {e}") | |
| print("Please ensure Neo4j is running and accessible.") | |
| return False | |
| return True | |
| def process_image_task(image_path: str, task_text: str, agent: CellposeAgent) -> tuple[str, str | None]: | |
| """ | |
| Process a user task with the CellposeAgent. | |
| The image_path is now passed SEPARATELY to agent.run() so it gets stored | |
| in global context. This prevents the LLM from needing to reproduce the | |
| exact path string (which it often corrupts). | |
| Args: | |
| image_path: Path to the uploaded image file | |
| task_text: User's text prompt/question | |
| agent: Initialized CellposeAgent instance | |
| Returns: | |
| tuple: (agent's text response, path to segmented image or None) | |
| """ | |
| if not image_path: | |
| return "⚠️ Please upload an image first.", None | |
| # Build task WITHOUT embedding the full path | |
| # The agent instructions tell it to use image_path="" and the tool will resolve from context | |
| if not task_text or task_text.strip() == "": | |
| task = "Analyze the uploaded image and recommend optimal segmentation parameters. Then run segmentation." | |
| else: | |
| # Don't include the path in the task - it will be available via context | |
| task = task_text | |
| try: | |
| # Pass image_path SEPARATELY - this stores it in global context | |
| # Tools will retrieve the correct path from context instead of relying on LLM | |
| result = agent.run(task=task, image_path=image_path) | |
| get_client().flush() | |
| # Extract output image path from agent's memory (more reliable than parsing text) | |
| output_image_path = None | |
| try: | |
| # Search through agent's memory steps in reverse order (most recent first) | |
| for step in reversed(agent.agent.memory.steps): | |
| if isinstance(step, ActionStep) and step.observations: | |
| try: | |
| obs_data = json.loads(step.observations) | |
| # Check if this step contains segmentation output | |
| if obs_data.get("status") == "success" and "output_path" in obs_data: | |
| candidate_path = obs_data["output_path"] | |
| # Verify the file exists | |
| if Path(candidate_path).exists(): | |
| output_image_path = candidate_path | |
| print(f"✓ Found segmentation output: {output_image_path}") | |
| break | |
| except (json.JSONDecodeError, Exception) as parse_error: | |
| continue | |
| if output_image_path is None: | |
| print("ℹ️ No segmentation output found in agent memory") | |
| except Exception as e: | |
| print(f"Warning: Could not extract output image from agent memory: {e}") | |
| return result, output_image_path | |
| except Exception as e: | |
| return f"❌ Error processing task: {str(e)}", None | |
| def create_gradio_interface(): | |
| """Creates and configures the Gradio interface.""" | |
| # Initialize the agent once at startup | |
| if not initialize_app(): | |
| raise RuntimeError("Failed to initialize application. Please check logs.") | |
| agent = CellposeAgent() | |
| print("✓ CellposeAgent initialized and ready.") | |
| with gr.Blocks(title="Cellpose-SAM Agent", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown( | |
| """ | |
| # 🔬 Cellpose-SAM Segmentation Agent | |
| Upload a microscopy image and ask the AI agent to recommend optimal segmentation parameters, | |
| run segmentation, or answer questions about the cellpose-sam pipeline. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| # Image upload | |
| image_input = gr.Image( | |
| label="Upload Microscopy Image", | |
| type="filepath", | |
| height=300 | |
| ) | |
| # Task input | |
| task_input = gr.Textbox( | |
| label="Task / Question", | |
| placeholder="e.g., 'What parameters would work best?' or 'Run segmentation' (image path handled automatically)", | |
| lines=3 | |
| ) | |
| # Submit button | |
| submit_btn = gr.Button("Run Agent", variant="primary", size="lg") | |
| # Example tasks | |
| gr.Markdown("### 💡 Example Tasks") | |
| gr.Examples( | |
| examples=[ | |
| ["What parameters would work best for this image?"], | |
| ["Analyze this image and run segmentation with optimal parameters."], | |
| ["What is the flow_threshold parameter and how does it affect segmentation?"], | |
| ["Run segmentation with diameter=30, flow_threshold=0.5, cellprob_threshold=0, min_size=20"], | |
| ], | |
| inputs=task_input, | |
| label="Click to use:" | |
| ) | |
| with gr.Column(scale=1): | |
| # Text output | |
| output = gr.Textbox( | |
| label="Agent Response", | |
| 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, output_image] | |
| ) | |
| gr.Markdown( | |
| """ | |
| --- | |
| ### 📚 What can this agent do? | |
| - **Parameter Recommendation**: Analyzes your image and suggests optimal segmentation parameters | |
| - **Automated Segmentation**: Runs the full cellpose-sam pipeline with parameter refinement | |
| - **Visual Analysis**: Uses vision-language models to assess segmentation quality | |
| - **Documentation Search**: Answers questions about parameters using RAG and knowledge graphs | |
| - **Iterative Refinement**: Automatically adjusts parameters based on visual feedback | |
| ### 🔍 How it works | |
| 1. Upload your microscopy image | |
| 2. The agent finds similar images and recommends parameters | |
| 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** | |
| ### ⚙️ Technical Note | |
| Image paths are handled automatically by the system - the agent doesn't need to | |
| reproduce exact file paths, preventing common path corruption issues with LLMs. | |
| """ | |
| ) | |
| return demo | |
| def main(): | |
| """Launch the Gradio application.""" | |
| demo = create_gradio_interface() | |
| demo.launch(server_port=7860) | |
| if __name__ == "__main__": | |
| main() |