Update app.py

app.py

@@ -281,7 +281,48 @@ def run_segmentation_editor(editor_data, model_choice):
 
     except Exception as e:
         return 0, None, f"Error during segmentation: {str(e)}", gr.update(visible=False), None, None, 0.0
-
+        
+def run_segmentation_with_cleanup(func):
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            # Force garbage collection on error
+            import gc
+            gc.collect()
+            return 0, None, f"Error: {str(e)}", gr.update(visible=False), None, None, 0.0
+    return wrapper
+
+run_segmentation_editor = run_segmentation_with_cleanup(run_segmentation_editor)
+
+def validate_and_prepare_image(image_np, max_side=384, max_pixels=384*384):
+    """
+    Validate image and prepare for processing with strict size limits
+    Returns: (processed_image, error_message)
+    """
+    if image_np is None:
+        return None, "No image provided"
+    
+    # Check if image is too small
+    h, w = image_np.shape[:2]
+    if h < 50 or w < 50:
+        return None, "Image too small (minimum 50x50 pixels)"
+    
+    # Check file size before processing
+    if h * w > 4000 * 4000:  # Reject images over 16 megapixels before resize
+        return None, "Image too large. Please resize to under 4000x4000 before uploading."
+    
+    # Apply safe resize
+    resized = safe_resize(image_np)
+    
+    # Ensure RGB format
+    if len(resized.shape) == 2:
+        resized = cv2.cvtColor(resized, cv2.COLOR_GRAY2RGB)
+    elif len(resized.shape) == 3 and resized.shape[2] == 4:
+        resized = cv2.cvtColor(resized, cv2.COLOR_RGBA2RGB)
+    
+    return resized, None
+    
 
 def update_viability_realtime(blue_threshold, stored_masks, stored_image_np):
     # avoid unpacking None (e.g. slider moved before segmentation)
@@ -310,28 +351,30 @@ def update_viability_realtime(blue_threshold, stored_masks, stored_image_np):
         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">
-<meta name="apple-mobile-web-app-capable" content="yes">
-<meta name="apple-mobile-web-app-status-bar-style" content="default">
-<meta name="apple-mobile-web-app-title" content="Cellpose Cell Counter">
-<meta name="viewport" content="width=device-width, initial-scale=1">
-<link rel="apple-touch-icon" href="/file=favicon.png">
-<link rel="icon" type="image/png" sizes="192x192" href="/file=favicon.png">
-<script>
-if ('serviceWorker' in navigator) {
-  window.addEventListener('load', function() {
-    navigator.serviceWorker.register('/file=static.zip/static/service-worker.js')
-      .then(function(registration) {
-        console.log('ServiceWorker registration successful');
-      }, function(err) {
-        console.log('ServiceWorker registration failed: ', err);
-      });
-  });
-}
-</script>
-"""
+def diagnose_image_upload(image):
+    """Helper function to diagnose upload issues"""
+    try:
+        if image is None:
+            return "No image received"
+        
+        img_np = np.array(image)
+        h, w = img_np.shape[:2]
+        channels = img_np.shape[2] if len(img_np.shape) == 3 else 1
+        dtype = img_np.dtype
+        size_mb = img_np.nbytes / (1024 * 1024)
+        
+        info = f"Image received successfully!\n"
+        info += f"Dimensions: {w}x{h}\n"
+        info += f"Channels: {channels}\n"
+        info += f"Data type: {dtype}\n"
+        info += f"Memory size: {size_mb:.2f} MB\n"
+        
+        if size_mb > 10:
+            info += "\n⚠️ WARNING: Image is large and may cause issues on mobile"
+        
+        return info
+    except Exception as e:
+        return f"Error analyzing image: {str(e)}"
 
 
 def pack_array(arr):
@@ -364,19 +407,30 @@ with gr.Blocks(
         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)
+                label="Draw selection on image",
+                type="pil",
+                image_mode="RGB",
+                height=384,  # Reduced from 512
+                width=384,   # ADD explicit width
+                brush=gr.Brush(colors=["#ff0000"], color_mode="fixed", default_size=20),
+                eraser=gr.Eraser(default_size=20)
+                )
+            
+            image_input = gr.Image(
+                type="pil",
+                label="Microscopy Image",
+                image_mode="RGB",
+                height=384,  # Reduced from 512
+                width=384    # ADD explicit width
                 )
-                model_dropdown1 = gr.Dropdown(
-                    choices=list(MODEL_OPTIONS.keys()),
-                    label="Select Model",
-                    value="Hemocytometer Model"
+            
+            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")
+                
+            segment_btn1 = gr.Button("🔬 Run Segmentation", variant="primary", size="lg")
 
             with gr.Column():
                 cell_count_output1 = gr.Number(label="Total Cells Detected", precision=0)
@@ -455,4 +509,9 @@ with gr.Blocks(
         """)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(
+        share=False,
+        show_error=True,  # Show detailed errors
+        max_file_size="10mb",  # Limit upload size
+        favicon_path="favicon.png"
+    )