Update app.py

app.py

@@ -16,40 +16,31 @@ MODEL_OPTIONS = {
 loaded_models = {}
 
 @spaces.GPU
-def segment_and_count(image_with_crop, model_choice):
-    # Extract PIL image and crop box
-    image = image_with_crop["image"]
-    crop_box = image_with_crop["crop"]  # Format: [x1, y1, x2, y2]
-
-    # Crop the image if a box was drawn
-    if crop_box is not None:
-        x1, y1, x2, y2 = map(int, crop_box)
-        image = image.crop((x1, y1, x2, y2))
-
-    # Convert to NumPy array
-    image_np = np.array(image)
-
-    # Ensure RGB format
-    if len(image_np.shape) == 2:
-        image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
-    elif image_np.shape[2] == 4:
-        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
-
-    # Load model
+def segment_and_count(edited_image, model_choice):
+    # Load selected model
     model_filename = MODEL_OPTIONS[model_choice]
     model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=model_filename)
-
+    
     if model_filename in loaded_models:
         model = loaded_models[model_filename]
     else:
         model = models.CellposeModel(gpu=True, pretrained_model=model_path)
         loaded_models[model_filename] = model
 
+    # Convert edited PIL image to numpy
+    image_np = np.array(edited_image)
+
+    # If grayscale, convert to RGB
+    if len(image_np.shape) == 2:
+        image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
+    elif len(image_np.shape) == 3 and image_np.shape[2] == 4:
+        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
+
     # Run Cellpose
     masks, flows, styles = model.eval(image_np, diameter=None, channels=[0, 0])
 
-    # Count cells
-    cell_count = len(np.unique(masks)) - 1  # background = 0
+    # Count unique cells
+    cell_count = len(np.unique(masks)) - 1
 
     # Overlay visualization
     overlay = image_np.copy().astype(np.float32)
@@ -58,41 +49,31 @@ def segment_and_count(image_with_crop, model_choice):
         colors = np.random.randint(0, 255, size=(masks.max() + 1, 3))
         colors[0] = [0, 0, 0]
         colored_mask = colors[masks]
-        overlay = (1 - 0.4) * overlay + 0.4 * colored_mask
+        alpha = 0.4
+        overlay = (1 - alpha) * overlay + alpha * colored_mask
 
     overlay = np.clip(overlay, 0, 255).astype(np.uint8)
     overlay_image = Image.fromarray(overlay)
 
     return cell_count, overlay_image
 
-# Build Gradio app using Blocks
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧪 Cell Counter with Cellpose")
+    gr.Markdown("## 🧪 Cell Counter with Cellpose + ImageEditor")
+    gr.Markdown("Upload a microscopy image, draw/crop a region using the editor, then select a model to count cells in that region.")
 
     with gr.Row():
-        image_input = gr.Image(
-            type="pil",
-            label="Microscopy Image",
-            tool="select"
-        )
-        model_dropdown = gr.Dropdown(
-            choices=list(MODEL_OPTIONS.keys()),
-            label="Select Model",
-            value="Hemocytometer Model"
-        )
-
-    run_button = gr.Button("Run Segmentation")
+        image_editor = gr.ImageEditor(label="Draw or Crop Region", type="pil")
+        model_selector = gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), value="Hemocytometer Model", label="Select Model")
 
     with gr.Row():
         count_output = gr.Number(label="Number of Cells")
-        overlay_output = gr.Image(type="pil", label="Segmented Overlay")
+        overlay_output = gr.Image(label="Segmented Overlay")
 
-    run_button.click(
-        fn=segment_and_count,
-        inputs=[image_input, model_dropdown],
-        outputs=[count_output, overlay_output]
-    )
+    image_editor.change(fn=segment_and_count, inputs=[image_editor, model_selector], outputs=[count_output, overlay_output])
+    model_selector.change(fn=segment_and_count, inputs=[image_editor, model_selector], outputs=[count_output, overlay_output])
 
 if __name__ == "__main__":
     demo.launch()
 
+