Update app.py

app.py

@@ -216,9 +216,34 @@ 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()
+    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
+
+    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
+    
+        
 
 @spaces.GPU
-def run_segmentation_editor(editor_data, model_choice):
+def run_segmentation_editor(editor_data, model_choice, min_cell_size):
     """
     Runs cell segmentation using ImageEditor data.
     Returns initial segmentation overlay, counts, confluency, and also masks/image for state.
@@ -252,6 +277,13 @@ def run_segmentation_editor(editor_data, model_choice):
         # Run Cellpose segmentation
         masks, flows, styles = model.eval(processed_image_np, diameter=None, channels=[0, 0])
 
+        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"
+        else:
+            removed_count = 0
+            filter_msg=""
+            
         cell_count = len(np.unique(masks)) - 1
         confluency = measure_confluency(masks, processed_image_np)
 
@@ -357,6 +389,15 @@ with gr.Blocks(
                     label="Select Model",
                     value="Hemocytometer Model"
                 )
+                min_size_slider1 = gr.slider(
+                    minimum=0,
+                    maximum=500,
+                    value=50,
+                    step=10,
+                    label="Minimum Cell Size (pixels)",
+    
+                )
+            
                 segment_btn1 = gr.Button("🔬 Run Segmentation", variant="primary", size="lg")
 
             with gr.Column():
@@ -392,7 +433,7 @@ with gr.Blocks(
         # Event handlers
         segment_btn1.click(
             fn=run_segmentation_editor,
-            inputs=[image_editor, model_dropdown1],
+            inputs=[image_editor, model_dropdown1, min_size_slider1],
             outputs=[cell_count_output1, overlay_output1, info_output1, viability_section1, masks_state, image_state, confluency_output1]
         ).then(  # Chain the initial viability assessment after segmentation
             fn=update_viability_realtime,