Update app.py

app.py

@@ -128,31 +128,64 @@ def measure_confluency(masks, image_np):
     cell_pixels = np.count_nonzero(masks)
     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()
+    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, q=25):
+    ids = np.unique(masks)
+    ids = ids[ids > 0]
+    if len(ids) == 0:
+        return 0
+    sizes = np.array([np.count_nonzero(masks == cid) for cid in ids])
+    return int(round(np.percentile(sizes, q)))
+
+
+
 @spaces.GPU
 def run_segmentation(image, model_choice,min_cell_size):
     image_np = np.array(image)
@@ -182,13 +215,41 @@ def run_segmentation(image, model_choice,min_cell_size):
         # 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=""
-            
+        ids = np.unique(masks_raw)
+        ids = ids[ids > 0]
+
+        sizes = np.array([np.count_nonzero(masks_raw == cid) for cid in ids])
+
+        print("num_cells:", len(ids))
+        print("mean:", sizes.mean())
+        print("median:", np.median(sizes))
+        print("p90:", np.percentile(sizes, 90))
+        print("max:", sizes.max())
+        # 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
+
+        if min_used > 0:
+            masks, removed_small = filter_mask_by_size(masks, min_used)
+
+        if max_cell_size > 0:
+            masks, removed_large = filter_mask_by_maxsize(masks, int(max_cell_size))
+
+        
+        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)
 
@@ -197,17 +258,24 @@ def run_segmentation(image, model_choice,min_cell_size):
         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 += f"Now adjust the Blue Threshold for viability assessment."
+        if region_coords:
+            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),
@@ -216,10 +284,22 @@ def run_segmentation(image, model_choice,min_cell_size):
             pack_array(masks),
             pack_array(processed_image_np),
             confluency,
+            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
+        return (
+            0,
+            None,
+            f"Error during segmentation: {str(e)}",
+            gr.update(visible=False),
+            None,
+            None,
+            0.0,
+            gr.update(),
+        )
+
+
 
 
 
@@ -294,11 +374,20 @@ 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(
+                    minimum=0,
+                    maximum=500,
+                    value=500,
+                    step=10,
+                    label="Maximum Cell Size (pixels)",
+                )
+    
+                
                 segment_btn1 = gr.Button("🔬 Run Segmentation", variant="primary", size="lg")
 
             with gr.Column():
@@ -334,8 +423,8 @@ with gr.Blocks(
         # Event handlers
         segment_btn1.click(
             fn=run_segmentation,
-            inputs=[img_input, model_dropdown1,min_size_slider1],
-            outputs=[cell_count_output1, overlay_output1, info_output1, viability_section1, masks_state, image_state, confluency_output1]
+            inputs=[img_input, 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,
             inputs=[blue_threshold1, masks_state, image_state],  # Pass stored state as inputs