Update app.py

app.py

@@ -216,69 +216,66 @@ def measure_confluency(masks, image_np):
     confluency = cell_pixels / tot_pixels * 100
     return confluency
 
-def filter_mask_by_size(masks,minimum_pixels):
-    filtered_masks=masks.copy()
+def filter_mask_by_size(masks, minimum_pixels):
+    filtered_masks = masks.copy()
     cell_ids = np.unique(masks)
     cell_ids = cell_ids[cell_ids > 0]
 
     removed_count = 0
-
+    
     for cell_id in cell_ids:
         cell_mask = (masks == cell_id)
         cell_pixels = np.count_nonzero(cell_mask)
-
         if cell_pixels < minimum_pixels:
             filtered_masks[cell_mask] = 0
-            removed_count +=1
+            removed_count += 1
 
     unique_ids = np.unique(filtered_masks)
     unique_ids = unique_ids[unique_ids > 0]
-    
+
     renumbered_masks = np.zeros_like(filtered_masks)
     for new_id, old_id in enumerate(unique_ids, start=1):
         renumbered_masks[filtered_masks == old_id] = new_id
 
-    
     return renumbered_masks, removed_count
-    
-def filter_mask_by_maxsize(masks,maximum_pixels):
-    filtered_masks=masks.copy()
+
+
+def filter_mask_by_maxsize(masks, maximum_pixels):
+    filtered_masks = masks.copy()
     cell_ids = np.unique(masks)
     cell_ids = cell_ids[cell_ids > 0]
 
     removed_count = 0
-
     for cell_id in cell_ids:
         cell_mask = (masks == cell_id)
         cell_pixels = np.count_nonzero(cell_mask)
-
         if cell_pixels > maximum_pixels:
             filtered_masks[cell_mask] = 0
-            removed_count +=1
+            removed_count += 1
 
     unique_ids = np.unique(filtered_masks)
     unique_ids = unique_ids[unique_ids > 0]
-    
+
     renumbered_masks = np.zeros_like(filtered_masks)
     for new_id, old_id in enumerate(unique_ids, start=1):
         renumbered_masks[filtered_masks == old_id] = new_id
 
-    
     return renumbered_masks, removed_count
-    
+
+
 def rec_min_size(masks):
-    num_cells = len(np.unique(masks)) - 1
+    ids = np.unique(masks)
+    ids = ids[ids > 0]
+    num_cells = len(ids)
     if num_cells <= 0:
         return 0
     mean_cell_size = np.count_nonzero(masks) / num_cells
     return int(round(mean_cell_size))
 
+
 @spaces.GPU
 def run_segmentation_editor(editor_data, model_choice, min_cell_size, max_cell_size):
-    """
-    Runs cell segmentation using ImageEditor data.
-    Returns initial segmentation overlay, counts, confluency, and also masks/image for state.
-    """
+
     try:
         model_filename = MODEL_OPTIONS[model_choice]
         model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=model_filename)
@@ -294,7 +291,6 @@ def run_segmentation_editor(editor_data, model_choice, min_cell_size, max_cell_s
         if region_np is None:
             return 0, None, "No image provided.", gr.update(visible=False), None, None, 0.0, gr.update()
 
-
         # Enforce mobile-safe size limit immediately
         region_np = safe_resize(region_np)
 
@@ -307,55 +303,59 @@ def run_segmentation_editor(editor_data, model_choice, min_cell_size, max_cell_s
             processed_image_np = region_np
 
         # Run Cellpose segmentation
-        masks, flows, styles = model.eval(processed_image_np, diameter=None, channels=[0, 0])
+        masks_raw, flows, styles = model.eval(processed_image_np, diameter=None, channels=[0, 0])
+
+        # Compute recommendation from RAW masks 
+        recommend_min = rec_min_size(masks_raw)
 
+        
+        # If user sets slider to 0, use the recommendation
+        min_used = recommend_min if (min_cell_size == 0) else int(min_cell_size)
+
+        # Apply filters
+        masks = masks_raw
         removed_small = 0
         removed_large = 0
-        
-        # Minimum size filtering
 
-        recommend_min = rec_min_size(masks) 
-        
-        if min_cell_size > 0:
-            masks, removed_small = filter_mask_by_size(masks, min_cell_size)
-            filter_msg = f"Removed {removed_small} small objects (< {min_cell_size} pixels).\n"
-        else:
-            removed_small = 0
-            filter_msg=""
-            
-        
+        if min_used > 0:
+            masks, removed_small = filter_mask_by_size(masks, min_used)
 
-        # Maximum size filtering
         if max_cell_size > 0:
-            masks, removed_large = filter_mask_by_maxsize(masks, max_cell_size)
-            filter_msg = f"Removed {removed_large} large objects (> {max_cell_size} pixels).\n"
-        else:
-            removed_large = 0
-            filter_msg=""
+            masks, removed_large = filter_mask_by_maxsize(masks, int(max_cell_size))
 
-        removed_count = removed_small + removed_large
         
+        filter_msg = ""
+        if removed_small:
+            filter_msg += f"Removed {removed_small} small objects (< {min_used} pixels).\n"
+        if removed_large:
+            filter_msg += f"Removed {removed_large} large objects (> {int(max_cell_size)} pixels).\n"
+
         cell_count = len(np.unique(masks)) - 1
         confluency = measure_confluency(masks, processed_image_np)
+
         # Create a basic segmentation overlay (without viability)
         segmentation_overlay = processed_image_np.copy().astype(np.float32)
         if masks.max() > 0:
             np.random.seed(42)  # For consistent random colors
             colors = np.random.randint(0, 255, size=(masks.max() + 1, 3))
-            colors[0] = [0, 0, 0]  # Background color
+            colors[0] = [0, 0, 0]
             colored_mask = colors[masks]
             alpha = 0.4
             segmentation_overlay = (1 - alpha) * segmentation_overlay + alpha * colored_mask
         segmentation_overlay = np.clip(segmentation_overlay, 0, 255).astype(np.uint8)
 
-        info_msg = f"Segmentation complete! Found {cell_count} cells.\n"
+        info_msg = ""
+        if filter_msg:
+            info_msg += filter_msg
+        info_msg += f"Segmentation complete! Found {cell_count} cells.\n"
         info_msg += f"Confluency: {confluency:.1f}%\n"
-        info_msg = filter_msg + info_msg
         if region_coords:
-            info_msg += f"Processed region: {region_coords[0]},{region_coords[1]} to {region_coords[2]},{region_coords[3]}\n"
-        info_msg += f"Now adjust the Blue Threshold for viability assessment."
+            info_msg += (
+                f"Processed region: {region_coords[0]},{region_coords[1]} "
+                f"to {region_coords[2]},{region_coords[3]}\n"
+            )
+        info_msg += "Now adjust the Blue Threshold for viability assessment."
 
-        # Return initial segmentation display and state variables (packed for Gradio State)
         return (
             cell_count,
             Image.fromarray(segmentation_overlay),
@@ -364,12 +364,20 @@ def run_segmentation_editor(editor_data, model_choice, min_cell_size, max_cell_s
             pack_array(masks),
             pack_array(processed_image_np),
             confluency,
-            gr.update(value = recommend_min)
+            gr.update(value=recommend_min),  # update slider display to recommended
         )
 
     except Exception as e:
-        return 0, None, f"Error during segmentation: {str(e)}", gr.update(visible=False), None, None, 0.0, gr.update()
-
+        return (
+            0,
+            None,
+            f"Error during segmentation: {str(e)}",
+            gr.update(visible=False),
+            None,
+            None,
+            0.0,
+            gr.update(),
+        )
 
 
 
@@ -448,9 +456,9 @@ with gr.Blocks(
                 min_size_slider1 = gr.Slider(
                     minimum=0,
                     maximum=500,
-                    value=50,
+                    value=0,
                     step=10,
-                    label="Minimum Cell Size (pixels)",
+                    label="Minimum Cell Size (pixels). Leave at zero for automated recommendation",
     
                 )
                 max_size_slider1 = gr.Slider(