Upload cellpose_viability_app.py

cellpose_viability_app.py

@@ -0,0 +1,539 @@
+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 = {}
+
+# Global variables to store segmentation results for real-time viability assessment
+current_masks = None
+current_image = None
+current_region_coords = None
+
+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
+    
+    if isinstance(editor_data, dict):
+        background = editor_data.get('background')
+        layers = editor_data.get('layers', [])
+        
+        if background is None:
+            return None, None
+            
+        background_np = np.array(background)
+        
+        if layers and len(layers) > 0:
+            selection_layer = layers[0]
+            selection_np = np.array(selection_layer)
+            
+            if len(selection_np.shape) == 3:
+                if selection_np.shape[2] == 4:  # RGBA
+                    mask = selection_np[:, :, 3] > 0
+                else:  # RGB
+                    mask = np.any(selection_np > 0, axis=2)
+            else:
+                mask = selection_np > 0
+            
+            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()
+                
+                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)
+                
+                region = background_np[y_min:y_max+1, x_min:x_max+1]
+                return region, (x_min, y_min, x_max, y_max)
+        
+        return background_np, None
+    
+    else:
+        if hasattr(editor_data, 'size'):
+            image_np = np.array(editor_data)
+            return image_np, None
+        else:
+            return None, None
+
+def calculate_blue_intensity(image_np, masks):
+    """Calculate blue intensity for each cell"""
+    if len(image_np.shape) == 2:
+        # Convert grayscale to RGB
+        image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
+    elif len(image_np.shape) == 3 and image_np.shape[2] == 4:
+        # Handle RGBA
+        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
+    
+    # Calculate blue intensity for each cell
+    cell_blue_intensities = {}
+    unique_cells = np.unique(masks)
+    
+    for cell_id in unique_cells:
+        if cell_id == 0:  # Skip background
+            continue
+            
+        cell_mask = masks == cell_id
+        cell_pixels = image_np[cell_mask]
+        
+        if len(cell_pixels) > 0:
+            # Calculate mean blue intensity (blue channel)
+            blue_intensity = np.mean(cell_pixels[:, 2])  # Blue channel is index 2
+            cell_blue_intensities[cell_id] = blue_intensity
+    
+    return cell_blue_intensities
+
+def create_viability_overlay(image_np, masks, blue_threshold):
+    """Create overlay showing live/dead cells based on blue threshold"""
+    
+    if masks is None or image_np is None:
+        return None, 0, 0
+    
+    # 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:
+        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
+    
+    # Calculate blue intensities
+    cell_blue_intensities = calculate_blue_intensity(image_np, masks)
+    
+    # Create overlay - start with original image
+    overlay = image_np.copy().astype(np.float32)
+    
+    live_count = 0
+    dead_count = 0
+    
+    # Color cells based on viability
+    for cell_id, blue_intensity in cell_blue_intensities.items():
+        cell_mask = masks == cell_id
+        
+        if blue_intensity > blue_threshold:
+            # Dead cell - color in red
+            overlay[cell_mask] = [255, 0, 0]  # Red
+            dead_count += 1
+        else:
+            # Live cell - color in green
+            overlay[cell_mask] = [0, 255, 0]  # Green
+            live_count += 1
+    
+    # Blend with original image
+    alpha = 0.6  # Increased alpha for better visibility
+    final_overlay = (1 - alpha) * image_np.astype(np.float32) + alpha * overlay
+    final_overlay = np.clip(final_overlay, 0, 255).astype(np.uint8)
+    
+    return Image.fromarray(final_overlay), live_count, dead_count
+
+@spaces.GPU
+def segment_cells(editor_data, model_choice, crop_coords=None):
+    """
+    First stage: Segment cells and store results for viability assessment
+    """
+    global current_masks, current_image, current_region_coords
+    
+    try:
+        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, None, f"No image provided. Debug info:\n{debug_info}", gr.update(visible=False)
+        
+        # Handle manual crop coordinates
+        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
+                    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
+        
+        # Process image format
+        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:
+            region_np = cv2.cvtColor(region_np, cv2.COLOR_RGBA2RGB)
+        
+        # Run Cellpose segmentation
+        masks, flows, styles = model.eval(region_np, diameter=None, channels=[0, 0])
+        
+        # Store results globally for viability assessment
+        current_masks = masks
+        current_image = region_np
+        current_region_coords = region_coords
+        
+        # Count unique cells
+        cell_count = len(np.unique(masks)) - 1
+        
+        # Create basic segmentation overlay
+        overlay = region_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]  # Background stays black
+            
+            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)
+        
+        info_msg = f"Segmentation complete! Found {cell_count} cells.\n"
+        if region_coords:
+            info_msg += f"Region: {region_coords}\n"
+        info_msg += f"Now use the blue threshold slider below for viability assessment."
+        
+        # Make viability section visible
+        viability_section_update = gr.update(visible=True)
+        
+        return cell_count, overlay_image, info_msg, viability_section_update
+        
+    except Exception as e:
+        current_masks = None
+        current_image = None
+        current_region_coords = None
+        return 0, None, f"Error: {str(e)}", gr.update(visible=False)
+
+def update_viability_assessment(blue_threshold):
+    """
+    Real-time viability assessment based on blue threshold
+    """
+    global current_masks, current_image
+    
+    if current_masks is None or current_image is None:
+        return None, 0, 0, "Please run segmentation first"
+    
+    try:
+        # Debug information
+        unique_cells = np.unique(current_masks)
+        total_cells = len(unique_cells) - 1  # Subtract background
+        
+        debug_info = f"Debug - Total cells found: {total_cells}\n"
+        debug_info += f"Image shape: {current_image.shape}\n"
+        debug_info += f"Masks shape: {current_masks.shape}\n"
+        debug_info += f"Unique mask values: {unique_cells[:10]}...\n"  # Show first 10
+        
+        # Create viability overlay
+        viability_overlay, live_count, dead_count = create_viability_overlay(
+            current_image, current_masks, blue_threshold
+        )
+        
+        total_count = live_count + dead_count
+        viability_info = f"Blue threshold: {blue_threshold}\n"
+        viability_info += f"Total cells: {total_count}\nLive (green): {live_count}\nDead (red): {dead_count}\n"
+        
+        if total_count > 0:
+            viability_percentage = (live_count / total_count) * 100
+            viability_info += f"Viability: {viability_percentage:.1f}%\n"
+        
+        viability_info += f"\n{debug_info}"
+        
+        return viability_overlay, live_count, dead_count, viability_info
+        
+    except Exception as e:
+        return None, 0, 0, f"Error in viability assessment: {str(e)}\nMasks available: {current_masks is not None}\nImage available: {current_image is not None}"
+
+@spaces.GPU
+def segment_with_coords(image, model_choice, crop_coords):
+    """Alternative function for manual coordinates input"""
+    global current_masks, current_image, current_region_coords
+    
+    if image is None:
+        return 0, None, "No image provided", gr.update(visible=False)
+    
+    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
+                    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]
+                    current_region_coords = (x_min, y_min, x_max, y_max)
+            except ValueError:
+                pass
+        
+        # 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])
+        
+        # Store results globally
+        current_masks = masks
+        current_image = image_np
+        
+        # Count unique cells
+        cell_count = len(np.unique(masks)) - 1
+        
+        # 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)
+        
+        info_msg = f"Segmentation complete! Found {cell_count} cells.\n"
+        info_msg += f"Now use the blue threshold slider below for viability assessment."
+        
+        return cell_count, overlay_image, info_msg, gr.update(visible=True)
+        
+    except Exception as e:
+        current_masks = None
+        current_image = None
+        return 0, None, f"Error: {str(e)}", gr.update(visible=False)
+
+# Create the Gradio interface
+with gr.Blocks(title="Cell Counter with Real-time Viability Assessment", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# Cell Counter with Real-time Viability Assessment")
+    gr.Markdown("Upload a microscopy image, run segmentation, then use the blue threshold slider for real-time viability assessment.")
+    
+    with gr.Tab("Image Editor (Draw Selection)"):
+        gr.Markdown("### Step 1: Draw selection and run 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("🔬 Run Segmentation", variant="primary", size="lg")
+                
+            with gr.Column():
+                cell_count_output1 = gr.Number(label="Total Cells Detected", precision=0)
+                overlay_output1 = gr.Image(type="pil", label="Segmentation Result")
+                info_output1 = gr.Textbox(label="Processing Info", lines=3)
+        
+        # Viability Assessment Section
+        with gr.Group(visible=False) as viability_section1:
+            gr.Markdown("### Step 2: Real-time Viability Assessment")
+            gr.Markdown("Adjust the blue threshold to classify cells as live (green) or dead (red). Higher values = more blue = more dead.")
+            
+            with gr.Row():
+                with gr.Column():
+                    blue_threshold_slider1 = gr.Slider(
+                        minimum=0,
+                        maximum=255,
+                        value=100,
+                        step=1,
+                        label="Blue Threshold",
+                        info="Higher values = more cells classified as dead"
+                    )
+                    
+                with gr.Column():
+                    live_count_output1 = gr.Number(label="Live Cells (Green)", precision=0)
+                    dead_count_output1 = gr.Number(label="Dead Cells (Red)", precision=0)
+            
+            viability_overlay1 = gr.Image(type="pil", label="Viability Assessment (Green=Live, Red=Dead)")
+            viability_info1 = gr.Textbox(label="Viability Statistics", lines=4)
+        
+        # Event handlers
+        segment_btn1.click(
+            fn=segment_cells,
+            inputs=[image_editor, model_dropdown1],
+            outputs=[cell_count_output1, overlay_output1, info_output1, viability_section1]
+        ).then(  # Chain the viability assessment after segmentation
+            fn=update_viability_assessment,
+            inputs=[blue_threshold_slider1],
+            outputs=[viability_overlay1, live_count_output1, dead_count_output1, viability_info1]
+        )
+        
+        blue_threshold_slider1.change(
+            fn=update_viability_assessment,
+            inputs=[blue_threshold_slider1],
+            outputs=[viability_overlay1, live_count_output1, dead_count_output1, viability_info1]
+        )
+    
+    with gr.Tab("Manual Coordinates"):
+        gr.Markdown("### Step 1: Upload image and run segmentation")
+        
+        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("🔬 Run Segmentation", variant="primary", size="lg")
+                
+            with gr.Column():
+                cell_count_output2 = gr.Number(label="Total Cells Detected", precision=0)
+                overlay_output2 = gr.Image(type="pil", label="Segmentation Result")
+                info_output2 = gr.Textbox(label="Processing Info", lines=3)
+        
+        # Viability Assessment Section
+        with gr.Group(visible=False) as viability_section2:
+            gr.Markdown("### Step 2: Real-time Viability Assessment")
+            gr.Markdown("Adjust the blue threshold to classify cells as live (green) or dead (red). Higher values = more blue = more dead.")
+            
+            with gr.Row():
+                with gr.Column():
+                    blue_threshold_slider2 = gr.Slider(
+                        minimum=0,
+                        maximum=255,
+                        value=100,
+                        step=1,
+                        label="Blue Threshold",
+                        info="Higher values = more cells classified as dead"
+                    )
+                    
+                with gr.Column():
+                    live_count_output2 = gr.Number(label="Live Cells (Green)", precision=0)
+                    dead_count_output2 = gr.Number(label="Dead Cells (Red)", precision=0)
+            
+            viability_overlay2 = gr.Image(type="pil", label="Viability Assessment (Green=Live, Red=Dead)")
+            viability_info2 = gr.Textbox(label="Viability Statistics", lines=4)
+        
+        # Event handlers
+        segment_btn2.click(
+            fn=segment_with_coords,
+            inputs=[image_input, model_dropdown2, coord_input],
+            outputs=[cell_count_output2, overlay_output2, info_output2, viability_section2]
+        ).then(  # Chain the viability assessment after segmentation
+            fn=update_viability_assessment,
+            inputs=[blue_threshold_slider2],
+            outputs=[viability_overlay2, live_count_output2, dead_count_output2, viability_info2]
+        )
+        
+        blue_threshold_slider2.change(
+            fn=update_viability_assessment,
+            inputs=[blue_threshold_slider2],
+            outputs=[viability_overlay2, live_count_output2, dead_count_output2, viability_info2]
+        )
+    
+    # Instructions
+    with gr.Accordion("Instructions", open=False):
+        gr.Markdown("""
+        ### How to use:
+        
+        1. **Upload and Segment**: 
+           - Upload your microscopy image
+           - Select a model (Hemocytometer Model for blood cells)
+           - Draw a selection region or specify coordinates
+           - Click "Run Segmentation"
+        
+        2. **Viability Assessment**:
+           - After segmentation, the viability section will appear
+           - Use the blue threshold slider to classify cells in real-time
+           - Green cells = Live (low blue intensity)
+           - Red cells = Dead (high blue intensity)
+           - Adjust the threshold until the classification looks correct
+        
+        3. **Interpreting Results**:
+           - The app calculates viability percentage automatically
+           - Live cells appear green, dead cells appear red
+           - Statistics show total counts and viability percentage
+        
+        ### Tips:
+        - Start with a threshold around 100 and adjust based on your staining
+        - Dead cells typically have higher blue intensity due to trypan blue uptake
+        - The threshold is applied in real-time without re-running segmentation
+        """)
+
+if __name__ == "__main__":
+    demo.launch()