Update app.py

app.py

@@ -80,4 +80,501 @@ def extract_region_from_editor(editor_data):
                 
                 pad = 5
                 h, w = background_np.shape[:2]
-                y_min = max(0, y_min - pad)
+                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]
+                print(f"DEBUG: extract_region_from_editor - Extracted region shape: {region.shape}")
+                return region, (x_min, y_min, x_max, y_max)
+        print("DEBUG: extract_region_from_editor - No selection or empty selection, returning full background.")
+        return background_np, None
+    else: # Assuming it's a direct PIL Image if not dict
+        if hasattr(editor_data, 'size'):
+            image_np = np.array(editor_data)
+            print(f"DEBUG: extract_region_from_editor - Direct PIL image, shape: {image_np.shape}")
+            return image_np, None
+        else:
+            print("DEBUG: extract_region_from_editor - editor_data is not dict and no size attribute, returning None.")
+            return None, None
+
+def calculate_blue_intensity(image_np, masks):
+    """Calculate blue intensity for each cell"""
+    # Ensure image is in correct format
+    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)
+    
+    print(f"DEBUG: Found {len(unique_cells)} unique mask values: {unique_cells}")
+    print(f"DEBUG: Image shape: {image_np.shape}, dtype: {image_np.dtype}")
+    print(f"DEBUG: Masks shape: {masks.shape}, dtype: {masks.dtype}")
+    
+    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 is index 2)
+            blue_intensity = np.mean(cell_pixels[:, 2])  # Blue channel
+            cell_blue_intensities[cell_id] = blue_intensity
+            # print(f"DEBUG: Cell {cell_id}: {len(cell_pixels)} pixels, blue intensity: {blue_intensity:.2f}") # Too verbose, uncomment if needed
+    print(f"DEBUG: Total cells with blue intensity: {len(cell_blue_intensities)}")
+    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:
+        print("DEBUG: create_viability_overlay - masks or image_np is None, returning None.")
+        return None, 0, 0
+    
+    print(f"DEBUG: create_viability_overlay called with threshold: {blue_threshold}")
+    print(f"DEBUG: Image shape: {image_np.shape}, Masks shape: {masks.shape}, Masks max: {masks.max()}")
+    
+    # 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)
+    
+    print(f"DEBUG: Got {len(cell_blue_intensities)} cells with blue intensities")
+    
+    # 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
+        
+        # print(f"DEBUG: Processing cell {cell_id}, blue intensity: {blue_intensity:.2f}, threshold: {blue_threshold}") # Too verbose
+        
+        if blue_intensity > blue_threshold:
+            # Dead cell - color in red
+            overlay[cell_mask] = [255, 0, 0]  # Red
+            dead_count += 1
+            # print(f"DEBUG: Cell {cell_id} classified as DEAD") # Too verbose
+        else:
+            # Live cell - color in green
+            overlay[cell_mask] = [0, 255, 0]  # Green
+            live_count += 1
+            # print(f"DEBUG: Cell {cell_id} classified as LIVE") # Too verbose
+    print(f"DEBUG: Final counts - Live: {live_count}, Dead: {dead_count}")
+    
+    # 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 removed from here, handled within extract_region_from_editor implicitly by ImageEditor
+    """
+    First stage: Segment cells and return results for viability assessment
+    """
+    global current_region_coords # Keep global for info display if needed later
+    
+    try:
+        print("DEBUG: --- Starting segment_cells ---")
+        debug_info = debug_editor_data(editor_data)
+        print(f"DEBUG: Editor data info:\n{debug_info}")
+        
+        # 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:
+            print("DEBUG: region_np is None, returning error and hiding viability section.")
+            current_region_coords = None
+            # Also return None for masks and image for the gr.State components
+            return 0, None, f"No image provided. Debug info:\n{debug_info}", gr.update(visible=False), None, None
+        
+        print(f"DEBUG: Image to segment shape (after extraction/cropping): {region_np.shape}, dtype: {region_np.dtype}")
+        
+        # Process image format
+        if len(region_np.shape) == 2:
+            region_np_processed = cv2.cvtColor(region_np, cv2.COLOR_GRAY2RGB)
+        elif len(region_np.shape) == 3 and region_np.shape[2] == 4:
+            region_np_processed = cv2.cvtColor(region_np, cv2.COLOR_RGBA2RGB)
+        else:
+            region_np_processed = region_np # Already RGB or suitable
+
+        print(f"DEBUG: Image before model.eval shape: {region_np_processed.shape}, dtype: {region_np_processed.dtype}")
+        
+        # Run Cellpose segmentation
+        masks, flows, styles = model.eval(region_np_processed, diameter=None, channels=[0, 0])
+        print(f"DEBUG: After model.eval: masks type: {type(masks)}, shape: {masks.shape if hasattr(masks, 'shape') else 'N/A'}")
+        
+        # Update global region coords for info display
+        current_region_coords = region_coords
+        
+        # Count unique cells
+        cell_count = len(np.unique(masks)) - 1
+        print(f"DEBUG: Cell count: {cell_count}")
+        
+        # Create basic segmentation overlay
+        overlay = region_np_processed.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)
+        print("DEBUG: segment_cells returning successfully, viability section visible.")
+        
+        # Return masks and original processed image for state components
+        return cell_count, overlay_image, info_msg, viability_section_update, masks, region_np_processed
+            
+    except Exception as e:
+        print(f"ERROR: An exception occurred in segment_cells: {e}")
+        current_region_coords = None
+        # Return None for masks and image in case of error
+        return 0, None, f"Error: {str(e)}", gr.update(visible=False), None, None
+
+def update_viability_assessment(blue_threshold, stored_masks, stored_image):
+    """
+    Real-time viability assessment based on blue threshold
+    """
+    
+    print("DEBUG: --- Starting update_viability_assessment ---")
+    print(f"DEBUG: At start of update_viability_assessment: stored_masks is None? {stored_masks is None}, stored_image is None? {stored_image is None}")
+    
+    if stored_masks is None or stored_image is None:
+        print("DEBUG: Passed-in stored_masks or stored_image is None, returning 'Please run segmentation first'.")
+        return None, 0, 0, "Please run segmentation first"
+    
+    try:
+        # Debug information
+        unique_cells = np.unique(stored_masks)
+        total_cells = len(unique_cells) - 1  # Subtract background
+                
+        debug_info = f"Debug - Total cells found: {total_cells}\n"
+        debug_info += f"Image shape: {stored_image.shape}\n"
+        debug_info += f"Masks shape: {stored_masks.shape}\n"
+        debug_info += f"Unique mask values (first 10): {unique_cells[:10]}...\n"
+                
+        # Create viability overlay
+        viability_overlay, live_count, dead_count = create_viability_overlay(
+            stored_image, stored_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}"
+        print("DEBUG: update_viability_assessment returning successfully.")
+                
+        return viability_overlay, live_count, dead_count, viability_info
+            
+    except Exception as e:
+        print(f"ERROR: An exception occurred in update_viability_assessment: {e}")
+        return None, 0, 0, f"Error in viability assessment: {str(e)}\nMasks available: {stored_masks is not None}\nImage available: {stored_image is not None}"
+
+@spaces.GPU
+def segment_with_coords(image, model_choice, crop_coords):
+    """Alternative function for manual coordinates input"""
+    global current_region_coords
+        
+    print("DEBUG: --- Starting segment_with_coords ---")
+    if image is None:
+        print("DEBUG: No image provided for segment_with_coords, returning error.")
+        current_region_coords = None
+        return 0, None, "No image provided", gr.update(visible=False), None, None
+        
+    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)
+        print(f"DEBUG: image_np from PIL: shape={image_np.shape}, dtype={image_np.dtype}")
+                
+        # 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_processed = image_np[y_min:y_max, x_min:x_max]
+                    current_region_coords = (x_min, y_min, x_max, y_max)
+                    print(f"DEBUG: Applied manual crop in segment_with_coords: {current_region_coords}, new image_np_processed shape: {image_np_processed.shape}")
+                else:
+                    image_np_processed = image_np # No valid crop, use original
+            except ValueError:
+                print(f"DEBUG: Invalid crop coordinates format in segment_with_coords: {crop_coords}")
+                image_np_processed = image_np # Invalid crop, use original
+        else:
+            image_np_processed = image_np # No crop coords, use original
+        
+        # Process image format
+        if len(image_np_processed.shape) == 2:
+            image_np_final = cv2.cvtColor(image_np_processed, cv2.COLOR_GRAY2RGB)
+        elif len(image_np_processed.shape) == 3 and image_np_processed.shape[2] == 4:
+            image_np_final = cv2.cvtColor(image_np_processed, cv2.COLOR_RGBA2RGB)
+        else:
+            image_np_final = image_np_processed
+        
+        print(f"DEBUG: Image before model.eval (segment_with_coords) shape: {image_np_final.shape}, dtype: {image_np_final.dtype}")
+        # Run Cellpose
+        masks, flows, styles = model.eval(image_np_final, diameter=None, channels=[0, 0])
+        print(f"DEBUG: After model.eval (segment_with_coords): masks type: {type(masks)}, shape: {masks.shape if hasattr(masks, 'shape') else 'N/A'}")
+        
+        # Count unique cells
+        cell_count = len(np.unique(masks)) - 1
+                
+        # Create overlay
+        overlay = image_np_final.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"
+        if current_region_coords: # Use global if set by crop
+            info_msg += f"Region: {current_region_coords}\n"
+        info_msg += f"Now use the blue threshold slider below for viability assessment."
+        print("DEBUG: segment_with_coords returning successfully, viability section visible.")
+                
+        return cell_count, overlay_image, info_msg, gr.update(visible=True), masks, image_np_final
+            
+    except Exception as e:
+        print(f"ERROR: An exception occurred in segment_with_coords: {e}")
+        current_region_coords = None
+        return 0, None, f"Error: {str(e)}", gr.update(visible=False), None, None
+
+# 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.")
+        
+    # Define State components to store masks and image data across function calls
+    # These are hidden from the UI but persist data per session.
+    masks_state = gr.State(value=None)
+    image_state = gr.State(value=None)
+
+    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_cells now returns masks and image_np which are stored in masks_state and image_state
+        segment_btn1.click(
+            fn=segment_cells,
+            inputs=[image_editor, model_dropdown1],
+            outputs=[cell_count_output1, overlay_output1, info_output1, viability_section1, masks_state, image_state]
+        ).then(  # Chain the viability assessment after segmentation
+            fn=update_viability_assessment,
+            inputs=[blue_threshold_slider1, masks_state, image_state], # Pass stored state as inputs
+            outputs=[viability_overlay1, live_count_output1, dead_count_output1, viability_info1]
+        )
+                
+        # blue_threshold_slider changes also need the stored state
+        blue_threshold_slider1.change(
+            fn=update_viability_assessment,
+            inputs=[blue_threshold_slider1, masks_state, image_state], # Pass stored state as inputs
+            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_with_coords now returns masks and image_np which are stored in masks_state and image_state
+        segment_btn2.click(
+            fn=segment_with_coords,
+            inputs=[image_input, model_dropdown2, coord_input],
+            outputs=[cell_count_output2, overlay_output2, info_output2, viability_section2, masks_state, image_state]
+        ).then(  # Chain the viability assessment after segmentation
+            fn=update_viability_assessment,
+            inputs=[blue_threshold_slider2, masks_state, image_state], # Pass stored state as inputs
+            outputs=[viability_overlay2, live_count_output2, dead_count_output2, viability_info2]
+        )
+                
+        # blue_threshold_slider changes also need the stored state
+        blue_threshold_slider2.change(
+            fn=update_viability_assessment,
+            inputs=[blue_threshold_slider2, masks_state, image_state], # Pass stored state as inputs
+            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()