Update app.py

app.py

@@ -163,32 +163,37 @@ def run_model_gpu1000(img):
     masks, flows, _ = model.eval(img)#, channels = [0,0])
     return masks, flows
 
+from zipfile import Zipfile
 #@spaces.GPU(duration=10)
 def cellpose_segment(filepath, resize = 1000):
 
-    for j in range((len(filepath))):
-        img_input = imread(filepath[j])
-        #img_input = np.array(img_pil)
-        img = image_resize(img_input, resize = resize)
+    zip_path = 'masks.zip'
+    with ZipFile(zip_path, 'w') as myzip:
+        for j in range((len(filepath))):
+            img_input = imread(filepath[j])
+            #img_input = np.array(img_pil)
+            img = image_resize(img_input, resize = resize)
+            
+            resize = np.max(img.shape)
+            if resize<1000:
+                masks, flows = run_model_gpu(img)
+            elif resize < 5000:
+                masks, flows = run_model_gpu60(img)
+            elif resize < 20000:
+                masks, flows = run_model_gpu240(img)
+            else:
+                raise ValueError("Image size must be less than 20,000")
         
-        resize = np.max(img.shape)
-        if resize<1000:
-            masks, flows = run_model_gpu(img)
-        elif resize < 5000:
-            masks, flows = run_model_gpu60(img)
-        elif resize < 20000:
-            masks, flows = run_model_gpu240(img)
-        else:
-            raise ValueError("Image size must be less than 20,000")
+            target_size = (img_input.shape[1], img_input.shape[0])
+            if (target_size[0]!=img.shape[1] or target_size[1]!=img.shape[0]):
+                # scale it back to keep the orignal size
+                masks = cv2.resize(masks.astype('uint16'), target_size, interpolation=cv2.INTER_NEAREST).astype('uint16')
     
-        target_size = (img_input.shape[1], img_input.shape[0])
-        if (target_size[0]!=img.shape[1] or target_size[1]!=img.shape[0]):
-            # scale it back to keep the orignal size
-            masks = cv2.resize(masks.astype('uint16'), target_size, interpolation=cv2.INTER_NEAREST).astype('uint16')
-
-        fname_masks = os.path.splitext(filepath[j])[0]+"_masks.tif"
-        imsave(fname_masks, masks)
+            fname_masks = os.path.splitext(filepath[j])[0]+"_masks.tif"
+            imsave(fname_masks, masks)
     
+            myzip.write(fname_masks)
+            
     
     #masks, flows, _ = model.eval(img, channels=[0,0])
     flows = flows[0]
@@ -215,7 +220,10 @@ def cellpose_segment(filepath, resize = 1000):
     fname_out  = os.path.splitext(filepath[-1])[0]+"_outlines.png"
     outpix.save(fname_out) #"outlines.png")
 
-    b1 = gr.DownloadButton(visible=True, value = fname_masks)
+    if len(filepath)>1:
+        b1 = gr.DownloadButton(visible=True, value = zip_path)
+    else:
+        b1 = gr.DownloadButton(visible=True, value = fname_masks)
     b2 = gr.DownloadButton(visible=True, value = fname_out) #"outlines.png")
     
     return outpix, flows, b1, b2