Update app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import spaces
 from cellpose import models
 import numpy as np
 import cv2
@@ -8,52 +9,68 @@ from PIL import Image
 import io 
 from huggingface_hub import hf_hub_download
 
+# Download model at startup
 HF_REPO_ID = "myang4218/cellposemodel" 
 HF_MODEL_FILENAME = "cellposemodel.npy"
-
-
-
 model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=HF_MODEL_FILENAME)
 print(f"Model downloaded to: {model_path}")
 
-# Load pretrained Cellpose model (nuclei or cyto depending on your data)
-model = models.CellposeModel(gpu=True, pretrained_model=model_path)
-
-
+@spaces.GPU
 def segment_and_count(image):
+    # Initialize model inside the GPU-decorated function
+    model = models.CellposeModel(gpu=True, pretrained_model=model_path)
+    
     # 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
-
-    # Overlay mask for display
-
-    overlay = image_np.copy()
-    alpha = 0.3  # Transparency factor (0.0 - 1.0, where 1.0 is fully opaque)
-
-    # Blend the red color with the original image where mask > 0
-    overlay[masks > 0] = (1 - alpha) * overlay[masks > 0] + alpha * np.array([255, 0, 0])
-
+    
+    # Create better overlay visualization
+    overlay = image_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)
-
+    
     return cell_count, overlay_image
 
 # Gradio interface
 demo = gr.Interface(
     fn=segment_and_count,
     inputs=gr.Image(type="pil"),
-    outputs=[gr.Number(label="Number of Cells"), gr.Image(type="pil", label="Segmented Overlay")],
+    outputs=[
+        gr.Number(label="Number of Cells"), 
+        gr.Image(type="pil", label="Segmented Overlay")
+    ],
     title="Cell Counter with Cellpose",
-    description="Upload a microscopy image and count the number of cells using Cellpose segmentation."
+    description="Upload a microscopy image and count the number of cells using Cellpose segmentation.",
 )
 
 if __name__ == "__main__":