Update app.py

app.py

@@ -98,10 +98,6 @@ def image_resize(img, resize=400):
     img = img.astype(np.uint8)
     return img
 
-
-def plot_masks(masks):
-    
-    return RGB
     
 #@spaces.GPU(duration=10)
 #def run_model_gpu(model, img):
@@ -132,10 +128,16 @@ def cellpose_segment(img_input):
     #flows = Image.fromarray(flows)    
     #masks = Image.fromarray(255. * crand[masks])
 
-    pil_masks = Image.fromarray((masks/masks.max()).astype('float32'))
-    pil_masks.save("test.tiff")
+    pil_masks = Image.fromarray(masks.astype('int32'))
+    pil_masks.save("masks.tiff")
+
+    pil_outpix = Image.fromarray(outpix)
+    pil_outpix.save("outlines.png")
+
+    b1 = gr.DownloadButton(visible=True, value = "masks.tiff")
+    b2 = gr.DownloadButton(visible=True, value = "outlines.png")
     
-    return outpix, overlay, flows, gr.DownloadButton(visible=True, value = "test.tiff")
+    return outpix, overlay, flows, b1, b2
 
 # Gradio Interface
 #iface = gr.Interface(
@@ -146,9 +148,7 @@ 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("Download segmentations", visible=False)
-    
+   
 with gr.Blocks(title = "Hello", 
                css=".gradio-container {background:purple;}") as demo:
 
@@ -156,7 +156,8 @@ 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)
+            down_btn = gr.DownloadButton("Download masks as TIFF", visible=False)
+            down_btn2 = gr.DownloadButton("Download outline image as PNG", visible=False)
 
         with gr.Column(scale=1):            
             img_outlines = gr.Image(label = "Output image", type = "pil")            
@@ -165,9 +166,9 @@ with gr.Blocks(title = "Hello",
             #masks = gr.Image(label = "Output image", type = "numpy")
     
     
-    send_btn.click(fn=cellpose_segment, inputs=[input_image], outputs=[img_outlines, img_overlay, flows, down_btn])
+    send_btn.click(fn=cellpose_segment, inputs=[input_image], outputs=[img_outlines, img_overlay, flows, down_btn, down_btn2])
 
-    down_btn.click(download_file, None, down_btn)
+    down_btn.click(None, None, None)