Update app.py

app.py

@@ -14,7 +14,6 @@ MODEL_OPTIONS = {
     "Hemocytometer Model": "hemocytometermodel.npy",
     "General Model": "generalmodel.npy"
 }
-
 loaded_models = {}
 
 # Global variables to store segmentation results for real-time viability assessment
@@ -26,9 +25,7 @@ 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():
@@ -40,30 +37,35 @@ def debug_editor_data(editor_data):
                     if hasattr(layer, 'size'):
                         info += f"      Size: {layer.size}\n"
                         info += f"      Mode: {layer.mode}\n"
-    else:
+    else: # Direct PIL Image
         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:
+        print("DEBUG: extract_region_from_editor - editor_data is None")
         return None, None
     
+    print(f"DEBUG: extract_region_from_editor - editor_data type: {type(editor_data)}")
+    
     if isinstance(editor_data, dict):
         background = editor_data.get('background')
         layers = editor_data.get('layers', [])
         
         if background is None:
+            print("DEBUG: extract_region_from_editor - background is None")
             return None, None
-            
+        
         background_np = np.array(background)
+        print(f"DEBUG: extract_region_from_editor - background_np shape: {background_np.shape}")
         
         if layers and len(layers) > 0:
             selection_layer = layers[0]
             selection_np = np.array(selection_layer)
+            print(f"DEBUG: extract_region_from_editor - selection_np shape: {selection_np.shape}")
             
             if len(selection_np.shape) == 3:
                 if selection_np.shape[2] == 4:  # RGBA
@@ -86,15 +88,17 @@ def extract_region_from_editor(editor_data):
                 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:
+    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):
@@ -106,7 +110,7 @@ def calculate_blue_intensity(image_np, masks):
     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)
@@ -118,7 +122,7 @@ def calculate_blue_intensity(image_np, 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]
         
@@ -126,8 +130,7 @@ def calculate_blue_intensity(image_np, masks):
             # 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}")
-    
+            # 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
 
@@ -135,18 +138,18 @@ 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: masks 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}")
+    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)
     
@@ -162,19 +165,18 @@ def create_viability_overlay(image_np, masks, blue_threshold):
     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}")
+        # 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")
+            # 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")
-    
+            # 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
@@ -192,7 +194,9 @@ def segment_cells(editor_data, model_choice, crop_coords=None):
     global current_masks, current_image, current_region_coords
     
     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]
@@ -208,9 +212,15 @@ def segment_cells(editor_data, model_choice, crop_coords=None):
         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_masks = None # Ensure globals are reset
+            current_image = None
+            current_region_coords = None
             return 0, None, f"No image provided. Debug info:\n{debug_info}", gr.update(visible=False)
         
-        # Handle manual crop coordinates
+        print(f"DEBUG: Image to segment shape (after extraction/cropping): {region_np.shape}, dtype: {region_np.dtype}")
+        
+        # Handle manual crop coordinates (logic already exists, no extra prints needed here unless it's the issue)
         if crop_coords and crop_coords.strip():
             try:
                 coords = [int(x.strip()) for x in crop_coords.split(',')]
@@ -221,10 +231,12 @@ def segment_cells(editor_data, model_choice, crop_coords=None):
                     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)
+                    current_region_coords = (x_min, y_min, x_max, y_max)
+                    print(f"DEBUG: Applied manual crop: {current_region_coords}, new region shape: {region_np.shape}")
             except ValueError:
+                print(f"DEBUG: Invalid crop coordinates format: {crop_coords}")
                 pass
         
         # Process image format
@@ -233,18 +245,24 @@ def segment_cells(editor_data, model_choice, crop_coords=None):
         elif len(region_np.shape) == 3 and region_np.shape[2] == 4:
             region_np = cv2.cvtColor(region_np, cv2.COLOR_RGBA2RGB)
         
+        print(f"DEBUG: Image before model.eval shape: {region_np.shape}, dtype: {region_np.dtype}")
         # Run Cellpose segmentation
         masks, flows, styles = model.eval(region_np, 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'}")
         
         # Store results globally for viability assessment
         current_masks = masks
         current_image = region_np
         current_region_coords = region_coords
         
+        print(f"DEBUG: Globals set - current_masks is None? {current_masks is None}, current_image is None? {current_image is None}")
+        print(f"DEBUG: Globals set - current_masks shape: {current_masks.shape if current_masks is not None else 'N/A'}")
+        
         # Count unique cells
         cell_count = len(np.unique(masks)) - 1
+        print(f"DEBUG: Cell count: {cell_count}")
         
-        # Create basic segmentation overlay
+        # Create basic segmentation overlay (logic unchanged)
         overlay = region_np.copy().astype(np.float32)
         
         if masks.max() > 0:
@@ -255,7 +273,7 @@ def segment_cells(editor_data, model_choice, crop_coords=None):
             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)
         
@@ -266,10 +284,12 @@ def segment_cells(editor_data, model_choice, crop_coords=None):
         
         # Make viability section visible
         viability_section_update = gr.update(visible=True)
+        print("DEBUG: segment_cells returning successfully, viability section visible.")
         
         return cell_count, overlay_image, info_msg, viability_section_update
         
     except Exception as e:
+        print(f"ERROR: An exception occurred in segment_cells: {e}")
         current_masks = None
         current_image = None
         current_region_coords = None
@@ -281,7 +301,11 @@ def update_viability_assessment(blue_threshold):
     """
     global current_masks, current_image
     
+    print("DEBUG: --- Starting update_viability_assessment ---")
+    print(f"DEBUG: At start of update_viability_assessment: current_masks is None? {current_masks is None}, current_image is None? {current_image is None}")
+    
     if current_masks is None or current_image is None:
+        print("DEBUG: current_masks or current_image is None, returning 'Please run segmentation first'.")
         return None, 0, 0, "Please run segmentation first"
     
     try:
@@ -292,7 +316,7 @@ def update_viability_assessment(blue_threshold):
         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
+        debug_info += f"Unique mask values (first 10): {unique_cells[:10]}...\n"
         
         # Create viability overlay
         viability_overlay, live_count, dead_count = create_viability_overlay(
@@ -306,22 +330,32 @@ def update_viability_assessment(blue_threshold):
         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: {current_masks is not None}\nImage available: {current_image is not None}"
 
+# The segment_with_coords function and Gradio UI setup remain the same
+# ... (rest of your app.py code) ...
+
 @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
-    
+
+    print("DEBUG: --- Starting segment_with_coords ---")
     if image is None:
+        print("DEBUG: No image provided for segment_with_coords, returning error.")
+        current_masks = None
+        current_image = None
+        current_region_coords = None
         return 0, None, "No image provided", gr.update(visible=False)
-    
+
     try:
         # Load the model
         model_filename = MODEL_OPTIONS[model_choice]
@@ -332,10 +366,11 @@ def segment_with_coords(image, model_choice, crop_coords):
         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:
@@ -347,51 +382,61 @@ def segment_with_coords(image, model_choice, crop_coords):
                     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)
+                    print(f"DEBUG: Applied manual crop in segment_with_coords: {current_region_coords}, new image_np shape: {image_np.shape}")
             except ValueError:
+                print(f"DEBUG: Invalid crop coordinates format in segment_with_coords: {crop_coords}")
                 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)
         
+        print(f"DEBUG: Image before model.eval (segment_with_coords) shape: {image_np.shape}, dtype: {image_np.dtype}")
         # Run Cellpose
         masks, flows, styles = model.eval(image_np, 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'}")
+
         # Store results globally
         current_masks = masks
         current_image = image_np
         
+        print(f"DEBUG: Globals set (segment_with_coords) - current_masks is None? {current_masks is None}, current_image is None? {current_image is None}")
+        print(f"DEBUG: Globals set (segment_with_coords) - current_masks shape: {current_masks.shape if current_masks is not None else 'N/A'}")
+
         # 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]
-            
+            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."
-        
+        print("DEBUG: segment_with_coords returning successfully, viability section visible.")
+
         return cell_count, overlay_image, info_msg, gr.update(visible=True)
-        
+
     except Exception as e:
+        print(f"ERROR: An exception occurred in segment_with_coords: {e}")
         current_masks = None
         current_image = None
+        current_region_coords = None
         return 0, None, f"Error: {str(e)}", gr.update(visible=False)
 
 # Create the Gradio interface
@@ -416,17 +461,17 @@ with gr.Blocks(title="Cell Counter with Real-time Viability Assessment", theme=g
                     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(
@@ -437,14 +482,14 @@ with gr.Blocks(title="Cell Counter with Real-time Viability Assessment", theme=g
                         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,
@@ -455,16 +500,16 @@ with gr.Blocks(title="Cell Counter with Real-time Viability Assessment", theme=g
             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")
@@ -479,17 +524,17 @@ with gr.Blocks(title="Cell Counter with Real-time Viability Assessment", theme=g
                     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(
@@ -500,14 +545,14 @@ with gr.Blocks(title="Cell Counter with Real-time Viability Assessment", theme=g
                         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,
@@ -518,36 +563,36 @@ with gr.Blocks(title="Cell Counter with Real-time Viability Assessment", theme=g
             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"
-        
+                
+        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
-        
+            - 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
-        
+            - 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
@@ -555,4 +600,4 @@ with gr.Blocks(title="Cell Counter with Real-time Viability Assessment", theme=g
         """)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()