Update app.py

app.py

@@ -16,21 +16,44 @@ MODEL_OPTIONS = {
 loaded_models = {}
 
 @spaces.GPU
-def segment_and_count(edited_image, model_choice):
-    # Load selected model
+def segment_and_count(edited_input, model_choice):
+    if edited_input is None:
+        return 0, None
+
+    # Handle Gradio ImageEditor output
+    if isinstance(edited_input, dict) and "image" in edited_input:
+        image = edited_input["image"]
+        
+        # Optional: apply cropping manually if edit info exists
+        edit_info = edited_input.get("edit", {})
+        crop_box = edit_info.get("crop")
+        if crop_box:
+            x = int(crop_box["x"])
+            y = int(crop_box["y"])
+            w = int(crop_box["width"])
+            h = int(crop_box["height"])
+            image = image.crop((x, y, x + w, y + h))
+    else:
+        image = edited_input  # fallback, already a PIL Image
+
+    # Convert to NumPy
+    image_np = np.array(image)
+
+    # Validate shape
+    if image_np.size == 0 or len(image_np.shape) < 2:
+        return 0, None
+
+    # 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
 
-    # Convert edited PIL image to numpy
-    image_np = np.array(edited_image)
-
-    # If grayscale, convert to RGB
+    # Ensure image is 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:
@@ -39,7 +62,7 @@ def segment_and_count(edited_image, model_choice):
     # Run Cellpose
     masks, flows, styles = model.eval(image_np, diameter=None, channels=[0, 0])
 
-    # Count unique cells
+    # Count cells
     cell_count = len(np.unique(masks)) - 1
 
     # Overlay visualization
@@ -57,6 +80,7 @@ def segment_and_count(edited_image, model_choice):
 
     return cell_count, overlay_image
 
+
 # Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("## 🧪 Cell Counter with Cellpose + ImageEditor")