Update app.py

app.py

@@ -152,23 +152,23 @@ def image_resize(img, resize=400):
 
     
 @spaces.GPU(duration=10)
-def run_model_gpu(img, max_iter):
-    masks, flows, _ = model.eval(img, niter = max_iter)#, channels = [0,0])
+def run_model_gpu(img, max_iter, flow_threshold, cellprob_threshold):
+    masks, flows, _ = model.eval(img, niter = max_iter, flow_threshold = flow_threshold, cellprob_threshold = cellprob_threshold)
     return masks, flows
 
 @spaces.GPU(duration=60)
-def run_model_gpu60(img, max_iter):
-    masks, flows, _ = model.eval(img, niter = max_iter)#, channels = [0,0])
+def run_model_gpu60(img, max_iter, flow_threshold, cellprob_threshold):
+    masks, flows, _ = model.eval(img, niter = max_iter, flow_threshold = flow_threshold, cellprob_threshold = cellprob_threshold)
     return masks, flows
 
 @spaces.GPU(duration=240)
-def run_model_gpu240(img, max_iter):
-    masks, flows, _ = model.eval(img, niter = max_iter)#, channels = [0,0])
+def run_model_gpu240(img, max_iter, flow_threshold, cellprob_threshold):
+    masks, flows, _ = model.eval(img, niter = max_iter, flow_threshold = flow_threshold, cellprob_threshold = cellprob_threshold)
     return masks, flows
 
 import datetime
 from zipfile import ZipFile
-def cellpose_segment(filepath, resize = 1000,max_iter = 250):
+def cellpose_segment(filepath, resize = 1000,max_iter = 250, flow_threshold= 0.4, cellprob_threshold = 0):
 
     zip_path = os.path.splitext(filepath[-1])[0]+"_masks.zip"
     #zip_path = 'masks.zip'
@@ -183,11 +183,11 @@ def cellpose_segment(filepath, resize = 1000,max_iter = 250):
             
             maxsize = np.max(img.shape)
             if maxsize<=1000:
-                masks, flows = run_model_gpu(img, max_iter)
+                masks, flows = run_model_gpu(img, max_iter, flow_threshold, cellprob_threshold)
             elif maxsize < 5000:
-                masks, flows = run_model_gpu60(img, max_iter)
+                masks, flows = run_model_gpu60(img, max_iter, flow_threshold, cellprob_threshold)
             elif maxsize < 20000:
-                masks, flows = run_model_gpu240(img, max_iter)
+                masks, flows = run_model_gpu240(img, max_iter, flow_threshold, cellprob_threshold)
             else:
                 raise ValueError("Image size must be less than 20,000")
 
@@ -310,6 +310,8 @@ with gr.Blocks(title = "Hello",
                     with gr.Row():
                         resize = gr.Number(label = 'max resize', value = 1000)
                         max_iter = gr.Number(label = 'max iterations', value = 250)
+                        flow_threshold = gr.Number(label = 'flow threshold', value = 0.4)
+                        cellprob_threshold = gr.Number(label = 'cellprob threshold', value = 0)
                         
                     up_btn = gr.UploadButton("Multi-file upload (png, jpg, tif etc)", visible=True, file_count = "multiple")                        
                     
@@ -336,7 +338,7 @@ with gr.Blocks(title = "Hello",
     input_image.upload(update_button, input_image, [input_image, up_btn, outlines, flows])
     up_btn.upload(update_image, up_btn, [input_image, up_btn, outlines, flows])
     
-    send_btn.click(cellpose_segment, [up_btn, resize, max_iter], [outlines, flows, down_btn, down_btn2])
+    send_btn.click(cellpose_segment, [up_btn, resize, max_iter, flow_threshold, cellprob_threshold], [outlines, flows, down_btn, down_btn2])
 
     #down_btn.click(download_function, None, [down_btn, down_btn2])