""" 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()