Update app.py

app.py

@@ -8,8 +8,7 @@ import tempfile
 from PIL import Image
 import io
 from huggingface_hub import hf_hub_download
-import base64
-
+import base64  # needed for mobile ImageEditor base64 handling
 
 HF_REPO_ID = "myang4218/cellposemodel"
 MODEL_OPTIONS = {
@@ -19,6 +18,34 @@ MODEL_OPTIONS = {
 
 loaded_models = {}
 
+# ---- mobile-safe size limits (aggressive for Safari) ----
+MAX_SIDE = 512           # shrink any side larger than this
+MAX_PIXELS = 512 * 512   # ~0.26 megapixels
+
+
+def safe_resize(image_np):
+    """
+    Downscale image to fit within MAX_SIDE and MAX_PIXELS while
+    preserving aspect ratio. Works for RGB / RGBA / grayscale.
+    """
+    h, w = image_np.shape[:2]
+    total = h * w
+
+    # already small enough
+    if max(h, w) <= MAX_SIDE and total <= MAX_PIXELS:
+        return image_np
+
+    # compute scale so BOTH constraints are satisfied
+    scale_side = MAX_SIDE / max(h, w)
+    scale_pixels = (MAX_PIXELS / total) ** 0.5
+    scale = min(scale_side, scale_pixels)
+
+    new_w = max(1, int(w * scale))
+    new_h = max(1, int(h * scale))
+
+    return cv2.resize(image_np, (new_w, new_h), interpolation=cv2.INTER_AREA)
+
+
 def extract_region_from_editor(editor_data):
     """Extract selected region from ImageEditor data (mobile-safe)."""
     if editor_data is None:
@@ -79,13 +106,13 @@ def extract_region_from_editor(editor_data):
             else:
                 try:
                     sel_np = np.array(sel)
-                except:
+                except Exception:
                     sel_np = None
 
             # If selection exists: compute bounding box
             if sel_np is not None:
                 if sel_np.ndim == 3:
-                    if sel_np.shape[2] == 4: 
+                    if sel_np.shape[2] == 4:
                         mask = sel_np[:, :, 3] > 0
                     else:
                         mask = np.any(sel_np > 0, axis=2)
@@ -96,7 +123,7 @@ def extract_region_from_editor(editor_data):
                 if coords[0].size > 0:
                     y0, y1 = coords[0].min(), coords[0].max()
                     x0, x1 = coords[1].min(), coords[1].max()
-                    return background_np[y0:y1+1, x0:x1+1], (x0, y0, x1, y1)
+                    return background_np[y0:y1 + 1, x0:x1 + 1], (x0, y0, x1, y1)
 
         # No selection → return full image
         return background_np, None
@@ -108,12 +135,12 @@ def extract_region_from_editor(editor_data):
 def classify_cells_by_blueness(image_np, masks, blue_threshold):
     """
     Classify cells as dead (blue) or alive based on single blueness metric
-   
+
     Args:
         image_np: RGB image array
         masks: Cellpose segmentation masks
         blue_threshold: Single threshold value (0-100) for blueness detection
-    
+
     Returns:
         dead_count, alive_count, colored_overlay
     """
@@ -125,62 +152,63 @@ def classify_cells_by_blueness(image_np, masks, blue_threshold):
 
     # Convert RGB to HSV
     hsv = cv2.cvtColor(image_np, cv2.COLOR_RGB2HSV)
-    
+
     # Calculate blueness index for each pixel
     hue = hsv[:, :, 0].astype(np.float32)
     saturation = hsv[:, :, 1].astype(np.float32)
-    
+
     # Hue score: peaks around 115 (blue in HSV), drops off towards edges
     # Handle hue wrap-around for blue detection (100-130 range)
     hue_distance = np.minimum(np.abs(hue - 115), 180 - np.abs(hue - 115))
     hue_score = np.maximum(0, 1 - hue_distance / 65)  # 65 gives good blue range
-    
+
     # Combine hue proximity with saturation intensity
     blueness = hue_score * (saturation / 255.0)
-    
+
     # Convert threshold from 0-100 to 0-1 scale
     threshold = blue_threshold / 100.0
-    
+
     # 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 average blueness for this cell
         cell_blueness = np.mean(blueness[cell_mask])
-        
+
         # Classify based on threshold
         if cell_blueness > threshold:
             dead_cells.append(cell_id)
         else:
             alive_cells.append(cell_id)
-    
+
     # Create colored overlay
-    overlay = image_np.copy().astype(np.float32) # Ensure float for blending
-    
+    overlay = image_np.copy().astype(np.float32)  # Ensure float for blending
+
     # 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
     final_overlay = (1 - alpha) * image_np.astype(np.float32) + alpha * overlay
     final_overlay = np.clip(final_overlay, 0, 255).astype(np.uint8)
-    
+
     return len(dead_cells), len(alive_cells), final_overlay
-    
+
+
 def measure_confluency(masks, image_np):
     """Calculate the percentage of image area covered by cells"""
     tot_pixels = image_np.shape[0] * image_np.shape[1]
@@ -198,28 +226,21 @@ def run_segmentation_editor(editor_data, model_choice):
     try:
         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
-        
+
         region_np, region_coords = extract_region_from_editor(editor_data)
-        
+
         if region_np is None:
             return 0, None, f"No image provided.", gr.update(visible=False), None, None, 0.0
-        
-        # Resize large images to prevent crashes
-        max_size = 256
-        if region_np.shape[0] > max_size or region_np.shape[1] > max_size:
-            h, w = region_np.shape[:2]
-            if h > w:
-                new_h, new_w = max_size, int(w * max_size / h)
-            else:
-                new_h, new_w = int(h * max_size / w), max_size
-            region_np = cv2.resize(region_np, (new_w, new_h), interpolation=cv2.INTER_AREA)
-        
+
+        # Enforce mobile-safe size limit immediately
+        region_np = safe_resize(region_np)
+
         # Process image format to RGB
         if len(region_np.shape) == 2:
             processed_image_np = cv2.cvtColor(region_np, cv2.COLOR_GRAY2RGB)
@@ -227,36 +248,45 @@ def run_segmentation_editor(editor_data, model_choice):
             processed_image_np = cv2.cvtColor(region_np, cv2.COLOR_RGBA2RGB)
         else:
             processed_image_np = region_np
-        
+
         # Run Cellpose segmentation
         masks, flows, styles = model.eval(processed_image_np, diameter=None, channels=[0, 0])
-        
+
         cell_count = len(np.unique(masks)) - 1
         confluency = measure_confluency(masks, processed_image_np)
-        
+
         # Create a basic segmentation overlay (without viability)
         segmentation_overlay = processed_image_np.copy().astype(np.float32)
         if masks.max() > 0:
-            np.random.seed(42) # For consistent random colors
+            np.random.seed(42)  # For consistent random colors
             colors = np.random.randint(0, 255, size=(masks.max() + 1, 3))
-            colors[0] = [0, 0, 0] # Background color
+            colors[0] = [0, 0, 0]  # Background color
             colored_mask = colors[masks]
             alpha = 0.4
             segmentation_overlay = (1 - alpha) * segmentation_overlay + alpha * colored_mask
         segmentation_overlay = np.clip(segmentation_overlay, 0, 255).astype(np.uint8)
-        
+
         info_msg = f"Segmentation complete! Found {cell_count} cells.\n"
         info_msg += f"Confluency: {confluency:.1f}%\n"
         if region_coords:
             info_msg += f"Processed region: {region_coords[0]},{region_coords[1]} to {region_coords[2]},{region_coords[3]}\n"
         info_msg += f"Now adjust the Blue Threshold for viability assessment."
-                
-        # Return initial segmentation display and state variables
-        return cell_count, Image.fromarray(segmentation_overlay), info_msg, gr.update(visible=True), pack_array(masks), pack_array(processed_image_np), confluency
-            
+
+        # Return initial segmentation display and state variables (packed for Gradio State)
+        return (
+            cell_count,
+            Image.fromarray(segmentation_overlay),
+            info_msg,
+            gr.update(visible=True),
+            pack_array(masks),
+            pack_array(processed_image_np),
+            confluency,
+        )
+
     except Exception as e:
         return 0, None, f"Error during segmentation: {str(e)}", gr.update(visible=False), None, None, 0.0
 
+
 @spaces.GPU
 def run_segmentation_manual(image, model_choice, crop_coords):
     """
@@ -265,29 +295,22 @@ def run_segmentation_manual(image, model_choice, crop_coords):
     """
     if image is None:
         return 0, None, "No image provided", gr.update(visible=False), None, None, 0.0
-        
+
     try:
         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
-                
+
         image_np = np.array(image)
-        
-        # Resize large images to prevent crashes
-        max_size = 256
-        if image_np.shape[0] > max_size or image_np.shape[1] > max_size:
-            h, w = image_np.shape[:2]
-            if h > w:
-                new_h, new_w = max_size, int(w * max_size / h)
-            else:
-                new_h, new_w = int(h * max_size / w), max_size
-            image_np = cv2.resize(image_np, (new_w, new_h), interpolation=cv2.INTER_AREA)
-                
+
+        # Enforce mobile-safe size limit immediately
+        image_np = safe_resize(image_np)
+
         # Apply crop if coordinates are provided
         if crop_coords and crop_coords.strip():
             try:
@@ -295,81 +318,89 @@ def run_segmentation_manual(image, model_choice, crop_coords):
                 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))
-                                        
+                    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))
+
                     processed_image_np = image_np[y_min:y_max, x_min:x_max]
                 else:
-                    processed_image_np = image_np # No valid crop, use original
+                    processed_image_np = image_np  # No valid crop, use original
             except ValueError:
-                processed_image_np = image_np # Invalid crop, use original
+                processed_image_np = image_np  # Invalid crop, use original
         else:
-            processed_image_np = image_np # No crop coords, use original
-        
+            processed_image_np = image_np  # No crop coords, use original
+
         # Process image format to RGB
         if len(processed_image_np.shape) == 2:
             processed_image_np = cv2.cvtColor(processed_image_np, cv2.COLOR_GRAY2RGB)
         elif len(processed_image_np.shape) == 3 and processed_image_np.shape[2] == 4:
             processed_image_np = cv2.cvtColor(processed_image_np, cv2.COLOR_RGBA2RGB)
-        
+
         # Run Cellpose
         masks, flows, styles = model.eval(processed_image_np, diameter=None, channels=[0, 0])
-        
-        cell_count = len(np.unique(masks)) - 1    
+
+        cell_count = len(np.unique(masks)) - 1
         confluency = measure_confluency(masks, processed_image_np)
-        
+
         # Create a basic segmentation overlay
         segmentation_overlay = processed_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]                        
+            colors[0] = [0, 0, 0]
             colored_mask = colors[masks]
             alpha = 0.4
             segmentation_overlay = (1 - alpha) * segmentation_overlay + alpha * colored_mask
         segmentation_overlay = np.clip(segmentation_overlay, 0, 255).astype(np.uint8)
-                
+
         info_msg = f"Segmentation complete! Found {cell_count} cells.\n"
         info_msg += f"Confluency: {confluency:.1f}%\n"
         if crop_coords:
             info_msg += f"Processed with coordinates: {crop_coords}\n"
         info_msg += f"Now adjust the Blue Threshold for viability assessment."
-                
-        return cell_count, Image.fromarray(segmentation_overlay), info_msg, gr.update(visible=True), pack_array(masks), pack_array(processed_image_np), confluency
-            
+
+        return (
+            cell_count,
+            Image.fromarray(segmentation_overlay),
+            info_msg,
+            gr.update(visible=True),
+            pack_array(masks),
+            pack_array(processed_image_np),
+            confluency,
+        )
+
     except Exception as e:
         return 0, None, f"Error during segmentation: {str(e)}", gr.update(visible=False), None, None, 0.0
 
+
 def update_viability_realtime(blue_threshold, stored_masks, stored_image_np):
+    # avoid unpacking None (e.g. slider moved before segmentation)
     if stored_masks is None or stored_image_np is None:
         return None, 0, 0, 0.0, "Please run segmentation first."
-    
+
     stored_masks = unpack_array(stored_masks)
     stored_image_np = unpack_array(stored_image_np)
 
-    if stored_masks is None or stored_image_np is None:
-        return None, 0, 0, 0.0, "Please run segmentation first."
-    
     try:
         dead_count, alive_count, viability_overlay_np = classify_cells_by_blueness(
             stored_image_np, stored_masks, blue_threshold
         )
-        
+
         total_count = alive_count + dead_count
         viability_percent = (alive_count / total_count * 100) if total_count > 0 else 0.0
         confluency = measure_confluency(stored_masks, stored_image_np)
-        
+
         overlay_image = Image.fromarray(viability_overlay_np)
         info_msg = f"Total cells: {total_count}\nLive (green): {alive_count}\nDead (red): {dead_count}\n"
         info_msg += f"Viability: {viability_percent:.1f}%\nConfluency: {confluency:.1f}%\nBlue threshold: {blue_threshold}%"
-        
+
         return overlay_image, alive_count, dead_count, viability_percent, info_msg
-            
+
     except Exception as e:
         return None, 0, 0, 0.0, f"Error updating viability: {str(e)}"
 
+
 # PWA Head HTML - Updated for nested zip structure
 pwa_head = """
 <link rel="manifest" href="/file=static.zip/static/manifest.json">
@@ -392,36 +423,42 @@ if ('serviceWorker' in navigator) {
 }
 </script>
 """
+
+
 def pack_array(arr):
     pil = Image.fromarray(arr.astype(np.uint8))
     buf = io.BytesIO()
     pil.save(buf, format="PNG")
     return buf.getvalue()
 
+
 def unpack_array(data):
     return np.array(Image.open(io.BytesIO(data)))
 
+
 # Gradio interface
 with gr.Blocks(
-    title="CellposeCellCounter", 
+    title="CellposeCellCounter",
     theme=gr.themes.Soft(),
     head=pwa_head  # Add PWA head elements
 ) as demo:
     gr.Markdown("# CellposeCellCounter")
     gr.Markdown("For accurate cell confluency, crop the image to display only desired area.")
-        
+
+    # Define State components to store masks and image data across function calls
     masks_state = gr.State(value=None)
     image_state = gr.State(value=None)
 
-
     with gr.Tab("Image Editor (Draw Selection)"):
         gr.Markdown("### Draw selection and run segmentation")
-                
+
         with gr.Row():
             with gr.Column():
                 image_editor = gr.ImageEditor(
                     label="Draw selection on image",
                     type="pil",
+                    image_mode="RGB",
+                    height=512,  # limit canvas size on client (important for mobile)
                     brush=gr.Brush(colors=["#ff0000"], color_mode="fixed", default_size=20),
                     eraser=gr.Eraser(default_size=20)
                 )
@@ -431,18 +468,18 @@ with gr.Blocks(
                     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)
                 confluency_output1 = gr.Number(label="Confluency (%)", precision=1)
                 overlay_output1 = gr.Image(type="pil", label="Segmentation Result")
                 info_output1 = gr.Textbox(label="Processing Info", lines=4)
-                
+
         # Viability Assessment Section
         with gr.Group(visible=False) as viability_section1:
             gr.Markdown("### Viability Assessment (Trypan Blue)")
             gr.Markdown("Adjust the threshold to classify cells as live (green) or dead (red).")
-                        
+
             with gr.Row():
                 with gr.Column():
                     blue_threshold1 = gr.Slider(
@@ -453,27 +490,26 @@ with gr.Blocks(
                         label="Blue Threshold (%)",
                         info="Higher values = more selective for blue cells"
                     )
-                                    
+
                 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_percent_output1 = gr.Number(label="Viability (%)", precision=1)
             viability_info1 = gr.Textbox(label="Analysis Results", lines=5)
-                
+
         # Event handlers
-        # segment_cells now returns masks and image_np which are stored in masks_state and image_state
         segment_btn1.click(
             fn=run_segmentation_editor,
             inputs=[image_editor, model_dropdown1],
             outputs=[cell_count_output1, overlay_output1, info_output1, viability_section1, masks_state, image_state, confluency_output1]
         ).then(  # Chain the initial viability assessment after segmentation
             fn=update_viability_realtime,
-            inputs=[blue_threshold1, masks_state, image_state], # Pass stored state as inputs
+            inputs=[blue_threshold1, masks_state, image_state],  # Pass stored state as inputs
             outputs=[viability_overlay1, live_count_output1, dead_count_output1, viability_percent_output1, viability_info1]
         )
-                
+
         # Slider changes update viability in real-time
         blue_threshold1.change(
             fn=update_viability_realtime,
@@ -483,10 +519,15 @@ with gr.Blocks(
 
     with gr.Tab("Manual Coordinates"):
         gr.Markdown("### Upload image and run segmentation")
-                
+
         with gr.Row():
             with gr.Column():
-                image_input = gr.Image(type="pil", label="Microscopy Image")
+                image_input = gr.Image(
+                    type="pil",
+                    label="Microscopy Image",
+                    image_mode="RGB",
+                    height=512  # limit canvas size here too
+                )
                 model_dropdown2 = gr.Dropdown(
                     choices=list(MODEL_OPTIONS.keys()),
                     label="Select Model",
@@ -498,18 +539,18 @@ with gr.Blocks(
                     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)
                 confluency_output2 = gr.Number(label="Confluency (%)", precision=1)
                 overlay_output2 = gr.Image(type="pil", label="Segmentation Result")
                 info_output2 = gr.Textbox(label="Processing Info", lines=4)
-                
+
         # Viability Assessment Section
         with gr.Group(visible=False) as viability_section2:
             gr.Markdown("### Viability Assessment (Trypan Blue)")
             gr.Markdown("Adjust the threshold to classify cells as live (green) or dead (red).")
-                        
+
             with gr.Row():
                 with gr.Column():
                     blue_threshold2 = gr.Slider(
@@ -520,15 +561,15 @@ with gr.Blocks(
                         label="Blue Threshold (%)",
                         info="Higher values = more selective for blue cells"
                     )
-                                    
+
                 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_percent_output2 = gr.Number(label="Viability (%)", precision=1)
             viability_info2 = gr.Textbox(label="Analysis Results", lines=5)
-                
+
         # Event handlers
         segment_btn2.click(
             fn=run_segmentation_manual,
@@ -536,10 +577,10 @@ with gr.Blocks(
             outputs=[cell_count_output2, overlay_output2, info_output2, viability_section2, masks_state, image_state, confluency_output2]
         ).then(  # Chain the initial viability assessment after segmentation
             fn=update_viability_realtime,
-            inputs=[blue_threshold2, masks_state, image_state], 
+            inputs=[blue_threshold2, masks_state, image_state],
             outputs=[viability_overlay2, live_count_output2, dead_count_output2, viability_percent_output2, viability_info2]
         )
-                
+
         # Slider changes update viability in real-time
         blue_threshold2.change(
             fn=update_viability_realtime,
@@ -551,17 +592,17 @@ with gr.Blocks(
     with gr.Accordion("Instructions", open=False):
         gr.Markdown("""
         ### How to use:
-                
+
         1. **Upload and Segment**:
             - Upload your microscopy image.
             - Select a Cellpose model (e.g., "Hemocytometer Model" for blood cells).
             - Draw a selection region using the Image Editor, or specify coordinates manually.
             - Click "Run Segmentation".
-                
+
         2. **Analysis Results**:
             - **Cell Count**: Total number of detected cells
             - **Confluency**: Percentage of image area covered by cells (useful for assessing cell density). Note that cell confluency is calculated per the entire area of the image input.
-                
+
         3. **Real-time Viability Assessment (Trypan Blue)**:
             - After segmentation, the viability section will become visible.
             - This tool is specifically designed for **Trypan Blue stained images**, where dead cells appear blue.
@@ -569,11 +610,11 @@ with gr.Blocks(
             - **Lower values (e.g., 10-20%)** are more sensitive and will classify more cells as blue/dead.
             - **Higher values (e.g., 30-50%)** are more selective and will only classify strongly blue cells as dead.
             - Green cells = Live, Red cells = Dead.
-                
+
         4. **Interpreting Results**:
             - The app calculates and displays the total, live, and dead cell counts, along with the viability percentage and confluency.
             - **Confluency** helps assess how densely packed your cells are, which is important for cell culture monitoring.
         """)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()