Update app.py

app.py

@@ -8,6 +8,7 @@ import tempfile
 from PIL import Image
 import io
 from huggingface_hub import hf_hub_download
+import base64
 
 
 HF_REPO_ID = "myang4218/cellposemodel"
@@ -251,7 +252,7 @@ def run_segmentation_editor(editor_data, model_choice):
         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), masks, 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
@@ -336,12 +337,14 @@ def run_segmentation_manual(image, model_choice, 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), masks, 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):
+    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)
@@ -408,9 +411,6 @@ with gr.Blocks(
     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)
-
     masks_state = pack_array(masks)
     image_state = pack_array(processed_image_np)