Update app.py

app.py

@@ -12,7 +12,6 @@ from pathlib import Path
 from PIL import Image
 
 # --- Safe Input Mocking ---
-# Fix: Mock input to prevent hanging on server consoles
 builtins.input = lambda *args: "y"
 
 # GenAI & ADK Imports
@@ -29,7 +28,7 @@ except ImportError as e:
     Sam3Model = None
     Sam3Processor = None
     root_agent = None
-    AnalysisDeps = None  # FIX: Correctly indented inside except block
+    AnalysisDeps = None
 
 # Optional: Distinctipy for better colors
 try:
@@ -100,6 +99,32 @@ def load_models():
         return f"⚠️ Model load failed: {e}"
 
 # --- Helpers ---
+def clean_layer_name(filename):
+    """
+    Converts 'data_blue_nuclei.npz' -> 'Nuclei'.
+    Removes standard color names and underscores.
+    """
+    # Remove prefix/suffix
+    raw = os.path.basename(filename).replace("data_", "").replace(".npz", "")
+    
+    # Split into parts
+    parts = raw.split('_')
+    
+    # Filter out common colors
+    colors = {
+        'blue', 'green', 'red', 'yellow', 'cyan', 'magenta', 
+        'orange', 'purple', 'white', 'black', 'gray', 'grey', 
+        'pink', 'brown', 'lime', 'teal'
+    }
+    
+    cleaned_parts = [p for p in parts if p.lower() not in colors]
+    
+    # Join back and Title Case
+    if not cleaned_parts:
+        return raw.replace("_", " ").title()
+        
+    return " ".join(cleaned_parts).title()
+
 def load_excel_data(logs_text):
     placeholder = pd.DataFrame({"Status": ["No Data Available"]})
     candidates = glob.glob("/tmp/*.xlsx") + glob.glob("*.xlsx")
@@ -134,15 +159,18 @@ def get_available_layers():
     files = glob.glob("/tmp/data_*.npz")
     layers = []
     for f in files:
-        name = os.path.basename(f).replace("data_", "").replace(".npz", "")
-        layers.append(name)
-    return sorted(layers)
+        # UPDATED: Use the clean name helper
+        layers.append(clean_layer_name(f))
+    
+    # Return unique sorted layers
+    return sorted(list(set(layers)))
 
 def update_opacity_sliders(layers):
     updates = []
     for i in range(4):
         if i < len(layers):
-            layer_name = layers[i].replace("_", " ").title()
+            # Layer name is already clean from get_available_layers
+            layer_name = layers[i]
             updates.append(gr.update(visible=True, label=f"{layer_name} Opacity", value=0.6))
         else:
             updates.append(gr.update(visible=False))
@@ -177,7 +205,8 @@ def generate_overlay(image_path_str, selected_layers, layer_opacities=None, forc
             base_w, base_h = base_img.size
             
             for file_path in all_layer_files:
-                layer_name = os.path.basename(file_path).replace("data_", "").replace(".npz", "")
+                # UPDATED: Use clean name as key
+                layer_name = clean_layer_name(file_path)
                 
                 # Skip if already cached
                 if layer_name in MASK_CACHE["layers"]:
@@ -269,13 +298,11 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
         yield [{"role": "assistant", "content": error_msg}], None, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
         return
 
-    # FIX: Safety check for AnalysisDeps availability
     if AnalysisDeps is None:
         error_msg = "❌ Project imports failed. 'AnalysisDeps' is missing. Check your 'cellemetry' package installation."
         yield [{"role": "assistant", "content": error_msg}], None, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
         return
 
-    # FIX: Wrap Agent/Runner initialization in try/except to catch crashes here
     try:
         deps = AnalysisDeps(
             sam_model=MODEL_CACHE["model"],
@@ -315,12 +342,11 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
 
     logs = [f"🔄 **Starting analysis** on {MODEL_CACHE['device']}..."]
     
-    # FIX: Encode spaces for Markdown
+    # Encode spaces for Markdown
     display_path = image_path_str.replace(" ", "%20")
 
     def yield_status(log_list):
         full_log = "\n\n".join(log_list)
-        # Use display_path
         user_msg = f"![](file={display_path})\n\n{user_prompt}"
         return [{"role": "user", "content": user_msg}, {"role": "assistant", "content": full_log}]
 
@@ -369,7 +395,6 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
     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
     
-    # Use display_path here too
     final_user_msg = f"![](file={display_path})\n\n{user_prompt}"
     final_history = [{"role": "user", "content": final_user_msg}, {"role": "assistant", "content": full_log_text}]
     slider_updates = update_opacity_sliders(layers)
@@ -396,26 +421,21 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
     if image_path and (not session_id or files):
         if not user_text: user_text = "Analyze this microscopy image."
         
-        # FIX: Encode spaces for Markdown
         display_path = image_path.replace(" ", "%20")
         
         history.append({"role": "user", "content": f"![](file={display_path})\n\n{user_text}"})
         history.append({"role": "assistant", "content": "🔄 Starting analysis (Model loading may take a moment)..."})
 
-        # Yield 1: Set initial visibility state ONCE
         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)
         
         final_result = None
-        # FIX: Catch potential generator crashes here to prevent asyncio loop death
         try:
             async for result in run_analysis(image_path, user_text, session_id):
                 final_result = result
                 updated_history = result[0].copy()
                 if files and len(updated_history) > 0:
-                    # Ensure the user message in history also uses the encoded path
                     updated_history[0] = {"role": "user", "content": f"![](file={display_path})\n\n{user_text}"}
                 
-                # Yield Loop: Pass gr.update() to prevent flickering
                 yield (updated_history, result[1], image_path, *result[2:], None, gr.update(), gr.update(), gr.update())
         except Exception as e:
             history.append({"role": "assistant", "content": f"❌ Critical Error: {str(e)}"})
@@ -427,7 +447,6 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
             if files and len(updated_history) > 0:
                 updated_history[0] = {"role": "user", "content": f"![](file={display_path})\n\n{user_text}"}
             
-            # Yield Final: Show Results
             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
 
@@ -445,7 +464,6 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
         content = types.Content(role="user", parts=[types.Part.from_text(text=user_text)])
         accumulated_response = ""
         
-        # 1. Stream response text
         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:
@@ -459,15 +477,12 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
             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
 
-        # 2. REFRESH DATA (Tables, Overlays, Layers)
         report_file, df_m, df_s, df_r = load_excel_data("") 
         layers = get_available_layers()
         
-        # Force overlay generation with new layers (using force_reload=True to clear cache)
         new_overlay = generate_overlay(current_img_path, layers, force_reload=True)
         slider_updates = update_opacity_sliders(layers)
         
-        # 3. Yield FINAL update with new data
         yield (
             history, 
             session_id,