Update app.py

app.py

@@ -3,11 +3,13 @@ import spaces
 from cellpose import models
 import numpy as np
 import cv2
+import matplotlib.pyplot as plt
+import tempfile
 from PIL import Image
+import io
 from huggingface_hub import hf_hub_download
 
 HF_REPO_ID = "myang4218/cellposemodel"
-
 MODEL_OPTIONS = {
     "Hemocytometer Model": "hemocytometermodel.npy",
     "General Model": "generalmodel.npy"
@@ -15,88 +17,278 @@ MODEL_OPTIONS = {
 
 loaded_models = {}
 
-@spaces.GPU
-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"]
+def extract_region_from_editor(editor_data):
+    """Extract the selected region from ImageEditor data"""
+    if editor_data is None:
+        return None, None
+    
+    # Get the background image and layers
+    background = editor_data.get('background')
+    layers = editor_data.get('layers', [])
+    
+    if background is None:
+        return None, None
+    
+    # Convert background to numpy array
+    background_np = np.array(background)
+    
+    # If there are layers (selections), process them
+    if layers:
+        # For simplicity, we'll use the first layer as the selection mask
+        # In a more complex implementation, you might want to combine multiple layers
+        selection_layer = layers[0]
         
-        # 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)
+        # The layer contains the selection information
+        # This is a simplified approach - you might need to adjust based on your specific needs
+        selection_image = selection_layer.get('image')
+        if selection_image is not None:
+            # Convert selection to mask
+            selection_np = np.array(selection_image)
+            
+            # Create a binary mask from the selection
+            if len(selection_np.shape) == 3:
+                # If it's an RGB image, convert to grayscale for mask
+                mask = cv2.cvtColor(selection_np, cv2.COLOR_RGB2GRAY)
+            else:
+                mask = selection_np
+            
+            # Find bounding box of the selection
+            coords = np.where(mask > 0)
+            if len(coords[0]) > 0:
+                y_min, y_max = coords[0].min(), coords[0].max()
+                x_min, x_max = coords[1].min(), coords[1].max()
+                
+                # Extract the region
+                region = background_np[y_min:y_max+1, x_min:x_max+1]
+                return region, (x_min, y_min, x_max, y_max)
+    
+    # If no selection, return the full image
+    return background_np, None
 
-    # Validate shape
-    if image_np.size == 0 or len(image_np.shape) < 2:
-        return 0, None
+def crop_image_with_coords(image_np, coords):
+    """Crop image using provided coordinates"""
+    if coords is None:
+        return image_np
+    
+    x_min, y_min, x_max, y_max = coords
+    return image_np[y_min:y_max+1, x_min:x_max+1]
 
-    # Load model
+@spaces.GPU
+def segment_and_count(editor_data, model_choice, crop_coords=None):
+    """
+    Segment and count cells in the selected region
+    
+    Args:
+        editor_data: Data from ImageEditor component
+        model_choice: Selected model for segmentation
+        crop_coords: Optional manual crop coordinates as "x_min,y_min,x_max,y_max"
+    """
+    # Load the 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
+    
+    # Extract region from editor
+    region_np, region_coords = extract_region_from_editor(editor_data)
+    
+    if region_np is None:
+        return 0, None, "No image provided"
+    
+    # If manual crop coordinates are provided, use them instead
+    if crop_coords and crop_coords.strip():
+        try:
+            coords = [int(x.strip()) for x in crop_coords.split(',')]
+            if len(coords) == 4:
+                x_min, y_min, x_max, y_max = coords
+                # Ensure coordinates are within image bounds
+                h, w = region_np.shape[:2]
+                x_min = max(0, min(x_min, w-1))
+                y_min = max(0, min(y_min, h-1))
+                x_max = max(x_min+1, min(x_max, w))
+                y_max = max(y_min+1, min(y_max, h))
+                
+                region_np = region_np[y_min:y_max, x_min:x_max]
+                region_coords = (x_min, y_min, x_max, y_max)
+        except ValueError:
+            pass  # Invalid coordinates, continue with current region
+    
+    # If grayscale, convert to RGB
+    if len(region_np.shape) == 2:
+        region_np = cv2.cvtColor(region_np, cv2.COLOR_GRAY2RGB)
+    elif len(region_np.shape) == 3 and region_np.shape[2] == 4:
+        # Handle RGBA images
+        region_np = cv2.cvtColor(region_np, cv2.COLOR_RGBA2RGB)
+    
+    # Run Cellpose on the selected region
+    masks, flows, styles = model.eval(region_np, diameter=None, channels=[0, 0])
+    
+    # Count unique cells
+    cell_count = len(np.unique(masks)) - 1  # subtract 1 for background
+    
+    # Create better overlay visualization
+    overlay = region_np.copy().astype(np.float32)
+    
+    # Create colored mask overlay
+    if masks.max() > 0:
+        # Generate random colors for each cell
+        np.random.seed(42)  # For reproducible colors
+        colors = np.random.randint(0, 255, size=(masks.max() + 1, 3))
+        colors[0] = [0, 0, 0]  # Background stays black
+        
+        # Create colored overlay
+        colored_mask = colors[masks]
+        
+        # Blend with original image
+        alpha = 0.4
+        overlay = (1 - alpha) * overlay + alpha * colored_mask
+    
+    # Ensure values are in valid range and convert to uint8
+    overlay = np.clip(overlay, 0, 255).astype(np.uint8)
+    
+    # Convert result to PIL Image for output
+    overlay_image = Image.fromarray(overlay)
+    
+    # Create info message
+    if region_coords:
+        info_msg = f"Processed region: {region_coords[0]},{region_coords[1]} to {region_coords[2]},{region_coords[3]}"
+    else:
+        info_msg = "Processed entire image"
+    
+    return cell_count, overlay_image, info_msg
 
-    # Ensure image is RGB
+# Alternative function for simple image input with coordinate textbox
+@spaces.GPU
+def segment_with_coords(image, model_choice, crop_coords):
+    """Segment using regular image input with manual coordinate specification"""
+    if image is None:
+        return 0, None, "No image provided"
+    
+    # Load the 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 PIL Image to numpy array
+    image_np = np.array(image)
+    
+    # Apply crop if coordinates are provided
+    if crop_coords and crop_coords.strip():
+        try:
+            coords = [int(x.strip()) for x in crop_coords.split(',')]
+            if len(coords) == 4:
+                x_min, y_min, x_max, y_max = coords
+                # Ensure coordinates are within image bounds
+                h, w = image_np.shape[:2]
+                x_min = max(0, min(x_min, w-1))
+                y_min = max(0, min(y_min, h-1))
+                x_max = max(x_min+1, min(x_max, w))
+                y_max = max(y_min+1, min(y_max, h))
+                
+                image_np = image_np[y_min:y_max, x_min:x_max]
+        except ValueError:
+            pass  # Invalid coordinates, use full image
+    
+    # Process image format
     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
+    
+    # Count unique cells
     cell_count = len(np.unique(masks)) - 1
-
-    # Overlay visualization
+    
+    # Create overlay
     overlay = image_np.copy().astype(np.float32)
+    
     if masks.max() > 0:
         np.random.seed(42)
         colors = np.random.randint(0, 255, size=(masks.max() + 1, 3))
         colors[0] = [0, 0, 0]
+        
         colored_mask = colors[masks]
         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, f"Processed with coordinates: {crop_coords}" if crop_coords else "Processed entire image"
 
-    return cell_count, overlay_image
-
-
-
-# Gradio UI
-with gr.Blocks() as demo:
-    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_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(label="Segmented Overlay")
-
-    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])
+# Create the Gradio interface with tabs for different input methods
+with gr.Blocks(title="Cell Counter with Region Selection") as demo:
+    gr.Markdown("# Cell Counter with Cellpose - Region Selection")
+    gr.Markdown("Upload a microscopy image and select a region to count cells using Cellpose segmentation.")
+    
+    with gr.Tab("Image Editor (Draw Selection)"):
+        gr.Markdown("Use the drawing tools to select a region of the image for segmentation.")
+        
+        with gr.Row():
+            with gr.Column():
+                image_editor = gr.ImageEditor(
+                    label="Draw selection on image",
+                    type="pil",
+                    brush=gr.Brush(colors=["#ff0000"], color_mode="fixed", default_size=20),
+                    eraser=gr.Eraser(default_size=20)
+                )
+                model_dropdown1 = gr.Dropdown(
+                    choices=list(MODEL_OPTIONS.keys()),
+                    label="Select Model",
+                    value="Hemocytometer Model"
+                )
+                segment_btn1 = gr.Button("Segment Selected Region", variant="primary")
+                
+            with gr.Column():
+                cell_count_output1 = gr.Number(label="Number of Cells")
+                overlay_output1 = gr.Image(type="pil", label="Segmented Overlay")
+                info_output1 = gr.Textbox(label="Processing Info")
+        
+        segment_btn1.click(
+            fn=segment_and_count,
+            inputs=[image_editor, model_dropdown1],
+            outputs=[cell_count_output1, overlay_output1, info_output1]
+        )
+    
+    with gr.Tab("Manual Coordinates"):
+        gr.Markdown("Upload an image and specify coordinates manually (format: x_min,y_min,x_max,y_max)")
+        
+        with gr.Row():
+            with gr.Column():
+                image_input = gr.Image(type="pil", label="Microscopy Image")
+                model_dropdown2 = gr.Dropdown(
+                    choices=list(MODEL_OPTIONS.keys()),
+                    label="Select Model",
+                    value="Hemocytometer Model"
+                )
+                coord_input = gr.Textbox(
+                    label="Crop Coordinates (optional)",
+                    placeholder="e.g., 100,100,400,400",
+                    info="Format: x_min,y_min,x_max,y_max"
+                )
+                segment_btn2 = gr.Button("Segment Region", variant="primary")
+                
+            with gr.Column():
+                cell_count_output2 = gr.Number(label="Number of Cells")
+                overlay_output2 = gr.Image(type="pil", label="Segmented Overlay")
+                info_output2 = gr.Textbox(label="Processing Info")
+        
+        segment_btn2.click(
+            fn=segment_with_coords,
+            inputs=[image_input, model_dropdown2, coord_input],
+            outputs=[cell_count_output2, overlay_output2, info_output2]
+        )
 
 if __name__ == "__main__":
     demo.launch()