Update app.py

app.py

@@ -128,12 +128,14 @@ def cellpose_segment(img_input):
     crand = .2 + .8 * np.random.rand(np.max(masks.flatten()).astype('int')+1,).astype('float32')
     crand[0] = 0
 
-    overlay = Image.fromarray(overlay)
-    flows = Image.fromarray(flows)
-    
-    masks = Image.fromarray(255. * crand[masks])
+    #overlay = Image.fromarray(overlay)
+    #flows = Image.fromarray(flows)    
+    #masks = Image.fromarray(255. * crand[masks])
+
+    pil_masks = Image.fromarray(masks.astype('int32'))
+    pil_masks.save("test.png")
     
-    return outpix, overlay, flows, masks
+    return outpix, overlay, flows, masks, gr.DownloadButton(visible=True, value = "test.png")
 
 # Gradio Interface
 #iface = gr.Interface(
@@ -144,6 +146,9 @@ def cellpose_segment(img_input):
 #    description="upload an image, then cellpose will segment it at a max size of 400x400 (for full functionality, 'pip install cellpose' locally)"
 #)
 
+def download_file():
+    return gr.DownloadButton(v"Download segmentations", visible=False)
+    
 with gr.Blocks(title = "Hello", 
                css=".gradio-container {background:purple;}") as demo:
 
@@ -151,16 +156,18 @@ with gr.Blocks(title = "Hello",
         with gr.Column(scale=1):
             input_image = gr.Image(label = "Input image", type = "numpy")
             send_btn = gr.Button("Run Cellpose-SAM")
+            down_btn = gr.DownloadButton("Download segmentations", visible=False)
 
         with gr.Column(scale=1):            
             img_outlines = gr.Image(label = "Output image", type = "pil")            
-            img_overlay = gr.Image(label = "Output image", type = "pil")            
-            flows = gr.Image(label = "Output image", type = "pil")
-            masks = gr.Image(label = "Output image", type = "pil")
-            
-
-    send_btn.click(fn=cellpose_segment, inputs=[input_image], outputs=[img_outlines, img_overlay, flows, masks])
+            img_overlay = gr.Image(label = "Output image", type = "numpy")            
+            flows = gr.Image(label = "Output image", type = "numpy")
+            masks = gr.Image(label = "Output image", type = "numpy")
+    
+    
+    send_btn.click(fn=cellpose_segment, inputs=[input_image], outputs=[img_outlines, img_overlay, flows, masks, down_btn])
 
+    down_btn.click(download_file, None, down_btn)