Update app.py

app.py

@@ -30,6 +30,13 @@ except ImportError as e:
     root_agent = None
     AnalysisDeps = None
 
+# Optional: Distinctipy for better colors
+try:
+    from distinctipy import distinctipy
+except ImportError:
+    distinctipy = None
+    print("⚠️ distinctipy not found. Using fallback colors.")
+
 # --- Global State ---
 MODEL_CACHE = {
     "model": None,
@@ -38,7 +45,6 @@ MODEL_CACHE = {
     "loaded": False
 }
 
-# In-memory cache for masks and base image to prevent disk I/O on slider updates
 MASK_CACHE = {
     "current_path": None,
     "base_image": None,
@@ -47,18 +53,28 @@ MASK_CACHE = {
 
 ACTIVE_RUNNER = None
 
-# FIX: Dynamic Color Palette (Removed hard-coded "nucleus" keys)
-# Colors: Green, Blue, Red, Yellow, Cyan, Magenta, Orange, Purple
-COLOR_PALETTE = [
-    (0, 255, 0),    # Green
-    (0, 100, 255),  # Blue
-    (255, 0, 0),    # Red
-    (255, 255, 0),  # Yellow
-    (0, 255, 255),  # Cyan
-    (255, 0, 255),  # Magenta
-    (255, 128, 0),  # Orange
-    (128, 0, 255)   # Purple
-]
+# --- Dynamic Color Helper ---
+def generate_color_palette(n=50):
+    """Generates a palette of N distinct colors [0-255]."""
+    if distinctipy:
+        print(f"🎨 Generating {n} distinct colors using distinctipy...")
+        colors = distinctipy.get_colors(n)
+        return [tuple(int(c * 255) for c in color) for color in colors]
+    
+    try:
+        import matplotlib.pyplot as plt
+        cmap = plt.get_cmap('tab20')
+        return [tuple(int(c * 255) for c in cmap(i % 20)[:3]) for i in range(n)]
+    except Exception:
+        pass
+
+    return [
+        (0, 255, 0), (0, 0, 255), (255, 0, 0), (255, 255, 0), 
+        (0, 255, 255), (255, 0, 255), (255, 128, 0), (128, 0, 255),
+        (0, 128, 0), (0, 0, 128), (128, 0, 0), (128, 128, 0)
+    ] * (n // 12 + 1)
+
+COLOR_PALETTE = generate_color_palette(50)
 
 def load_models():
     """Initialize SAM3 model. Now called AFTER app startup."""
@@ -84,7 +100,6 @@ def load_models():
 
 # --- Helpers ---
 def load_excel_data(logs_text):
-    """Finds and loads the Excel report, transposing for better display."""
     placeholder = pd.DataFrame({"Status": ["No Data Available"]})
     candidates = glob.glob("/tmp/*.xlsx") + glob.glob("*.xlsx")
     
@@ -133,24 +148,40 @@ def update_opacity_sliders(layers):
     return updates
 
 # --- OPTIMIZED OVERLAY GENERATION ---
-def generate_overlay(image_path_str, selected_layers, layer_opacities=None):
+def generate_overlay(image_path_str, selected_layers, layer_opacities=None, force_reload=False):
+    """
+    Regenerates overlay.
+    force_reload: If True, clears the layer cache to pick up new files from agent.
+    """
     if not image_path_str:
         return None
     
-    # Check cache
-    if MASK_CACHE["current_path"] != image_path_str or MASK_CACHE["base_image"] is None:
+    # Force reload if requested (e.g., after follow-up analysis)
+    if force_reload:
+        MASK_CACHE["layers"] = {}
+
+    # Check cache loading
+    if MASK_CACHE["current_path"] != image_path_str or MASK_CACHE["base_image"] is None or not MASK_CACHE["layers"]:
         print(f"🔄 Caching masks for {os.path.basename(image_path_str)}...")
         try:
-            base_img = Image.open(image_path_str).convert("RGBA")
-            MASK_CACHE["base_image"] = base_img
-            MASK_CACHE["current_path"] = image_path_str
-            MASK_CACHE["layers"] = {}
+            if MASK_CACHE["current_path"] != image_path_str or MASK_CACHE["base_image"] is None:
+                base_img = Image.open(image_path_str).convert("RGBA")
+                MASK_CACHE["base_image"] = base_img
+                MASK_CACHE["current_path"] = image_path_str
+            else:
+                base_img = MASK_CACHE["base_image"]
 
+            # Always scan for new layers if we are here
             all_layer_files = glob.glob("/tmp/data_*.npz")
             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", "")
+                
+                # Skip if already cached
+                if layer_name in MASK_CACHE["layers"]:
+                    continue
+                    
                 try:
                     data = np.load(file_path)
                     masks = data['masks'] if 'masks' in data else data[data.files[0]]
@@ -178,20 +209,18 @@ def generate_overlay(image_path_str, selected_layers, layer_opacities=None):
     base_image = MASK_CACHE["base_image"]
     overlay_accum = Image.new('RGBA', base_image.size, (0, 0, 0, 0))
     
-    # Get all available layers to ensure consistent coloring regardless of selection
     all_known_layers = sorted(MASK_CACHE["layers"].keys())
 
     for layer_name in selected_layers:
         if layer_name in MASK_CACHE["layers"]:
             mask_bool = MASK_CACHE["layers"][layer_name]
             
-            # FIX: Assign color based on sorted index instead of hard-coded map
-            # This ensures "Mitochondria" gets a color even if we didn't plan for it
+            # Use the global generated palette
             if layer_name in all_known_layers:
                 color_idx = all_known_layers.index(layer_name) % len(COLOR_PALETTE)
                 color = COLOR_PALETTE[color_idx]
             else:
-                color = (255, 255, 0) # Fallback Yellow
+                color = (255, 255, 0)
 
             opacity = 0.6
             if layer_opacities and layer_name in layer_opacities:
@@ -336,12 +365,14 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
     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):
         if not user_text: user_text = "Analyze this microscopy image."
         
         history.append({"role": "user", "content": f"![](file={image_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
@@ -350,15 +381,20 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
             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}"}
+            
+            # Yield Loop: Pass gr.update() to prevent flickering
             yield (updated_history, result[1], image_path, *result[2:], None, gr.update(), gr.update(), gr.update())
         
         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 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
 
+    # CASE 2: FOLLOW-UP ANALYSIS
     elif session_id and user_text:
         history.append({"role": "user", "content": user_text})
         history.append({"role": "assistant", "content": "💭 Thinking..."})
@@ -371,6 +407,8 @@ 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:
@@ -382,6 +420,32 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
         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
+
+        # 2. REFRESH DATA (Tables, Overlays, Layers)
+        # Scan for potential new files created by the agent
+        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, 
+            current_img_path, 
+            new_overlay,                            # Updated Image
+            gr.CheckboxGroup(value=layers, choices=layers), # Updated Layers
+            report_file,                            # Updated Excel
+            df_m, df_s, df_r,                       # Updated Tables
+            *slider_updates,                        # Updated Sliders
+            None,                                   # Clear input
+            gr.update(),                            # Welcome (no change)
+            gr.update(),                            # Loading (no change)
+            gr.update()                             # Results (no change)
+        )
         return
 
     else:
@@ -396,8 +460,8 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
 custom_css = """
 /* 1. Global Margin Setting */
 #main_container {
-    margin-left: 10% !important;
-    margin-right: 10% !important;
+    margin-left: 20% !important;
+    margin-right: 20% !important;
     width: auto !important;
 }