Update app.py

app.py

@@ -19,7 +19,6 @@ from google.adk.runners import InMemoryRunner
 from google.genai import types
 
 # Project Imports
-# Wrap imports to prevent immediate crash if dependencies are missing
 try:
     from cellemetry import root_agent
     from cellemetry.config import AnalysisDeps
@@ -39,12 +38,17 @@ MODEL_CACHE = {
     "loaded": False
 }
 
-# Store the active runner globally
-ACTIVE_RUNNER = None
+# OPTIMIZATION: In-memory cache for masks and base image to prevent disk I/O on slider updates
+MASK_CACHE = {
+    "current_path": None,
+    "base_image": None,  # PIL RGBA Image
+    "layers": {}         # Dict of 'layer_name': numpy_boolean_mask
+}
 
+ACTIVE_RUNNER = None
 
 def load_models():
-    """Initialize SAM3 model. Now called AFTER app startup to prevent timeout."""
+    """Initialize SAM3 model. Now called AFTER app startup."""
     if MODEL_CACHE["loaded"]:
         return
         
@@ -53,7 +57,6 @@ def load_models():
     MODEL_CACHE["device"] = device
     
     try:
-        # Check if imports succeeded
         if Sam3Model is None:
             raise ImportError("Sam3Model not found. Please check requirements.")
 
@@ -84,9 +87,7 @@ def load_excel_data(logs_text):
             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
@@ -118,64 +119,101 @@ def update_opacity_sliders(layers):
             updates.append(gr.update(visible=False))
     return updates
 
+# --- OPTIMIZED OVERLAY GENERATION ---
 def generate_overlay(image_path_str, selected_layers, layer_opacities=None):
+    """Regenerates overlay using in-memory caching for speed."""
     if not image_path_str:
         return None
     
-    base_image = Image.open(image_path_str).convert("RGBA")
+    # 1. Check if we need to load data into cache
+    # We reload if the path changes OR if the cache is empty
+    if MASK_CACHE["current_path"] != image_path_str or MASK_CACHE["base_image"] is None:
+        print(f"🔄 Caching masks for {os.path.basename(image_path_str)}...")
+        try:
+            # Load Base Image
+            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"] = {} # Clear old layers
+
+            # Pre-load and resize ALL available layers to match base image
+            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", "")
+                try:
+                    data = np.load(file_path)
+                    masks = data['masks'] if 'masks' in data else data[data.files[0]]
+                    
+                    if masks.size > 0:
+                        if masks.ndim == 3:
+                            combined_mask = np.max(masks, axis=0)
+                        else:
+                            combined_mask = masks
+                        
+                        # Resize boolean mask to match base image ONCE
+                        # We use PIL for high-quality resizing of the mask
+                        mask_pil = Image.fromarray(combined_mask.astype(np.uint8) * 255)
+                        if mask_pil.size != (base_w, base_h):
+                            mask_pil = mask_pil.resize((base_w, base_h), Image.Resampling.NEAREST)
+                        
+                        # Store as boolean numpy array for fast processing
+                        MASK_CACHE["layers"][layer_name] = np.array(mask_pil, dtype=bool)
+                except Exception as e:
+                    print(f"Failed to cache layer {layer_name}: {e}")
+        except Exception as e:
+            print(f"Failed to load base image: {e}")
+            return None
+
+    # 2. Fast Composition using Cache
+    if MASK_CACHE["base_image"] is None:
+        return None
 
+    base_image = MASK_CACHE["base_image"]
+    # Start with a transparent overlay layer
+    overlay_accum = Image.new('RGBA', base_image.size, (0, 0, 0, 0))
+    
     color_map = {
         "green_cell": (0, 255, 0),
         "blue_nucleus": (0, 0, 255),
         "cell": (0, 255, 0),
         "nucleus": (0, 0, 255),
     }
-    
-    overlay = Image.new('RGBA', base_image.size, (0, 0, 0, 0))
-    
-    for layer_name in selected_layers:
-        file_path = f"/tmp/data_{layer_name}.npz"
-        if os.path.exists(file_path):
-            try:
-                data = np.load(file_path)
-                masks = data['masks'] if 'masks' in data else data[data.files[0]]
-                
-                if masks.size > 0:
-                    if masks.ndim == 3:
-                        combined_mask = np.max(masks, axis=0)
-                    else:
-                        combined_mask = masks
-                    
-                    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))
-                    
-                    opacity = 0.5
-                    if layer_opacities and layer_name in layer_opacities:
-                        opacity = layer_opacities[layer_name]
-                    
-                    mask_overlay = np.zeros((*combined_mask.shape, 4), dtype=np.uint8)
-                    mask_overlay[combined_mask] = (*color, int(255 * opacity))
-                    
-                    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}")
+    # Iterate through requested layers
+    for layer_name in selected_layers:
+        if layer_name in MASK_CACHE["layers"]:
+            mask_bool = MASK_CACHE["layers"][layer_name]
+            
+            # Get settings
+            color = color_map.get(layer_name.lower(), (255, 255, 0))
+            opacity = 0.6 # Default
+            if layer_opacities and layer_name in layer_opacities:
+                opacity = layer_opacities[layer_name]
+            
+            # Create a solid color block
+            # Optimization: We construct the RGBA buffer directly
+            layer_rgba = np.zeros((mask_bool.shape[0], mask_bool.shape[1], 4), dtype=np.uint8)
+            
+            # Apply color and opacity only where mask is True
+            layer_rgba[mask_bool] = (*color, int(255 * opacity))
+            
+            # Composite using PIL (fast C implementation)
+            layer_img = Image.fromarray(layer_rgba, 'RGBA')
+            overlay_accum = Image.alpha_composite(overlay_accum, layer_img)
 
-    result = Image.alpha_composite(base_image, overlay)
+    # Final Composite
+    result = Image.alpha_composite(base_image, overlay_accum)
     return result.convert("RGB")
 
 # --- Core Logic ---
 async def run_analysis(image_path_str, user_prompt, session_id_state):
-    """Runs analysis. Triggers model load if not yet ready."""
+    """Runs analysis."""
     waiting_df = pd.DataFrame({"Status": ["Waiting..."]})
     empty_slider_updates = [gr.update()] * 4
     
-    # Lazy Load: Ensure model is loaded before inference
+    # Lazy Load
     if not MODEL_CACHE["loaded"]:
         yield [], None, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
         load_models() 
@@ -184,15 +222,19 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
         yield [], None, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
         return
 
-    # Cleanup
+    # Cleanup Files AND Cache
     for f in glob.glob("/tmp/out_*.png") + glob.glob("/tmp/data_*.npz") + glob.glob("/tmp/*.xlsx"):
         try: os.remove(f)
         except: pass
+    
+    # Reset Cache for new run
+    MASK_CACHE["current_path"] = None 
+    MASK_CACHE["base_image"] = None
+    MASK_CACHE["layers"] = {}
 
     # 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
@@ -229,7 +271,6 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
     
     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}]
 
@@ -272,6 +313,8 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
     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 generation (will trigger cache population)
     initial_overlay = generate_overlay(image_path_str, layers)
 
     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."
@@ -303,7 +346,6 @@ 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."
         
-        # Add preview image using Markdown
         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)..."})
 
@@ -361,15 +403,11 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
 
 # CSS to handle margins and width consistency
 custom_css = """
-/* 1. Global Margin Setting */
 #main_container {
-    margin-left: 20% !important;
-    margin-right: 20% !important;
+    margin-left: 10% !important;
+    margin-right: 10% !important;
     width: auto !important;
 }
-
-/* 2. Fix Tab Width Consistency (Right Panel) */
-/* We enforce a minimum width so the panel doesn't shrink when switching to empty dataframes */
 .right-panel {
     min-width: 600px !important;
     flex-grow: 2 !important;
@@ -380,7 +418,6 @@ with gr.Blocks(title="Cellemetry Agent", css=custom_css) as demo:
     session_id_state = gr.State(None)
     current_image_path = gr.State(None)
     
-    # WRAPPER to apply 20% margins
     with gr.Column(elem_id="main_container"):
         
         with gr.Row():
@@ -400,7 +437,6 @@ with gr.Blocks(title="Cellemetry Agent", css=custom_css) as demo:
                 )
 
             # --- RIGHT COLUMN (Results) ---
-            # Added 'right-panel' class for CSS width enforcement
             with gr.Column(scale=2, elem_classes=["right-panel"]):
                 
                 # Welcome overlay