Update app.py

app.py

@@ -126,15 +126,20 @@ def calculate_blue_intensity(image_np, masks):
 
 def create_viability_overlay(image_np, masks, blue_threshold):
     """Create overlay showing live/dead cells based on blue threshold"""
-    global current_masks, current_image
     
     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
+    # Create overlay - start with original image
     overlay = image_np.copy().astype(np.float32)
     
     live_count = 0
@@ -154,11 +159,11 @@ def create_viability_overlay(image_np, masks, blue_threshold):
             live_count += 1
     
     # Blend with original image
-    alpha = 0.5
-    overlay = (1 - alpha) * image_np + alpha * overlay
-    overlay = np.clip(overlay, 0, 255).astype(np.uint8)
+    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(overlay), live_count, dead_count
+    return Image.fromarray(final_overlay), live_count, dead_count
 
 @spaces.GPU
 def segment_cells(editor_data, model_choice, crop_coords=None):
@@ -261,22 +266,34 @@ def update_viability_assessment(blue_threshold):
         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"Total cells: {total_count}\nLive (green): {live_count}\nDead (red): {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"\nViability: {viability_percentage:.1f}%"
+            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)}"
+        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):
@@ -414,6 +431,10 @@ with gr.Blocks(title="Cell Counter with Real-time Viability Assessment", theme=g
             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(
@@ -473,6 +494,10 @@ with gr.Blocks(title="Cell Counter with Real-time Viability Assessment", theme=g
             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(