Create app.py

app.py

@@ -0,0 +1,458 @@
+import gradio as gr
+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"
+}
+
+loaded_models = {}
+
+def debug_editor_data(editor_data):
+    """Debug function to understand the structure of ImageEditor data"""
+    if editor_data is None:
+        return "No data provided"
+    
+    info = f"Type: {type(editor_data)}\n"
+    
+    if isinstance(editor_data, dict):
+        info += f"Keys: {list(editor_data.keys())}\n"
+        for key, value in editor_data.items():
+            info += f"  {key}: {type(value)}\n"
+            if key == 'layers' and value:
+                info += f"    Layers count: {len(value)}\n"
+                for i, layer in enumerate(value):
+                    info += f"    Layer {i}: {type(layer)}\n"
+                    if hasattr(layer, 'size'):
+                        info += f"      Size: {layer.size}\n"
+                        info += f"      Mode: {layer.mode}\n"
+    else:
+        if hasattr(editor_data, 'size'):
+            info += f"Size: {editor_data.size}\n"
+            info += f"Mode: {editor_data.mode}\n"
+    
+    return info
+
+def extract_region_from_editor(editor_data):
+    """Extract the selected region from ImageEditor data"""
+    if editor_data is None:
+        return None, None
+    
+    # ImageEditor can return different formats depending on version
+    # Let's handle the most common cases
+    
+    if isinstance(editor_data, dict):
+        # Case 1: Dictionary with 'background' 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 drawn layers, try to extract selection
+        if layers and len(layers) > 0:
+            # Get the first layer
+            selection_layer = layers[0]
+            
+            # Convert to numpy array
+            selection_np = np.array(selection_layer)
+            
+            # Find non-transparent/non-black pixels as selection
+            if len(selection_np.shape) == 3:
+                # For RGB, look for non-black pixels
+                if selection_np.shape[2] == 4:  # RGBA
+                    # Use alpha channel
+                    mask = selection_np[:, :, 3] > 0
+                else:  # RGB
+                    # Use non-black pixels
+                    mask = np.any(selection_np > 0, axis=2)
+            else:
+                # Grayscale
+                mask = selection_np > 0
+            
+            # Find bounding box of the selection
+            coords = np.where(mask)
+            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()
+                
+                # Add some padding to ensure we don't get tiny regions
+                pad = 5
+                h, w = background_np.shape[:2]
+                y_min = max(0, y_min - pad)
+                y_max = min(h, y_max + pad)
+                x_min = max(0, x_min - pad)
+                x_max = min(w, x_max + pad)
+                
+                # 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
+    
+    else:
+        # Case 2: Direct PIL Image
+        if hasattr(editor_data, 'size'):  # Check if it's a PIL Image
+            image_np = np.array(editor_data)
+            return image_np, None
+        else:
+            return None, 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]
+
+def classify_cells_by_color(image_np, masks, hue_min, hue_max, sat_min, val_min, val_max):
+    """
+    Classify cells as dead (blue) or alive based on trypan blue color threshold
+    
+    Args:
+        image_np: RGB image array
+        masks: Cellpose segmentation masks
+        hue_min, hue_max: Hue range for blue color (0-179 in OpenCV)
+        sat_min: Minimum saturation (0-255)
+        val_min, val_max: Value (brightness) range (0-255)
+    
+    Returns:
+        dead_count, alive_count, colored_overlay
+    """
+    # Convert RGB to HSV
+    hsv = cv2.cvtColor(image_np, cv2.COLOR_RGB2HSV)
+    
+    # Create color mask for trypan blue
+    if hue_min <= hue_max:
+        color_mask = cv2.inRange(hsv, (hue_min, sat_min, val_min), (hue_max, 255, val_max))
+    else:
+        # Handle hue wrap-around (e.g., 170-10 for red-blue boundary)
+        mask1 = cv2.inRange(hsv, (hue_min, sat_min, val_min), (179, 255, val_max))
+        mask2 = cv2.inRange(hsv, (0, sat_min, val_min), (hue_max, 255, val_max))
+        color_mask = cv2.bitwise_or(mask1, mask2)
+    
+    # Get unique cell IDs (excluding background)
+    cell_ids = np.unique(masks)
+    cell_ids = cell_ids[cell_ids > 0]  # Remove background (0)
+    
+    dead_cells = []
+    alive_cells = []
+    
+    # Classify each cell
+    for cell_id in cell_ids:
+        # Get mask for this specific cell
+        cell_mask = (masks == cell_id)
+        
+        # Calculate overlap with color mask
+        overlap = np.sum(color_mask[cell_mask] > 0)
+        total_pixels = np.sum(cell_mask)
+        
+        # If >30% of cell pixels are blue, classify as dead
+        if overlap / total_pixels > 0.3:
+            dead_cells.append(cell_id)
+        else:
+            alive_cells.append(cell_id)
+    
+    # Create colored overlay
+    overlay = image_np.copy()
+    
+    # Color dead cells red, alive cells green
+    for cell_id in dead_cells:
+        cell_mask = (masks == cell_id)
+        overlay[cell_mask] = [255, 0, 0]  # Red for dead
+    
+    for cell_id in alive_cells:
+        cell_mask = (masks == cell_id)
+        overlay[cell_mask] = [0, 255, 0]  # Green for alive
+    
+    # Blend with original image
+    alpha = 0.4
+    overlay = (1 - alpha) * image_np + alpha * overlay
+    overlay = np.clip(overlay, 0, 255).astype(np.uint8)
+    
+    return len(dead_cells), len(alive_cells), overlay
+
+@spaces.GPU
+def segment_and_count_viability(editor_data, model_choice, hue_min, hue_max, sat_min, val_min, val_max, crop_coords=None):
+    """
+    Segment cells and analyze viability using trypan blue color threshold
+    """
+    try:
+        # Debug info
+        debug_info = debug_editor_data(editor_data)
+        
+        # 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, 0, 0, 0.0, None, f"No image provided. Debug info:\n{debug_info}"
+        
+        # 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])
+        
+        # Classify cells by color
+        dead_count, alive_count, viability_overlay = classify_cells_by_color(
+            region_np, masks, hue_min, hue_max, sat_min, val_min, val_max
+        )
+        
+        total_count = dead_count + alive_count
+        viability_percent = (alive_count / total_count * 100) if total_count > 0 else 0.0
+        
+        # Convert result to PIL Image for output
+        overlay_image = Image.fromarray(viability_overlay)
+        
+        # Create info message
+        info_msg = f"Total cells: {total_count}\nAlive cells: {alive_count}\nDead cells: {dead_count}\nViability: {viability_percent:.1f}%"
+        if region_coords:
+            info_msg += f"\nProcessed region: {region_coords[0]},{region_coords[1]} to {region_coords[2]},{region_coords[3]}"
+        
+        return total_count, alive_count, dead_count, viability_percent, overlay_image, info_msg
+        
+    except Exception as e:
+        return 0, 0, 0, 0.0, None, f"Error occurred: {str(e)}\nDebug info:\n{debug_info if 'debug_info' in locals() else 'No debug info'}"
+
+@spaces.GPU
+def segment_with_coords_viability(image, model_choice, hue_min, hue_max, sat_min, val_min, val_max, crop_coords):
+    """Segment using regular image input with manual coordinate specification and viability analysis"""
+    if image is None:
+        return 0, 0, 0, 0.0, None, "No image provided"
+    
+    try:
+        # 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])
+        
+        # Classify cells by color
+        dead_count, alive_count, viability_overlay = classify_cells_by_color(
+            image_np, masks, hue_min, hue_max, sat_min, val_min, val_max
+        )
+        
+        total_count = dead_count + alive_count
+        viability_percent = (alive_count / total_count * 100) if total_count > 0 else 0.0
+        
+        overlay_image = Image.fromarray(viability_overlay)
+        
+        info_msg = f"Total cells: {total_count}\nAlive cells: {alive_count}\nDead cells: {dead_count}\nViability: {viability_percent:.1f}%"
+        if crop_coords:
+            info_msg += f"\nProcessed with coordinates: {crop_coords}"
+        
+        return total_count, alive_count, dead_count, viability_percent, overlay_image, info_msg
+        
+    except Exception as e:
+        return 0, 0, 0, 0.0, None, f"Error occurred: {str(e)}"
+
+# Create the Gradio interface with tabs for different input methods
+with gr.Blocks(title="Cell Viability Counter with Trypan Blue Analysis") as demo:
+    gr.Markdown("# Cell Viability Counter with Trypan Blue Analysis")
+    gr.Markdown("Upload a trypan blue stained microscopy image and analyze cell viability. Dead cells appear blue, alive cells are unstained.")
+    
+    with gr.Tab("Image Editor (Draw Selection)"):
+        gr.Markdown("Use the drawing tools to select a region of the image for segmentation and viability analysis.")
+        
+        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"
+                )
+                
+                # Color threshold controls
+                gr.Markdown("### Trypan Blue Color Threshold")
+                with gr.Row():
+                    hue_min1 = gr.Slider(0, 179, value=100, step=1, label="Hue Min (Blue)")
+                    hue_max1 = gr.Slider(0, 179, value=130, step=1, label="Hue Max (Blue)")
+                
+                with gr.Row():
+                    sat_min1 = gr.Slider(0, 255, value=50, step=1, label="Saturation Min")
+                    val_min1 = gr.Slider(0, 255, value=50, step=1, label="Value Min")
+                    val_max1 = gr.Slider(0, 255, value=255, step=1, label="Value Max")
+                
+                segment_btn1 = gr.Button("Analyze Viability", variant="primary")
+                
+            with gr.Column():
+                with gr.Row():
+                    total_count_output1 = gr.Number(label="Total Cells")
+                    alive_count_output1 = gr.Number(label="Alive Cells")
+                    dead_count_output1 = gr.Number(label="Dead Cells")
+                
+                viability_output1 = gr.Number(label="Viability (%)", precision=1)
+                overlay_output1 = gr.Image(type="pil", label="Viability Analysis (Green=Alive, Red=Dead)")
+                info_output1 = gr.Textbox(label="Analysis Results")
+        
+        segment_btn1.click(
+            fn=segment_and_count_viability,
+            inputs=[image_editor, model_dropdown1, hue_min1, hue_max1, sat_min1, val_min1, val_max1],
+            outputs=[total_count_output1, alive_count_output1, dead_count_output1, viability_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="Trypan Blue Stained 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"
+                )
+                
+                # Color threshold controls
+                gr.Markdown("### Trypan Blue Color Threshold")
+                with gr.Row():
+                    hue_min2 = gr.Slider(0, 179, value=100, step=1, label="Hue Min (Blue)")
+                    hue_max2 = gr.Slider(0, 179, value=130, step=1, label="Hue Max (Blue)")
+                
+                with gr.Row():
+                    sat_min2 = gr.Slider(0, 255, value=50, step=1, label="Saturation Min")
+                    val_min2 = gr.Slider(0, 255, value=50, step=1, label="Value Min")
+                    val_max2 = gr.Slider(0, 255, value=255, step=1, label="Value Max")
+                
+                segment_btn2 = gr.Button("Analyze Viability", variant="primary")
+                
+            with gr.Column():
+                with gr.Row():
+                    total_count_output2 = gr.Number(label="Total Cells")
+                    alive_count_output2 = gr.Number(label="Alive Cells")
+                    dead_count_output2 = gr.Number(label="Dead Cells")
+                
+                viability_output2 = gr.Number(label="Viability (%)", precision=1)
+                overlay_output2 = gr.Image(type="pil", label="Viability Analysis (Green=Alive, Red=Dead)")
+                info_output2 = gr.Textbox(label="Analysis Results")
+        
+        segment_btn2.click(
+            fn=segment_with_coords_viability,
+            inputs=[image_input, model_dropdown2, hue_min2, hue_max2, sat_min2, val_min2, val_max2, coord_input],
+            outputs=[total_count_output2, alive_count_output2, dead_count_output2, viability_output2, overlay_output2, info_output2]
+        )
+    
+    # Add usage instructions
+    with gr.Tab("Usage Instructions"):
+        gr.Markdown("""
+        ## How to Use This Tool
+        
+        ### 1. Image Preparation
+        - Upload a trypan blue stained microscopy image
+        - Dead cells should appear blue, alive cells should be unstained (clear)
+        
+        ### 2. Color Threshold Adjustment
+        - **Hue Min/Max**: Adjust to capture the blue color range of trypan blue (typically 100-130)
+        - **Saturation Min**: Minimum color saturation to distinguish from gray areas (typically 50+)
+        - **Value Min/Max**: Brightness range for blue detection (typically 50-255)
+        
+        ### 3. Default Settings
+        - The default settings work well for most trypan blue images
+        - Fine-tune based on your specific staining conditions and lighting
+        
+        ### 4. Output
+        - **Green overlay**: Live cells
+        - **Red overlay**: Dead cells
+        - **Viability %**: (Alive cells / Total cells) × 100
+        
+        ### 5. Tips
+        - Use consistent lighting and staining protocols for best results
+        - Adjust thresholds if cells are misclassified
+        - The tool works best with clear contrast between stained and unstained cells
+        """)
+
+if __name__ == "__main__":
+    demo.launch()