Update app.py

app.py

@@ -3,78 +3,88 @@ 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
 
+# Hugging Face model repo
 HF_REPO_ID = "myang4218/cellposemodel"
 
+# Model filename options
 MODEL_OPTIONS = {
     "Hemocytometer Model": "hemocytometermodel.npy",
     "General Model": "generalmodel.npy"
 }
 
+# Cache loaded models
 loaded_models = {}
 
 @spaces.GPU
-def segment_and_count(image, model_choice):
-    
+def segment_and_count(image, box, model_choice):
+    # Download and 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 PIL Image to numpy array
-    image_np = np.array(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:
-        # Handle RGBA images
-        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
-    
-    # Run Cellpose
-    masks, flows, styles = model.eval(image_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 = image_np.copy().astype(np.float32)
-    
-    # Create colored mask overlay
+
+    # Get ROI coordinates from selection box
+    if box is None:
+        # No box selected – process entire image
+        roi_pil = image
+        x, y = 0, 0
+    else:
+        x, y = int(box["x"]), int(box["y"])
+        w, h = int(box["width"]), int(box["height"])
+        roi_pil = image.crop((x, y, x + w, y + h))
+
+    # Convert cropped region to NumPy
+    roi_np = np.array(roi_pil)
+
+    # Convert to RGB if needed
+    if len(roi_np.shape) == 2:
+        roi_np = cv2.cvtColor(roi_np, cv2.COLOR_GRAY2RGB)
+    elif roi_np.shape[2] == 4:
+        roi_np = cv2.cvtColor(roi_np, cv2.COLOR_RGBA2RGB)
+
+    # Run Cellpose on cropped ROI
+    masks, flows, styles = model.eval(roi_np, diameter=None, channels=[0, 0])
+    cell_count = len(np.unique(masks)) - 1  # Exclude background
+
+    # Create colored overlay for ROI
+    overlay_roi = roi_np.copy().astype(np.float32)
     if masks.max() > 0:
-        # Generate random colors for each cell
-        np.random.seed(42)  # For reproducible colors
+        np.random.seed(42)
         colors = np.random.randint(0, 255, size=(masks.max() + 1, 3))
-        colors[0] = [0, 0, 0]  # Background stays black
-        
-        # Create colored overlay
+        colors[0] = [0, 0, 0]
         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)
-    
+        overlay_roi = (1 - alpha) * overlay_roi + alpha * colored_mask
+
+    overlay_roi = np.clip(overlay_roi, 0, 255).astype(np.uint8)
+
+    # Prepare final overlay image
+    full_overlay = np.array(image).astype(np.float32)
+    if full_overlay.shape[2] == 4:
+        full_overlay = cv2.cvtColor(full_overlay, cv2.COLOR_RGBA2RGB)
+
+    # Paste the overlay ROI back onto the full image
+    h, w = overlay_roi.shape[:2]
+    full_overlay[y:y + h, x:x + w] = overlay_roi
+
+    # Convert back to PIL for output
+    overlay_image = Image.fromarray(full_overlay.astype(np.uint8))
+
     return cell_count, overlay_image
 
+
 # Gradio interface
 demo = gr.Interface(
     fn=segment_and_count,
     inputs=[
-        gr.Image(type="pil", label="Microscopy Image"),
+        gr.Image(type="pil", label="Microscopy Image", tool="select"),
         gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), label="Select Model", value="Hemocytometer Model")
     ],
     outputs=[
@@ -82,8 +92,8 @@ demo = gr.Interface(
         gr.Image(type="pil", label="Segmented Overlay")
     ],
     title="Cell Counter with Cellpose",
-    description="Upload a microscopy image and select a model to count the number of cells using Cellpose segmentation."
+    description="Upload a microscopy image and optionally select a region to segment cells using Cellpose."
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()