Update app.py

app.py

@@ -129,9 +129,32 @@ 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(image, model_choice):
+def run_segmentation(image, model_choice,min_cell_size):
     image_np = np.array(image)
     image_np = safe_resize(image_np)
     
@@ -159,6 +182,13 @@ def run_segmentation(image, 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)
 
@@ -261,6 +291,14 @@ 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():
@@ -296,7 +334,7 @@ with gr.Blocks(
         # Event handlers
         segment_btn1.click(
             fn=run_segmentation,
-            inputs=[img_input, model_dropdown1],
+            inputs=[img_input, 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,