Update app.py

app.py

@@ -241,7 +241,31 @@ def filter_mask_by_size(masks,minimum_pixels):
     
     return renumbered_masks, removed_count
     
+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
+
+    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
     if num_cells <= 0:
@@ -250,7 +274,7 @@ def rec_min_size(masks):
     return int(round(mean_cell_size))
 
 @spaces.GPU
-def run_segmentation_editor(editor_data, model_choice, min_cell_size):
+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.
@@ -285,21 +309,34 @@ def run_segmentation_editor(editor_data, model_choice, min_cell_size):
         # Run Cellpose segmentation
         masks, flows, styles = model.eval(processed_image_np, diameter=None, channels=[0, 0])
 
-
+        removed_small = 0
+        removed_large = 0
+        
         # Minimum size filtering
 
         recommend_min = rec_min_size(masks) 
         
         if min_cell_size > 0:
-            masks, removed_count = filter_mask_by_size(masks, min_cell_size)
-            filter_msg = f"Removed {removed_count} small objects (< {min_cell_size} pixels).\n"
+            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_count = 0
+            removed_small = 0
             filter_msg=""
             
+        
+
+        # 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=""
+
+        removed_count = removed_small + removed_large
+        
         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:
@@ -313,6 +350,7 @@ def run_segmentation_editor(editor_data, model_choice, min_cell_size):
 
         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."
@@ -415,6 +453,13 @@ with gr.Blocks(
                     label="Minimum Cell Size (pixels)",
     
                 )
+                max_size_slider1 = gr.Slider(
+                    minimum=0,
+                    maximum=500,
+                    value=500,
+                    step=10,
+                    label="Maximum Cell Size (pixels)",
+                )
             
                 segment_btn1 = gr.Button("🔬 Run Segmentation", variant="primary", size="lg")
 
@@ -451,7 +496,7 @@ with gr.Blocks(
         # Event handlers
         segment_btn1.click(
             fn=run_segmentation_editor,
-            inputs=[image_editor, model_dropdown1, min_size_slider1],
+            inputs=[image_editor, model_dropdown1, min_size_slider1, max_size_slider1],
             outputs=[cell_count_output1, overlay_output1, info_output1, viability_section1, masks_state, image_state, confluency_output1, min_size_slider1]
         ).then(  # Chain the initial viability assessment after segmentation
             fn=update_viability_realtime,