Update app.py

app.py

@@ -17,34 +17,39 @@ loaded_models = {}
 
 @spaces.GPU
 def segment_and_count(image_with_crop, model_choice):
-    # Extract image and optional crop region
+    # Extract PIL image and crop box
     image = image_with_crop["image"]
-    crop_coords = image_with_crop.get("crop")
+    crop_box = image_with_crop["crop"]  # Format: [x1, y1, x2, y2]
 
-    if crop_coords:
-        # Crop the image if a region was selected
-        x0, y0, x1, y1 = map(int, crop_coords)
-        image = image.crop((x0, y0, x1, y1))
+    # 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
     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
 
-    image_np = np.array(image)
-
-    # Convert grayscale or RGBA to RGB
-    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)
-
     # Run Cellpose
     masks, flows, styles = model.eval(image_np, diameter=None, channels=[0, 0])
-    cell_count = len(np.unique(masks)) - 1  # exclude background
+
+    # Count cells
+    cell_count = len(np.unique(masks)) - 1  # background = 0
 
     # Overlay visualization
     overlay = image_np.copy().astype(np.float32)
@@ -60,13 +65,34 @@ def segment_and_count(image_with_crop, model_choice):
 
     return cell_count, overlay_image
 
-# Gradio Blocks Interface
+# Build Gradio app using Blocks
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧫 Cell Counter with Cellpose")
+    gr.Markdown("## 🧪 Cell Counter with Cellpose")
+
+    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")
 
     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")
+        count_output = gr.Number(label="Number of Cells")
+        overlay_output = gr.Image(type="pil", label="Segmented Overlay")
+
+    run_button.click(
+        fn=segment_and_count,
+        inputs=[image_input, model_dropdown],
+        outputs=[count_output, overlay_output]
+    )
 
-    run_button = gr.Button("
+if __name__ == "__main__":
+    demo.launch()