Update app.py

app.py

@@ -149,45 +149,41 @@ def image_resize(img, resize=400):
 
     
 @spaces.GPU(duration=10)
-def run_model_gpu(img):
-    masks, flows, _ = model.eval(img, niter = 250)#, channels = [0,0])
+def run_model_gpu(img, max_iter):
+    masks, flows, _ = model.eval(img, niter = max_iter)#, channels = [0,0])
     return masks, flows
 
 @spaces.GPU(duration=60)
 def run_model_gpu60(img):
-    masks, flows, _ = model.eval(img, niter = 500)#, channels = [0,0])
+    masks, flows, _ = model.eval(img, niter = max_iter)#, channels = [0,0])
     return masks, flows
 
 @spaces.GPU(duration=240)
 def run_model_gpu240(img):
-    masks, flows, _ = model.eval(img, niter = 500)#, channels = [0,0])
+    masks, flows, _ = model.eval(img, niter = max_iter)#, channels = [0,0])
     return masks, flows
 
-@spaces.GPU(duration=1000)
-def run_model_gpu1000(img):
-    masks, flows, _ = model.eval(img, niter = 500)#, channels = [0,0])
-    return masks, flows
 
 
 from zipfile import ZipFile
-def cellpose_segment(filepath, resize = 1000):
+def cellpose_segment(filepath, resize = 1000,max_iter = 250):
 
     zip_path = os.path.splitext(filepath[-1])[0]+"_masks.zip"
     #zip_path = 'masks.zip'
     with ZipFile(zip_path, 'w') as myzip:
         for j in range((len(filepath))):
-            print(j)
+            print(filepath[j])
             img_input = imread(filepath[j])
             #img_input = np.array(img_pil)
             img = image_resize(img_input, resize = resize)
             
             maxsize = np.max(img.shape)
             if maxsize<1000:
-                masks, flows = run_model_gpu(img)
+                masks, flows = run_model_gpu(img, max_iter=max_iter)
             elif maxsize < 5000:
-                masks, flows = run_model_gpu60(img)
+                masks, flows = run_model_gpu60(img, max_iter=max_iter)
             elif maxsize < 20000:
-                masks, flows = run_model_gpu240(img)
+                masks, flows = run_model_gpu240(img, max_iter=max_iter)
             else:
                 raise ValueError("Image size must be less than 20,000")
         
@@ -315,7 +311,9 @@ with gr.Blocks(title = "Hello",
                 with gr.Column(scale=1):                    
                     with gr.Row():
                         resize = gr.Number(label = 'max resize', value = 1000)
-                        up_btn = gr.UploadButton("Multi-file upload (png, jpg, tif etc)", visible=True, file_count = "multiple")                        
+                        max_iter = gr.Number(label = 'max iterations', value = 250)
+                        
+                    up_btn = gr.UploadButton("Multi-file upload (png, jpg, tif etc)", visible=True, file_count = "multiple")                        
                     
                     #gr.HTML("""<h4 style="color:white;"> Note2: Only the first image of a tif will display the segmentations, but you can download segmentations for all planes. </h4>""")
                     
@@ -340,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], [outlines, flows, down_btn, down_btn2])
+    send_btn.click(cellpose_segment, [up_btn, resize, max_iter], [outlines, flows, down_btn, down_btn2])
 
     #down_btn.click(download_function, None, [down_btn, down_btn2])