Update app.py

app.py

@@ -19,37 +19,52 @@ from google.adk.runners import InMemoryRunner
 from google.genai import types
 
 # Project Imports
-from cellemetry import root_agent
-from cellemetry.config import AnalysisDeps
-from transformers import Sam3Processor, Sam3Model
+# Wrap imports to prevent immediate crash if dependencies are missing
+try:
+    from cellemetry import root_agent
+    from cellemetry.config import AnalysisDeps
+    from transformers import Sam3Processor, Sam3Model
+except ImportError as e:
+    print(f"⚠️ Import Error (Non-fatal for UI startup): {e}")
+    Sam3Model = None
+    Sam3Processor = None
+    root_agent = None
+    AnalysisDeps = None
 
 # --- Global State ---
 MODEL_CACHE = {
     "model": None,
     "processor": None,
-    "device": "cpu"
+    "device": "cpu",
+    "loaded": False
 }
 
-# Store the active runner globally so we can access it in Q&A turns
+# Store the active runner globally
 ACTIVE_RUNNER = None
 
 
 def load_models():
-    """Initialize SAM3 model."""
-    if MODEL_CACHE["model"] is not None:
+    """Initialize SAM3 model. Now called AFTER app startup."""
+    if MODEL_CACHE["loaded"]:
         return
+        
     print("--- Loading SAM3 Model ---")
     device = "cuda" if torch.cuda.is_available() else "cpu"
     MODEL_CACHE["device"] = device
     
     try:
+        # Check if imports succeeded
+        if Sam3Model is None:
+            raise ImportError("Sam3Model not found. Please check requirements.")
+
         MODEL_CACHE["model"] = Sam3Model.from_pretrained("facebook/sam3").to(device)
         MODEL_CACHE["processor"] = Sam3Processor.from_pretrained("facebook/sam3")
+        MODEL_CACHE["loaded"] = True
         print(f"✅ SAM3 loaded on {device}")
+        return f"✅ SAM3 loaded on {device}"
     except Exception as e:
         print(f"⚠️ SAM3 load failed: {e}")
-
-load_models()
+        return f"⚠️ Model load failed: {e}"
 
 # --- Helpers ---
 def load_excel_data(logs_text):
@@ -65,14 +80,11 @@ def load_excel_data(logs_text):
     try:
         xls = pd.ExcelFile(report_file, engine='openpyxl')
         
-        # Helper to transpose and fix index so row labels are visible
         def process_sheet(sheet_name):
             if sheet_name in xls.sheet_names:
                 df = pd.read_excel(xls, sheet_name)
                 if not df.empty and len(df.columns) > 0:
-                    # Set first column as index, transpose, then reset index to make it a column again
                     df = df.set_index(df.columns[0]).T.reset_index()
-                    # Rename the new first column (formerly the index) to 'Metric'
                     df.rename(columns={df.columns[0]: "Metric"}, inplace=True)
                     return df
             return placeholder
@@ -87,7 +99,6 @@ def load_excel_data(logs_text):
         return report_file, placeholder, placeholder, placeholder
 
 def get_available_layers():
-    """Scans /tmp for .npz files and returns a list of available layer names."""
     files = glob.glob("/tmp/data_*.npz")
     layers = []
     for f in files:
@@ -96,9 +107,8 @@ def get_available_layers():
     return sorted(layers)
 
 def update_opacity_sliders(layers):
-    """Returns updated slider configurations based on available layers."""
     updates = []
-    for i in range(4):  # We have 4 sliders
+    for i in range(4):
         if i < len(layers):
             layer_name = layers[i].replace("_", " ").title()
             updates.append(gr.update(visible=True, label=f"{layer_name} Opacity", value=0.6))
@@ -106,22 +116,12 @@ def update_opacity_sliders(layers):
             updates.append(gr.update(visible=False))
     return updates
 
-def collect_layer_opacities(layers, op1, op2, op3, op4):
-    """Collects opacity values into a dictionary."""
-    opacities = {}
-    opacity_values = [op1, op2, op3, op4]
-    for i, layer in enumerate(layers[:4]):  # Only use first 4 layers
-        opacities[layer] = opacity_values[i]
-    return opacities
-
 def generate_overlay(image_path_str, selected_layers, layer_opacities=None):
-    """Regenerates the overlay image with adjustable opacity for each layer."""
     if not image_path_str:
         return None
     
     base_image = Image.open(image_path_str).convert("RGBA")
 
-    # Default colors for different layers (can expand as needed)
     color_map = {
         "green_cell": (0, 255, 0),
         "blue_nucleus": (0, 0, 255),
@@ -129,7 +129,6 @@ def generate_overlay(image_path_str, selected_layers, layer_opacities=None):
         "nucleus": (0, 0, 255),
     }
     
-    # Create overlay layer
     overlay = Image.new('RGBA', base_image.size, (0, 0, 0, 0))
     
     for layer_name in selected_layers:
@@ -145,58 +144,59 @@ def generate_overlay(image_path_str, selected_layers, layer_opacities=None):
                     else:
                         combined_mask = masks
                     
-                    # Resize if mask dimensions differ from image
                     h_mask, w_mask = combined_mask.shape
                     if overlay.size != (w_mask, h_mask):
                         overlay = overlay.resize((w_mask, h_mask), Image.Resampling.LANCZOS)
                         base_image = base_image.resize((w_mask, h_mask), Image.Resampling.LANCZOS)
 
                     combined_mask = combined_mask.astype(bool)
+                    color = color_map.get(layer_name.lower(), (255, 255, 0))
                     
-                    # Get color for this layer
-                    color = color_map.get(layer_name.lower(), (255, 255, 0))  # Default to yellow
-                    
-                    # Get opacity (default 0.5)
                     opacity = 0.5
                     if layer_opacities and layer_name in layer_opacities:
                         opacity = layer_opacities[layer_name]
                     
-                    # Create colored mask with opacity
                     mask_overlay = np.zeros((*combined_mask.shape, 4), dtype=np.uint8)
                     mask_overlay[combined_mask] = (*color, int(255 * opacity))
                     
-                    # Composite onto overlay
                     mask_image = Image.fromarray(mask_overlay, 'RGBA')
                     overlay = Image.alpha_composite(overlay, mask_image)
-            
             except Exception as e:
                 print(f"Error loading layer {layer_name}: {e}")
 
-    # Composite overlay onto base image
     result = Image.alpha_composite(base_image, overlay)
     return result.convert("RGB")
 
 # --- Core Logic ---
 async def run_analysis(image_path_str, user_prompt, session_id_state):
-    """
-    Runs the initial analysis using the Agent.
-    Updates the global ACTIVE_RUNNER and returns a session ID.
-    """
+    """Runs analysis. Triggers model load if not yet ready."""
     waiting_df = pd.DataFrame({"Status": ["Waiting..."]})
-    empty_slider_updates = [gr.update()] * 4  # Placeholder for 4 sliders
+    empty_slider_updates = [gr.update()] * 4
     
+    # Lazy Load: Ensure model is loaded before inference
+    if not MODEL_CACHE["loaded"]:
+        yield [], None, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
+        # We can't yield a log easily here without breaking the tuple structure, so we just wait
+        load_models() 
+
     if not image_path_str:
-        # Return empty state
         yield [], None, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
         return
 
-    # Cleanup previous run files
+    # Cleanup
     for f in glob.glob("/tmp/out_*.png") + glob.glob("/tmp/data_*.npz") + glob.glob("/tmp/*.xlsx"):
         try: os.remove(f)
         except: pass
 
-    # Setup Dependencies
+    # Setup
     image_path = Path(image_path_str)
+    
+    # Check if model loaded successfully
+    if MODEL_CACHE["model"] is None:
+        error_msg = "❌ Model failed to load. Please check logs."
+        yield [{"role": "assistant", "content": error_msg}], None, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
+        return
+
     deps = AnalysisDeps(
         sam_model=MODEL_CACHE["model"],
         sam_processor=MODEL_CACHE["processor"],
@@ -205,11 +205,9 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
         pixel_size_microns=None 
     )
     
-    # Initialize Runner
     global ACTIVE_RUNNER
     ACTIVE_RUNNER = InMemoryRunner(agent=root_agent, app_name="cellemetry_demo")
     
-    # Create Session
     session = await ACTIVE_RUNNER.session_service.create_session(
         app_name="cellemetry_demo",
         user_id="demo_user",
@@ -217,8 +215,6 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
     )
     
     session_id = session.id
-
-    # Prepare Input
     image_bytes = image_path.read_bytes()
     content = types.Content(
         role="user",
@@ -230,15 +226,10 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
 
     logs = [f"🔄 **Starting analysis** on {MODEL_CACHE['device']}..."]
     
-    # Helper to format output for the chatbot (UPDATED for Gradio 5.0 Messages format)
     def yield_status(log_list):
         full_log = "\n\n".join(log_list)
-        # Use Markdown to render the image in the chat
         user_msg = f"![](file={image_path_str})\n\n{user_prompt}"
-        return [
-            {"role": "user", "content": user_msg},
-            {"role": "assistant", "content": full_log}
-        ]
+        return [{"role": "user", "content": user_msg}, {"role": "assistant", "content": full_log}]
 
     yield yield_status(logs), session_id, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
 
@@ -264,7 +255,6 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
                                 logs[-1] = f"✅ **{author}**: {part.text}"
                             else:
                                 logs.append(f"✅ **{author}**: {part.text}")
-                        
                         yield yield_status(logs), session_id, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
 
     except Exception as e:
@@ -272,99 +262,203 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
         yield yield_status(logs), session_id, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
         return
 
-    # Finalize
     logs.append("\n✅ **Analysis Complete!** Loading results...")
     yield yield_status(logs), session_id, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
     
     await asyncio.sleep(0.5) 
     
-    # Load Data
     full_log_text = "\n".join(logs)
     report_file, df_m, df_s, df_r = load_excel_data(full_log_text)
     layers = get_available_layers()
     initial_overlay = generate_overlay(image_path_str, layers)
 
-    # Add completion summary
     completion_msg = f"\n\n---\n\n✨ **Analysis finished!** Found {len(layers)} layer(s). Results are now available in the Segmentation and Quantitative Results tabs."
     full_log_text += completion_msg
     
-    # Final Yield with all data
     final_user_msg = f"![](file={image_path_str})\n\n{user_prompt}"
-    final_history = [
-        {"role": "user", "content": final_user_msg},
-        {"role": "assistant", "content": full_log_text}
-    ]
+    final_history = [{"role": "user", "content": final_user_msg}, {"role": "assistant", "content": full_log_text}]
     slider_updates = update_opacity_sliders(layers)
+    
     yield final_history, session_id, initial_overlay, gr.CheckboxGroup(choices=layers, value=layers), report_file, df_m, df_s, df_r, *slider_updates
 
 async def unified_chat_handler(message, history, session_id, current_img_path):
-    """
-    Unified handler for both initial analysis and follow-up questions.
-    message: dict with 'text' and optionally 'files' keys (Gradio MultimodalTextbox format)
-    """
-    waiting_df = pd.DataFrame({"Status": ["Waiting..."]})
-    empty_slider_updates = [gr.update()] * 4
-    
-    # Ensure history is a list
     if history is None:
         history = []
 
-    # Extract text and files from message
     user_text = message.get("text", "").strip() if isinstance(message, dict) else str(message).strip()
     files = message.get("files", []) if isinstance(message, dict) else []
     
-    # Determine if we have an image
     image_path = None
     if files:
         image_path = files[0] if isinstance(files[0], str) else files[0].get("path")
     elif current_img_path:
         image_path = current_img_path
     
-    # Case 1: Initial analysis (has image, no session or new image)
+    waiting_df = pd.DataFrame({"Status": ["Waiting..."]})
+    empty_slider_updates = [gr.update()] * 4
+
+    # Case 1: Initial analysis
     if image_path and (not session_id or files):
-        # This is a new analysis request
-        if not user_text:
-            user_text = "Analyze this microscopy image."
+        if not user_text: user_text = "Analyze this microscopy image."
         
-        # Add user message with Markdown image preview
         history.append({"role": "user", "content": f"![](file={image_path})\n\n{user_text}"})
-        
-        # Add loading message
-        history.append({"role": "assistant", "content": "🔄 Starting analysis..."})
+        history.append({"role": "assistant", "content": "🔄 Starting analysis (Model loading may take a moment)..."})
 
-        # Hide welcome, show loading, hide results
         yield history, session_id, image_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *empty_slider_updates, None, gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
         
-        # Run full analysis
         final_result = None
         async for result in run_analysis(image_path, user_text, session_id):
-            # result = (history, session_id, overlay, checkboxes, download, df_m, df_s, df_r, slider1-4)
             final_result = result
-            # Update the history from run_analysis but preserve our image indicator
             updated_history = result[0].copy()
             if files and len(updated_history) > 0:
                 updated_history[0] = {"role": "user", "content": f"![](file={image_path})\n\n{user_text}"}
-            
-            # Keep loading overlay visible during processing
             yield (updated_history, result[1], image_path, *result[2:], None, gr.update(visible=False), gr.update(visible=True), gr.update(visible=False))
         
-        # Final yield: hide welcome, hide loading, show results
         if final_result:
             updated_history = final_result[0].copy()
             if files and len(updated_history) > 0:
                 updated_history[0] = {"role": "user", "content": f"![](file={image_path})\n\n{user_text}"}
             yield (updated_history, final_result[1], image_path, *final_result[2:], None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=True))
-        
         return
 
-    # Case 2: Follow-up question (has session, no new image)
+    # Case 2: Follow-up
     elif session_id and user_text:
-        # Add user message to history
         history.append({"role": "user", "content": user_text})
         history.append({"role": "assistant", "content": "💭 Thinking..."})
-        
-        # Keep all overlays in their current state
         yield history, session_id, current_img_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *empty_slider_updates, None, gr.update(), gr.update(), gr.update()
         
         if not ACTIVE_RUNNER:
-            history[-1]["content"] = "⚠️ Session expired. Please upload an image and start a new analysis."
+            history[-1]["content"] = "⚠️ Session expired."
+            yield history, None, current_img_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *empty_slider_updates, None, gr.update(), gr.update(), gr.update()
+            return
+        
+        content = types.Content(role="user", parts=[types.Part.from_text(text=user_text)])
+        accumulated_response = ""
+        try:
+            async for event in ACTIVE_RUNNER.run_async(user_id="demo_user", session_id=session_id, new_message=content):
+                if event.content and event.content.parts:
+                    for part in event.content.parts:
+                        if hasattr(part, 'text') and part.text:
+                            accumulated_response += part.text
+                            history[-1]["content"] = accumulated_response
+                            yield history, session_id, current_img_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *empty_slider_updates, None, gr.update(), gr.update(), gr.update()
+        except Exception as e:
+            history[-1]["content"] = f"❌ Error: {e}"
+            yield history, session_id, current_img_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *empty_slider_updates, None, gr.update(), gr.update(), gr.update()
+        return
+
+    else:
+        if not history:
+            history = [{"role": "assistant", "content": "👋 Welcome! Upload a microscopy image and describe what you'd like to analyze."}]
+        else:
+            history.append({"role": "assistant", "content": "⚠️ Please provide a question or upload a new image."})
+        yield history, session_id, current_img_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *empty_slider_updates, None, gr.update(), gr.update(), gr.update()
+
+# --- UI Layout ---
+with gr.Blocks(title="Cellemetry Agent") as demo:
+    session_id_state = gr.State(None)
+    current_image_path = gr.State(None)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            chatbot = gr.Chatbot(
+                label="Agent Conversation", 
+                height=600, 
+                value=[{"role": "assistant", "content": "👋 Welcome to Cellemetry! Upload a microscopy image and describe what you'd like to analyze."}],
+                show_label=True,
+                type="messages" 
+            )
+            chat_input = gr.MultimodalTextbox(
+                file_types=["image"],
+                placeholder="Upload an image and describe your analysis...",
+                show_label=False,
+                submit_btn="Send"
+            )
+
+        with gr.Column(scale=2):
+            with gr.Column(visible=True, elem_id="welcome-overlay") as welcome_overlay:
+                gr.HTML(f"""
+                    <div style="display: flex; flex-direction: column; align-items: center; justify-content: center; min-height: 780px; padding: 40px; background: #f8f9fa; border-radius: 8px; border: 2px solid #3498db;">
+                        <div style="text-align: center;">
+                            <div style='text-align: center;'>
+                                <img src="https://raw.githubusercontent.com/hmgill/Cellemetry/main/logo.png" alt="Logo" style="height:200px; display: block; margin: 0 auto;">
+                            </div>
+                            <h2 style="color: #333; margin: 20px 0 10px; font-weight: 600; font-size: 28px;">Welcome to Cellemetry</h2>
+                            <p style="color: #666; font-size: 16px; max-width: 400px; margin: 0 auto 30px; line-height: 1.6;">Upload a microscopy image to get started with AI-powered cell analysis and segmentation</p>
+                            <div style="padding: 20px; background: #fff; border-radius: 8px; border-left: 4px solid #3498db; box-shadow: 0 2px 4px rgba(0,0,0,0.05);">
+                                <p style="color: #555; margin: 0; font-size: 14px;">👈 Use the chat on the left to begin</p>
+                            </div>
+                        </div>
+                    </div>
+                    """
+                )
+            
+            with gr.Column(visible=False, elem_id="loading-overlay") as loading_overlay:
+                gr.HTML("""
+                    <div style="display: flex; flex-direction: column; align-items: center; justify-content: center; height: 780px; background: rgba(255, 255, 255, 0.95); border-radius: 8px; box-shadow: 0 4px 6px rgba(0,0,0,0.1);">
+                        <div style="text-align: center;">
+                            <div style="border: 8px solid #f3f3f3; border-top: 8px solid #3498db; border-radius: 50%; width: 60px; height: 60px; animation: spin 1s linear infinite; margin: 0 auto 20px;"></div>
+                            <h3 style="color: #555; margin: 0;">⚙️ Analysis in Progress</h3>
+                            <p style="color: #888; margin-top: 10px;">Please wait while we process your microscopy image...</p>
+                        </div>
+                        <style>@keyframes spin { 0% { transform: rotate(0deg); } 100% { transform: rotate(360deg); } }</style>
+                    </div>
+                """)
+            
+            with gr.Tabs(visible=False) as results_tabs:
+                with gr.Tab("🔍 Segmentation"):
+                    with gr.Row():
+                        with gr.Column(scale=3):
+                            overlay_output = gr.Image(label="Segmentation Result", height=780, type="pil")
+                        with gr.Column(scale=1):
+                            gr.Markdown("**Layer Controls**")
+                            layer_checkboxes = gr.CheckboxGroup(label="Visible Layers", choices=[], value=[], interactive=True)
+                            gr.Markdown("**Opacity Controls**")
+                            opacity_slider_1 = gr.Slider(minimum=0, maximum=1, value=0.6, step=0.1, label="Layer 1 Opacity", visible=False)
+                            opacity_slider_2 = gr.Slider(minimum=0, maximum=1, value=0.6, step=0.1, label="Layer 2 Opacity", visible=False)
+                            opacity_slider_3 = gr.Slider(minimum=0, maximum=1, value=0.6, step=0.1, label="Layer 3 Opacity", visible=False)
+                            opacity_slider_4 = gr.Slider(minimum=0, maximum=1, value=0.6, step=0.1, label="Layer 4 Opacity", visible=False)
+                
+                with gr.Tab("📊 Quantitative Results"):
+                    download_btn = gr.File(label="Download Excel Report")
+                    with gr.Tabs():
+                        with gr.Tab("Morphology"):
+                            tbl_morph = gr.Dataframe(interactive=False, wrap=True)
+                        with gr.Tab("Spatial"):
+                            tbl_spatial = gr.Dataframe(interactive=False, wrap=True)
+                        with gr.Tab("Relational"):
+                            tbl_rel = gr.Dataframe(interactive=False, wrap=True)
+
+    def regenerate_overlay_with_opacity(img_path, selected_layers, op1, op2, op3, op4):
+        if not img_path or not selected_layers: return None
+        opacities = {}
+        opacity_values = [op1, op2, op3, op4]
+        all_layers = get_available_layers()
+        for i, layer in enumerate(all_layers[:4]):
+            opacities[layer] = opacity_values[i]
+        return generate_overlay(img_path, selected_layers, opacities)
+
+    chat_input.submit(
+        fn=unified_chat_handler,
+        inputs=[chat_input, chatbot, session_id_state, current_image_path],
+        outputs=[chatbot, session_id_state, current_image_path, overlay_output, layer_checkboxes, download_btn, tbl_morph, tbl_spatial, tbl_rel, opacity_slider_1, opacity_slider_2, opacity_slider_3, opacity_slider_4, chat_input, welcome_overlay, loading_overlay, results_tabs]
+    )
+    
+    for component in [layer_checkboxes, opacity_slider_1, opacity_slider_2, opacity_slider_3, opacity_slider_4]:
+        component.change(
+            fn=regenerate_overlay_with_opacity,
+            inputs=[current_image_path, layer_checkboxes, opacity_slider_1, opacity_slider_2, opacity_slider_3, opacity_slider_4],
+            outputs=[overlay_output]
+        )
+
+    # TRIGGER MODEL LOAD AFTER UI LAUNCH
+    demo.load(load_models)
+
+if __name__ == "__main__":
+    demo.queue().launch(
+        ssr_mode=False, 
+        theme=gr.themes.Soft(), 
+        server_name="0.0.0.0", 
+        server_port=7860,
+        allowed_paths=["."] 
+    )